fix(draw): better queue negotiator

4 files changed, 122 insertions, 67 deletions

diff --git a/stockton-render/src/draw/context.rs b/stockton-render/src/draw/context.rs
index 3de857a..5448357 100644
--- a/stockton-render/src/draw/context.rs
+++ b/stockton-render/src/draw/context.rs
@@ -19,10 +19,10 @@ use super::{
     buffer::ModifiableBuffer,
     draw_buffers::{DrawBuffers, UvPoint},
     pipeline::CompletePipeline,
-    queue_negotiator::QueueNegotiator,
+    queue_negotiator::{QueueNegotiator, DrawQueue},
     render::do_render,
     target::{SwapchainProperties, TargetChain},
-    texture::{resolver::FsResolver, TextureLoadConfig, TextureRepo},
+    texture::{resolver::FsResolver, TexLoadQueue, TextureLoadConfig, TextureRepo},
     ui::{
         do_render as do_render_ui, ensure_textures as ensure_textures_ui, UiPipeline, UiPoint,
         UiTextures,
@@ -104,32 +104,30 @@ impl<'a, M: 'static + MinBspFeatures<VulkanSystem>> RenderingContext<'a, M> {
         // TODO: Properly figure out which adapter to use
         let adapter = adapters.remove(0);
 
-        let mut draw_queue_negotiator = QueueNegotiator::find(&adapter, |family| {
-            surface.supports_queue_family(family) && family.queue_type().supports_graphics()
-        })
-        .context("Error creating draw queue negotiator")?;
-
-        let mut tex_queue_negotiator =
-            QueueNegotiator::find(&adapter, TextureRepo::queue_family_filter)
-                .context("Error creating texture queue negotiator")?;
-        debug!(
-            "Using draw queue family {:?}",
-            draw_queue_negotiator.family_id()
-        );
-        debug!(
-            "Using tex queue family {:?}",
-            tex_queue_negotiator.family_id()
-        );
+        let (mut queue_negotiator, surface) = {
+            let dq: DrawQueue = DrawQueue { surface };
+
+            let qn = QueueNegotiator::find(&adapter, &[
+                &dq,
+                &TexLoadQueue,
+            ])
+            .context("Error creating draw queue negotiator")?;
+
+            (qn, dq.surface)
+        };
 
         // Device & Queue groups
         let (device_lock, mut queue_groups) = {
+            let (df, dqs) = queue_negotiator.family_spec::<DrawQueue>(&adapter.queue_families, 1).ok_or(EnvironmentError::NoSuitableFamilies)?;
+            let (tf, tqs) = queue_negotiator.family_spec::<TexLoadQueue>(&adapter.queue_families, 2).ok_or(EnvironmentError::NoSuitableFamilies)?;
+
             let gpu = unsafe {
                 adapter
                     .physical_device
                     .open(
                         &[
-                            (draw_queue_negotiator.family(&adapter), &[1.0]),
-                            (tex_queue_negotiator.family(&adapter), &[1.0]),
+                            (df, dqs.as_slice()),
+                            (tf, tqs.as_slice()),
                         ],
                         hal::Features::empty(),
                     )
@@ -153,7 +151,7 @@ impl<'a, M: 'static + MinBspFeatures<VulkanSystem>> RenderingContext<'a, M> {
         // Command pool
         let mut cmd_pool = unsafe {
             device.create_command_pool(
-                draw_queue_negotiator.family_id(),
+                queue_negotiator.family::<DrawQueue>().ok_or(EnvironmentError::NoSuitableFamilies)?,
                 CommandPoolCreateFlags::RESET_INDIVIDUAL,
             )
         }
@@ -174,9 +172,9 @@ impl<'a, M: 'static + MinBspFeatures<VulkanSystem>> RenderingContext<'a, M> {
         debug!("Creating 3D Texture Repo");
         let tex_repo = TextureRepo::new(
             device_lock.clone(),
-            tex_queue_negotiator.family_id(),
-            tex_queue_negotiator
-                .get_queue(&mut queue_groups)
+            queue_negotiator.family::<TexLoadQueue>().ok_or(EnvironmentError::NoQueues)?,
+            queue_negotiator
+                .get_queue::<TexLoadQueue>(&mut queue_groups)
                 .ok_or(EnvironmentError::NoQueues)
                 .context("Error getting 3D texture loader queue")?,
             &adapter,
@@ -191,9 +189,9 @@ impl<'a, M: 'static + MinBspFeatures<VulkanSystem>> RenderingContext<'a, M> {
         debug!("Creating UI Texture Repo");
         let ui_tex_repo = TextureRepo::new(
             device_lock.clone(),
-            tex_queue_negotiator.family_id(),
-            tex_queue_negotiator
-                .get_queue(&mut queue_groups)
+            queue_negotiator.family::<TexLoadQueue>().ok_or(EnvironmentError::NoQueues)?,
+            queue_negotiator
+                .get_queue::<TexLoadQueue>(&mut queue_groups)
                 .ok_or(EnvironmentError::NoQueues)
                 .context("Error getting UI texture loader queue")?,
             &adapter,
@@ -249,8 +247,8 @@ impl<'a, M: 'static + MinBspFeatures<VulkanSystem>> RenderingContext<'a, M> {
             device: device_lock,
             adapter,
 
-            queue: draw_queue_negotiator
-                .get_queue(&mut queue_groups)
+            queue: queue_negotiator
+                .get_queue::<DrawQueue>(&mut queue_groups)
                 .ok_or(EnvironmentError::NoQueues)
                 .context("Error getting draw queue")?,
 
diff --git a/stockton-render/src/draw/queue_negotiator.rs b/stockton-render/src/draw/queue_negotiator.rs
index 4ad6823..4f00c9e 100644
--- a/stockton-render/src/draw/queue_negotiator.rs
+++ b/stockton-render/src/draw/queue_negotiator.rs
@@ -1,65 +1,117 @@
 use crate::{error::EnvironmentError, types::*};
-use anyhow::Result;
+use anyhow::{Result, Error};
 use hal::queue::family::QueueFamilyId;
 use std::sync::{Arc, RwLock};
+use std::collections::HashMap;
+use std::any::TypeId;
 
 pub struct QueueNegotiator {
-    family_id: QueueFamilyId,
-    already_allocated: Vec<Arc<RwLock<QueueT>>>,
-    next_share: usize,
+    family_ids: HashMap<TypeId, QueueFamilyId>,
+    already_allocated: HashMap<TypeId, (Vec<Arc<RwLock<QueueT>>>, usize)>,
+}
+
+pub trait QueueFamilySelector: 'static {
+    fn is_suitable(&self, family: &QueueFamilyT) -> bool;
+    
+    fn get_type_id_self(&self) -> TypeId{
+        TypeId::of::<Self>()
+    }
+
+    fn get_type_id() -> TypeId where Self: Sized {
+        TypeId::of::<Self>()
+    }
 }
 
 impl QueueNegotiator {
-    pub fn find<F: FnMut(&&QueueFamilyT) -> bool>(adapter: &Adapter, filter: F) -> Result<Self> {
-        let family = adapter
-            .queue_families
-            .iter()
-            .find(filter)
-            .ok_or(EnvironmentError::NoSuitableFamilies)?;
+    pub fn find(adapter: &Adapter, stacks: &[&dyn QueueFamilySelector]) -> Result<Self> {
+        let mut families = HashMap::new();
+        for filter in stacks {
+            let candidates: Vec<&QueueFamilyT> = adapter
+                .queue_families
+                .iter()
+                .filter(|x| filter.is_suitable(*x))
+                .collect();
+            
+            if candidates.len() == 0 {
+                return Err(Error::new(EnvironmentError::NoSuitableFamilies));
+            }
+
+            // Prefer using unique families
+            let family = match candidates.iter().filter(|x| !families.values().any(|y| *y == x.id())).next() {
+                Some(x) => *x,
+                None => candidates[0]
+            };
+
+            families.insert(filter.get_type_id_self(), family.id());
+        }
 
         Ok(QueueNegotiator {
-            family_id: family.id(),
-            already_allocated: Vec::with_capacity(family.max_queues()),
-            next_share: 0,
+            family_ids: families,
+            already_allocated: HashMap::new(),
         })
     }
 
-    pub fn family<'a>(&self, adapter: &'a Adapter) -> &'a QueueFamilyT {
-        adapter
-            .queue_families
-            .iter()
-            .find(|x| x.id() == self.family_id)
-            .unwrap()
-    }
-
-    pub fn family_id(&self) -> QueueFamilyId {
-        self.family_id
-    }
+    pub fn get_queue<T: QueueFamilySelector>(&mut self, groups: &mut Vec<QueueGroup>) -> Option<Arc<RwLock<QueueT>>> {
+        let tid = T::get_type_id();
+        let family_id = self.family_ids.get(&tid)?;
 
-    pub fn get_queue(&mut self, groups: &mut Vec<QueueGroup>) -> Option<Arc<RwLock<QueueT>>> {
         match groups
             .iter()
-            .position(|x| x.queues.len() > 0 && x.family == self.family_id)
+            .position(|x| x.queues.len() > 0 && x.family == *family_id)
         {
             Some(idx) => {
                 // At least one remaining queue
                 let queue = groups[idx].queues.pop().unwrap();
                 let queue = Arc::new(RwLock::new(queue));
 
-                self.already_allocated.push(queue.clone());
+                self.add_to_allocated::<T>(queue.clone());
 
                 Some(queue)
             }
-            None => {
-                if self.already_allocated.len() == 0 {
-                    return None;
+            None => match self.already_allocated.get_mut(&tid) {
+                Some((queues, next_share)) => {
+                    let queue = (&queues[*next_share]).clone();
+
+                    *next_share += 1;
+
+                    Some(queue)
                 }
+                None => None
+            }
+        }
+    }
 
-                let queue = self.already_allocated[self.next_share].clone();
-                self.next_share = (self.next_share + 1) % self.already_allocated.len();
+    pub fn family<T: QueueFamilySelector>(&self) -> Option<QueueFamilyId> {
+        self.family_ids.get(&TypeId::of::<T>()).cloned()
+    }
 
-                Some(queue)
+    fn add_to_allocated<T: QueueFamilySelector>(&mut self, queue: Arc<RwLock<QueueT>>) {
+        let tid = T::get_type_id();
+        match self.already_allocated.get_mut(&tid) {
+            None => {
+                self.already_allocated.insert(tid, (vec![queue], 0));
+            },
+            Some(x) => {
+                x.0.push(queue);
             }
         }
     }
+
+    pub fn family_spec<'a, T: QueueFamilySelector>(&self, queue_families: &'a Vec<QueueFamilyT>, count: usize) -> Option<(&'a QueueFamilyT, Vec<f32>)> {
+        let qf_id = self.family::<T>()?;
+
+        let qf = queue_families.iter().find(|x| x.id() == qf_id)?;
+        let v = vec![1.0; count];
+
+        Some((qf, v))
+    }
+}
+
+pub struct DrawQueue {
+    pub surface: SurfaceT
+}
+impl QueueFamilySelector for DrawQueue {
+    fn is_suitable(&self, family: &QueueFamilyT) -> bool {
+        self.surface.supports_queue_family(family) && family.queue_type().supports_graphics()
+    }
 }
diff --git a/stockton-render/src/draw/texture/mod.rs b/stockton-render/src/draw/texture/mod.rs
index 62b28c5..5720349 100644
--- a/stockton-render/src/draw/texture/mod.rs
+++ b/stockton-render/src/draw/texture/mod.rs
@@ -12,7 +12,7 @@ pub use self::block::TexturesBlock;
 pub use self::image::LoadableImage;
 pub use self::load::TextureLoadConfig;
 pub use self::loader::BlockRef;
-pub use self::repo::TextureRepo;
+pub use self::repo::{TextureRepo, TexLoadQueue};
 
 /// The size of each pixel in an image
 pub const PIXEL_SIZE: usize = std::mem::size_of::<u8>() * 4;
diff --git a/stockton-render/src/draw/texture/repo.rs b/stockton-render/src/draw/texture/repo.rs
index ef35610..d2aa9f2 100644
--- a/stockton-render/src/draw/texture/repo.rs
+++ b/stockton-render/src/draw/texture/repo.rs
@@ -4,6 +4,7 @@ use super::{
     loader::{BlockRef, LoaderRequest, TextureLoader, TextureLoaderRemains, NUM_SIMULTANEOUS_CMDS},
     resolver::TextureResolver,
 };
+use crate::draw::queue_negotiator::QueueFamilySelector;
 use crate::error::LockPoisoned;
 use crate::types::*;
 
@@ -43,10 +44,6 @@ pub struct TextureRepo<'a> {
 }
 
 impl<'a> TextureRepo<'a> {
-    pub fn queue_family_filter(family: &&QueueFamilyT) -> bool {
-        family.queue_type().supports_transfer() && family.max_queues() >= NUM_SIMULTANEOUS_CMDS
-    }
-
     pub fn new<R: 'static + TextureResolver + Send + Sync>(
         device_lock: Arc<RwLock<DeviceT>>,
         family: QueueFamilyId,
@@ -198,3 +195,11 @@ impl<'a> TextureRepo<'a> {
         }
     }
 }
+
+pub struct TexLoadQueue;
+
+impl QueueFamilySelector for TexLoadQueue {
+    fn is_suitable(&self, family: &QueueFamilyT) -> bool {
+        family.queue_type().supports_transfer() && family.max_queues() >= NUM_SIMULTANEOUS_CMDS
+    }
+}
\ No newline at end of file