4 files changed, 126 insertions, 115 deletions

diff --git a/stockton-skeleton/src/context.rs b/stockton-skeleton/src/context.rs
index 348ddbd..bef9f9d 100644
--- a/stockton-skeleton/src/context.rs
+++ b/stockton-skeleton/src/context.rs
@@ -31,6 +31,7 @@ use crate::{
     draw_passes::Singular,
     error::{EnvironmentError, LockPoisoned},
     mem::MemoryPool,
+    queue_negotiator::QueueFamilyNegotiator,
     types::*,
 };
 
@@ -95,25 +96,18 @@ impl RenderingContext {
         let adapter = adapters.remove(0);
 
         // Queue Negotiator
-        let mut queue_families_specs = Vec::new();
-        let (mut queue_negotiator, surface) = {
+        let (queue_negotiator, surface) = {
             let dq: DrawQueue = DrawQueue { surface };
 
-            let mut qn = QueueNegotiator::default();
+            let mut qn = QueueFamilyNegotiator::new();
 
             // Draw Queue
-            qn.find(&adapter, &dq)
+            qn.find(&adapter, &dq, 1)
                 .context("Couldn't find draw queue family")?;
-            queue_families_specs.push(
-                qn.family_spec::<DrawQueue>(&adapter.queue_families, 1)
-                    .context("Couldn't find draw queue family")?,
-            );
 
             // Auxiliary queues for DP
-            queue_families_specs.extend(
-                IDP::find_aux_queues(&adapter, &mut qn)
-                    .context("Level pass couldn't populate queue negotiator")?,
-            );
+            IDP::find_aux_queues(&adapter, &mut qn)
+                .context("Level pass couldn't populate queue family negotiator")?;
 
             (qn, dq.surface)
         };
@@ -123,30 +117,19 @@ impl RenderingContext {
             // TODO: This sucks, but hal is restrictive on how we can pass this specific argument.
 
             // Deduplicate families & convert to specific type.
-            let mut queue_families_specs_real = Vec::with_capacity(queue_families_specs.len());
-            for (qf, ns) in queue_families_specs.iter_mut() {
-                if let Some(existing_family_spec) = queue_families_specs_real
-                    .iter()
-                    .position(|(qf2, _): &(&QueueFamilyT, &[f32])| qf2.id() == qf.id())
-                {
-                    ns.extend(queue_families_specs_real[existing_family_spec].1.iter());
-                    queue_families_specs_real[existing_family_spec] = (*qf, ns.as_slice());
-                } else {
-                    queue_families_specs_real.push((*qf, ns.as_slice()))
-                }
-            }
+            let open_spec = queue_negotiator.get_open_spec(&adapter);
 
             let gpu = unsafe {
                 adapter
                     .physical_device
-                    .open(queue_families_specs_real.as_slice(), hal::Features::empty())
+                    .open(&open_spec.as_vec(), hal::Features::empty())
                     .context("Error opening logical device")?
             };
 
             (Arc::new(RwLock::new(gpu.device)), gpu.queue_groups)
         };
 
-        queue_negotiator.set_queue_groups(queue_groups);
+        let mut queue_negotiator = queue_negotiator.finish(queue_groups);
 
         // Context properties
         let properties = ContextProperties::find_best(&adapter, &surface)
diff --git a/stockton-skeleton/src/draw_passes/cons.rs b/stockton-skeleton/src/draw_passes/cons.rs
index cc4f142..51c06a6 100644
--- a/stockton-skeleton/src/draw_passes/cons.rs
+++ b/stockton-skeleton/src/draw_passes/cons.rs
@@ -2,7 +2,7 @@
 //! Note that this can be extended to an arbitrary amount of draw passes.
 
 use super::{Beginning, DrawPass, End, IntoDrawPass, Middle, Singular};
-use crate::{context::RenderingContext, queue_negotiator::QueueNegotiator, types::*};
+use crate::{context::RenderingContext, queue_negotiator::QueueFamilyNegotiator, types::*};
 use stockton_types::Session;
 
 use anyhow::Result;
@@ -73,11 +73,11 @@ macro_rules! into_shared_impl {
 
         fn find_aux_queues<'a>(
             adapter: &'a Adapter,
-            queue_negotiator: &mut QueueNegotiator,
-        ) -> Result<Vec<(&'a QueueFamilyT, Vec<f32>)>> {
-            let mut v = IA::find_aux_queues(adapter, queue_negotiator)?;
-            v.extend(IB::find_aux_queues(adapter, queue_negotiator)?);
-            Ok(v)
+            queue_negotiator: &mut QueueFamilyNegotiator,
+        ) -> Result<()> {
+            IA::find_aux_queues(adapter, queue_negotiator)?;
+            IB::find_aux_queues(adapter, queue_negotiator)?;
+            Ok(())
         }
     };
 }
diff --git a/stockton-skeleton/src/draw_passes/mod.rs b/stockton-skeleton/src/draw_passes/mod.rs
index 5a92a45..d830fe5 100644
--- a/stockton-skeleton/src/draw_passes/mod.rs
+++ b/stockton-skeleton/src/draw_passes/mod.rs
@@ -1,8 +1,7 @@
 //! Traits and common draw passes.
 use std::ops::Range;
 
-use super::{queue_negotiator::QueueNegotiator, RenderingContext};
-use crate::types::*;
+use crate::{queue_negotiator::QueueFamilyNegotiator, types::*, RenderingContext};
 use hal::{
     image::Layout,
     pass::{AttachmentLoadOp, AttachmentOps, AttachmentStoreOp},
@@ -45,11 +44,10 @@ pub trait IntoDrawPass<T: DrawPass<P>, P: PassPosition> {
 
     /// This function should ask the queue negotatior to find families for any auxilary operations this draw pass needs to perform
     /// For example, .find(&TexLoadQueue)
-    /// It should return then call .family_spec for each queue type negotiated and return the results.
     fn find_aux_queues<'a>(
         adapter: &'a Adapter,
-        queue_negotiator: &mut QueueNegotiator,
-    ) -> Result<Vec<(&'a QueueFamilyT, Vec<f32>)>>;
+        queue_negotiator: &mut QueueFamilyNegotiator,
+    ) -> Result<()>;
 }
 
 /// Used so that draw passes can determine what state shared resources are in and how they should be left.
diff --git a/stockton-skeleton/src/queue_negotiator.rs b/stockton-skeleton/src/queue_negotiator.rs
index f3e38fd..b78fe33 100644
--- a/stockton-skeleton/src/queue_negotiator.rs
+++ b/stockton-skeleton/src/queue_negotiator.rs
@@ -1,70 +1,76 @@
 //! Used for requesting appropriate queue families, and sharing/allocating queues as necessary.
-//! This is created by `RenderingContext`, and should mostly be accessed from [`crate::draw_passes::IntoDrawPass`].
+//! You'll mostly use these from [`crate::draw_passes::IntoDrawPass`].
 //!
-//! For example, to use a `TexLoadQueue` in your drawpass:
+//! For example, to use a `TexLoadQueue` in your drawpass, first find the family during the init phase:
 //! ```
-//! # use crate::{types::*, texture::TexLoadQueue};
+//! # use stockton_skeleton::{types::*, texture::TexLoadQueue, queue_negotiator::*};
+//! # use anyhow::Result;
 //! fn find_aux_queues<'c>(
 //!    adapter: &'c Adapter,
-//!    queue_negotiator: &mut QueueNegotiator,
-//! ) -> Result<Vec<(&'c QueueFamilyT, Vec<f32>)>> {
-//!     queue_negotiator.find(adapter, &TexLoadQueue)?;
+//!    queue_negotiator: &mut QueueFamilyNegotiator,
+//! ) -> Result<()> {
+//!     queue_negotiator.find(adapter, &TexLoadQueue, 1)?;
 //!
-//!     Ok(vec![queue_negotiator
-//!         .family_spec::<TexLoadQueue>(&adapter.queue_families, 1)
-//!         .ok_or(EnvironmentError::NoSuitableFamilies)?])
+//!     Ok(())
 //! }
 //! ```
 //!
 //! Then get your queue in [`crate::draw_passes::IntoDrawPass::init`]
+//!
 //! ```
-//! # use crate::{types::*, context::RenderingContext, texture::TexLoadQueue};
+//! # use stockton_skeleton::{types::*, context::RenderingContext, texture::TexLoadQueue, queue_negotiator::*};
+//! # use anyhow::Result;
 //! # use stockton_types::Session;
-//! fn init(
-//!     self,
-//!     session: &mut Session,
-//!     context: &mut RenderingContext,
-//! ) -> Result<LevelDrawPass<'a, M>> {
-//!     let queue = context.queue_negotiator_mut().get_queue::<TexLoadQueue>().unwrap();
-//! }
+//! # struct YourDrawPass;
+//! # fn init(
+//! #    context: &mut RenderingContext,
+//! # ) -> Result<()> {
+//! let queue = context.queue_negotiator_mut().get_queue::<TexLoadQueue>()?;
+//!  // ...
+//! # Ok(())
+//! # }
 //! ```
 
 use crate::{
     error::{EnvironmentError, UsageError},
     types::*,
 };
-
 use anyhow::{bail, Error, Result};
 use hal::queue::family::QueueFamilyId;
 use std::{
     any::TypeId,
-    collections::HashMap,
+    collections::hash_map::{Entry, HashMap},
     sync::{Arc, RwLock},
 };
 
 /// A queue, possibly shared between threads.
 pub type SharedQueue = Arc<RwLock<QueueT>>;
 
-/// Used to find appropriate queue families and share queues from them as needed.
-pub struct QueueNegotiator {
-    family_ids: HashMap<TypeId, QueueFamilyId>,
-    already_allocated: HashMap<TypeId, (Vec<SharedQueue>, usize)>,
-    all: Vec<QueueGroup>,
+/// Used to find appropriate queue families during init phase.
+pub struct QueueFamilyNegotiator {
+    /// Family and count being used for each selector
+    family_ids: HashMap<TypeId, (usize, QueueFamilyId)>,
 }
 
-/// Can be used to select an appropriate queue family
-pub trait QueueFamilySelector: 'static {
-    /// Return true if the given family is suitable
-    fn is_suitable(&self, family: &QueueFamilyT) -> bool;
-}
+impl QueueFamilyNegotiator {
+    /// Create a new, empty, QueueFamilyNegotiator
+    pub fn new() -> Self {
+        QueueFamilyNegotiator {
+            family_ids: HashMap::new(),
+        }
+    }
 
-impl QueueNegotiator {
     /// Attempt to find an appropriate queue family using the given selector.
-    /// Returns early if the *type* of the selector has already been allocated a family.
+    /// If T has already been used in a different call for find, it will request the sum of the `count` values from both calls.
     /// This should usually be called by [`crate::draw_passes::IntoDrawPass::find_aux_queues`].
-    pub fn find<T: QueueFamilySelector>(&mut self, adapter: &Adapter, filter: &T) -> Result<()> {
-        if self.family_ids.contains_key(&TypeId::of::<T>()) {
-            return Ok(());
+    pub fn find<'a, T: QueueFamilySelector>(
+        &mut self,
+        adapter: &'a Adapter,
+        filter: &T,
+        mut count: usize,
+    ) -> Result<()> {
+        if let Entry::Occupied(e) = self.family_ids.entry(TypeId::of::<T>()) {
+            count = count + e.get().0;
         }
 
         let candidates: Vec<&QueueFamilyT> = adapter
@@ -80,26 +86,87 @@ impl QueueNegotiator {
         // Prefer using unique families
         let family = match candidates
             .iter()
-            .find(|x| !self.family_ids.values().any(|y| *y == x.id()))
+            .find(|x| !self.family_ids.values().any(|y| y.1 == x.id()))
         {
             Some(x) => *x,
             None => candidates[0],
         };
 
-        self.family_ids.insert(TypeId::of::<T>(), family.id());
+        self.family_ids
+            .insert(TypeId::of::<T>(), (count, family.id()));
 
         Ok(())
     }
 
+    /// Used to get a spec passed to [`hal::adapter::PhysicalDevice::open`]
+    pub(crate) fn get_open_spec<'a>(&self, adapter: &'a Adapter) -> AdapterOpenSpec<'a> {
+        // Deduplicate families & convert to specific type.
+        let mut spec = Vec::with_capacity(self.family_ids.len());
+        for (count, qf_id) in self.family_ids.values() {
+            if let Some(existing_family_spec) = spec
+                .iter()
+                .position(|(qf2_id, _): &(&QueueFamilyT, Vec<f32>)| qf2_id.id() == *qf_id)
+            {
+                for _ in 0..*count {
+                    spec[existing_family_spec].1.push(1.0);
+                }
+            } else {
+                let family = adapter
+                    .queue_families
+                    .iter()
+                    .find(|x| x.id() == *qf_id)
+                    .unwrap();
+                spec.push((family, vec![1.0; *count]))
+            }
+        }
+        AdapterOpenSpec(spec.into_boxed_slice())
+    }
+
+    /// Finish selecting our queue families, and turn this into a `QueueNegotiator`
+    pub fn finish<'a>(self, queue_groups: Vec<QueueGroup>) -> QueueNegotiator {
+        QueueNegotiator {
+            family_ids: self.family_ids,
+            already_allocated: HashMap::new(),
+            all: queue_groups,
+        }
+    }
+}
+
+/// Used internally in calls to [`hal::adapter::PhysicalDevice::open`]
+pub(crate) struct AdapterOpenSpec<'a>(Box<[(&'a QueueFamilyT, Vec<f32>)]>);
+
+impl<'a> AdapterOpenSpec<'a> {
+    pub fn as_vec(&self) -> Vec<(&'a QueueFamilyT, &[f32])> {
+        let mut v = Vec::with_capacity(self.0.len());
+        for (qf, cs) in self.0.iter() {
+            v.push((*qf, cs.as_slice()));
+        }
+
+        v
+    }
+}
+
+/// Used to share queues from families selected during init phase.
+pub struct QueueNegotiator {
+    family_ids: HashMap<TypeId, (usize, QueueFamilyId)>,
+    already_allocated: HashMap<TypeId, (Vec<SharedQueue>, usize)>,
+    all: Vec<QueueGroup>,
+}
+
+/// Can be used to select an appropriate queue family
+pub trait QueueFamilySelector: 'static {
+    /// Return true if the given family is suitable
+    fn is_suitable(&self, family: &QueueFamilyT) -> bool;
+}
+
+impl QueueNegotiator {
     /// Get a (possibly shared) queue. You should prefer to call this once and store the result.
-    /// You should already have called [`self::QueueNegotiator::find`] and [`self::QueueNegotiator::family_spec`],
-    /// otherwise this will return an error.
+    /// You should already have called [`self::QueueFamilyNegotiator::find`], otherwise this will return an error.
     ///
-    /// Round-robin allocation is used to try to fairly distribute work between each queue.
     /// The family of the queue returned is guaranteed to meet the spec of the `QueueFamilySelector` originally used by `find`.
     pub fn get_queue<T: QueueFamilySelector>(&mut self) -> Result<Arc<RwLock<QueueT>>> {
         let tid = TypeId::of::<T>();
-        let family_id = self
+        let (_, family_id) = self
             .family_ids
             .get(&tid)
             .ok_or(UsageError::QueueNegotiatorMisuse)?;
@@ -131,36 +198,9 @@ impl QueueNegotiator {
         }
     }
 
-    /// Convenience function to get a queue spec for the given selector.
-    /// You should probably call this from [`crate::draw_passes::IntoDrawPass::find_aux_queues`].
-    /// `count` is the maximum number of individual queues to request. You may get less than this, in which case they will be shared.
-    /// This will return an error if you haven't called [`self::QueueNegotiator::find`] beforehand, or if there were no suitable queue families.
-    pub fn family_spec<'a, T: QueueFamilySelector>(
-        &self,
-        queue_families: &'a [QueueFamilyT],
-        count: usize,
-    ) -> Result<(&'a QueueFamilyT, Vec<f32>)> {
-        let qf_id = self
-            .family::<T>()
-            .ok_or(UsageError::QueueNegotiatorMisuse)?;
-
-        let qf = queue_families
-            .iter()
-            .find(|x| x.id() == qf_id)
-            .ok_or(EnvironmentError::NoSuitableFamilies)?;
-        let v = vec![1.0; count];
-
-        Ok((qf, v))
-    }
-
     /// Get the queue family ID being used by the given selector
     pub fn family<T: QueueFamilySelector>(&self) -> Option<QueueFamilyId> {
-        self.family_ids.get(&TypeId::of::<T>()).cloned()
-    }
-
-    /// Used internally to get the queue groups from the adapter.
-    pub(crate) fn set_queue_groups(&mut self, queue_groups: Vec<QueueGroup>) {
-        self.all = queue_groups
+        self.family_ids.get(&TypeId::of::<T>()).map(|x| x.1)
     }
 
     /// Used internally to mark that we've started sharing a queue
@@ -177,16 +217,6 @@ impl QueueNegotiator {
     }
 }
 
-impl Default for QueueNegotiator {
-    fn default() -> Self {
-        QueueNegotiator {
-            family_ids: HashMap::new(),
-            already_allocated: HashMap::new(),
-            all: vec![],
-        }
-    }
-}
-
 /// A queue suitable for drawing to a given surface with.
 pub struct DrawQueue {
     pub surface: SurfaceT,