1 files changed, 357 insertions, 154 deletions

diff --git a/src/descriptor_set/mod.rs b/src/descriptor_set/mod.rs
index b69deee..82d4199 100644
--- a/src/descriptor_set/mod.rs
+++ b/src/descriptor_set/mod.rs
@@ -22,7 +22,7 @@
 //! that are often used together in the same set so that you can keep the same set binding through
 //! multiple draws.
 //!
-//! # Example
+//! # Examples
 //!
 //! > **Note**: This section describes the simple way to bind resources. There are more optimized
 //! > ways.
@@ -52,236 +52,439 @@
 //!
 //! - A `DescriptorSetLayout` is a Vulkan object that describes to the Vulkan implementation the
 //!   layout of a future descriptor set. When you allocate a descriptor set, you have to pass an
-//!   instance of this object. This is represented with the `DescriptorSetLayout` type in
+//!   instance of this object. This is represented with the [`DescriptorSetLayout`] type in
 //!   vulkano.
 //! - A `DescriptorPool` is a Vulkan object that holds the memory of descriptor sets and that can
 //!   be used to allocate and free individual descriptor sets. This is represented with the
-//!   `UnsafeDescriptorPool` type in vulkano.
+//!   [`DescriptorPool`] type in vulkano.
 //! - A `DescriptorSet` contains the bindings to resources and is allocated from a pool. This is
-//!   represented with the `UnsafeDescriptorSet` type in vulkano.
+//!   represented with the [`UnsafeDescriptorSet`] type in vulkano.
 //!
 //! In addition to this, vulkano defines the following:
 //!
-//! - The `DescriptorPool` trait can be implemented on types from which you can allocate and free
-//!   descriptor sets. However it is different from Vulkan descriptor pools in the sense that an
-//!   implementation of the `DescriptorPool` trait can manage multiple Vulkan descriptor pools.
-//! - The `StdDescriptorPool` type is a default implementation of the `DescriptorPool` trait.
-//! - The `DescriptorSet` trait is implemented on types that wrap around Vulkan descriptor sets in
+//! - The [`DescriptorSetAllocator`] trait can be implemented on types from which you can allocate
+//!   and free descriptor sets. However it is different from Vulkan descriptor pools in the sense
+//!   that an implementation of the [`DescriptorSetAllocator`] trait can manage multiple Vulkan
+//!   descriptor pools.
+//! - The [`StandardDescriptorSetAllocator`] type is a default implementation of the
+//!   [`DescriptorSetAllocator`] trait.
+//! - The [`DescriptorSet`] trait is implemented on types that wrap around Vulkan descriptor sets in
 //!   a safe way. A Vulkan descriptor set is inherently unsafe, so we need safe wrappers around
 //!   them.
-//! - The `SimpleDescriptorSet` type is a default implementation of the `DescriptorSet` trait.
-//! - The `DescriptorSetsCollection` trait is implemented on collections of types that implement
-//!   `DescriptorSet`. It is what you pass to the draw functions.
-
-pub use self::collection::DescriptorSetsCollection;
-pub use self::fixed_size_pool::FixedSizeDescriptorSetsPool;
-use self::layout::DescriptorSetLayout;
-pub use self::persistent::PersistentDescriptorSet;
-pub use self::persistent::PersistentDescriptorSetBuildError;
-pub use self::persistent::PersistentDescriptorSetError;
-use self::sys::UnsafeDescriptorSet;
-use crate::buffer::BufferAccess;
-use crate::descriptor_set::layout::{DescriptorBufferDesc, DescriptorDescTy};
-use crate::device::DeviceOwned;
-use crate::image::view::ImageViewAbstract;
-use crate::SafeDeref;
-use crate::VulkanObject;
-use smallvec::SmallVec;
-use std::hash::Hash;
-use std::hash::Hasher;
-use std::sync::Arc;
+//! - The [`DescriptorSetsCollection`] trait is implemented on collections of types that implement
+//!   [`DescriptorSet`]. It is what you pass to the draw functions.
+//!
+//! [`DescriptorPool`]: pool::DescriptorPool
+//! [`DescriptorSetAllocator`]: allocator::DescriptorSetAllocator
+//! [`StandardDescriptorSetAllocator`]: allocator::StandardDescriptorSetAllocator
+
+pub(crate) use self::update::{check_descriptor_write, DescriptorWriteInfo};
+pub use self::{
+    collection::DescriptorSetsCollection,
+    persistent::PersistentDescriptorSet,
+    update::{DescriptorSetUpdateError, WriteDescriptorSet, WriteDescriptorSetElements},
+};
+use self::{layout::DescriptorSetLayout, sys::UnsafeDescriptorSet};
+use crate::{
+    buffer::{view::BufferView, Subbuffer},
+    descriptor_set::layout::DescriptorType,
+    device::DeviceOwned,
+    image::view::ImageViewAbstract,
+    sampler::Sampler,
+    DeviceSize, OomError, VulkanObject,
+};
+use ahash::HashMap;
+use smallvec::{smallvec, SmallVec};
+use std::{
+    error::Error,
+    fmt::{Display, Error as FmtError, Formatter},
+    hash::{Hash, Hasher},
+    ops::Range,
+    ptr,
+    sync::Arc,
+};
 
+pub mod allocator;
 mod collection;
-pub mod fixed_size_pool;
 pub mod layout;
 pub mod persistent;
 pub mod pool;
 pub mod sys;
+mod update;
 
 /// Trait for objects that contain a collection of resources that will be accessible by shaders.
 ///
 /// Objects of this type can be passed when submitting a draw command.
-pub unsafe trait DescriptorSet: DeviceOwned {
+pub unsafe trait DescriptorSet: DeviceOwned + Send + Sync {
     /// Returns the inner `UnsafeDescriptorSet`.
     fn inner(&self) -> &UnsafeDescriptorSet;
 
     /// Returns the layout of this descriptor set.
     fn layout(&self) -> &Arc<DescriptorSetLayout>;
 
+    /// Returns the variable descriptor count that this descriptor set was allocated with.
+    fn variable_descriptor_count(&self) -> u32;
+
     /// Creates a [`DescriptorSetWithOffsets`] with the given dynamic offsets.
-    fn offsets<I>(self, dynamic_offsets: I) -> DescriptorSetWithOffsets
+    fn offsets(
+        self: Arc<Self>,
+        dynamic_offsets: impl IntoIterator<Item = u32>,
+    ) -> DescriptorSetWithOffsets
     where
-        Self: Sized + Send + Sync + 'static,
-        I: IntoIterator<Item = u32>,
+        Self: Sized + 'static,
     {
         DescriptorSetWithOffsets::new(self, dynamic_offsets)
     }
 
-    /// Returns the number of buffers within this descriptor set.
-    fn num_buffers(&self) -> usize;
+    /// Returns the resources bound to this descriptor set.
+    fn resources(&self) -> &DescriptorSetResources;
+}
 
-    /// Returns the `index`th buffer of this descriptor set, or `None` if out of range. Also
-    /// returns the index of the descriptor that uses this buffer.
-    ///
-    /// The valid range is between 0 and `num_buffers()`.
-    fn buffer(&self, index: usize) -> Option<(&dyn BufferAccess, u32)>;
+impl PartialEq for dyn DescriptorSet {
+    #[inline]
+    fn eq(&self, other: &Self) -> bool {
+        self.inner() == other.inner()
+    }
+}
 
-    /// Returns the number of images within this descriptor set.
-    fn num_images(&self) -> usize;
+impl Eq for dyn DescriptorSet {}
 
-    /// Returns the `index`th image of this descriptor set, or `None` if out of range. Also returns
-    /// the index of the descriptor that uses this image.
-    ///
-    /// The valid range is between 0 and `num_images()`.
-    fn image(&self, index: usize) -> Option<(&dyn ImageViewAbstract, u32)>;
+impl Hash for dyn DescriptorSet {
+    fn hash<H: Hasher>(&self, state: &mut H) {
+        self.inner().hash(state);
+    }
 }
 
-unsafe impl<T> DescriptorSet for T
-where
-    T: SafeDeref,
-    T::Target: DescriptorSet,
-{
-    #[inline]
-    fn inner(&self) -> &UnsafeDescriptorSet {
-        (**self).inner()
+pub(crate) struct DescriptorSetInner {
+    layout: Arc<DescriptorSetLayout>,
+    variable_descriptor_count: u32,
+    resources: DescriptorSetResources,
+}
+
+impl DescriptorSetInner {
+    pub(crate) fn new(
+        handle: ash::vk::DescriptorSet,
+        layout: Arc<DescriptorSetLayout>,
+        variable_descriptor_count: u32,
+        descriptor_writes: impl IntoIterator<Item = WriteDescriptorSet>,
+    ) -> Result<Self, DescriptorSetUpdateError> {
+        assert!(
+            !layout.push_descriptor(),
+            "the provided descriptor set layout is for push descriptors, and cannot be used to \
+            build a descriptor set object",
+        );
+
+        let max_count = layout.variable_descriptor_count();
+
+        assert!(
+            variable_descriptor_count <= max_count,
+            "the provided variable_descriptor_count ({}) is greater than the maximum number of \
+            variable count descriptors in the layout ({})",
+            variable_descriptor_count,
+            max_count,
+        );
+
+        let mut resources = DescriptorSetResources::new(&layout, variable_descriptor_count);
+
+        let descriptor_writes = descriptor_writes.into_iter();
+        let (lower_size_bound, _) = descriptor_writes.size_hint();
+        let mut descriptor_write_info: SmallVec<[_; 8]> = SmallVec::with_capacity(lower_size_bound);
+        let mut write_descriptor_set: SmallVec<[_; 8]> = SmallVec::with_capacity(lower_size_bound);
+
+        for write in descriptor_writes {
+            let layout_binding =
+                check_descriptor_write(&write, &layout, variable_descriptor_count)?;
+
+            resources.update(&write);
+            descriptor_write_info.push(write.to_vulkan_info(layout_binding.descriptor_type));
+            write_descriptor_set.push(write.to_vulkan(handle, layout_binding.descriptor_type));
+        }
+
+        if !write_descriptor_set.is_empty() {
+            for (info, write) in descriptor_write_info
+                .iter()
+                .zip(write_descriptor_set.iter_mut())
+            {
+                match info {
+                    DescriptorWriteInfo::Image(info) => {
+                        write.descriptor_count = info.len() as u32;
+                        write.p_image_info = info.as_ptr();
+                    }
+                    DescriptorWriteInfo::Buffer(info) => {
+                        write.descriptor_count = info.len() as u32;
+                        write.p_buffer_info = info.as_ptr();
+                    }
+                    DescriptorWriteInfo::BufferView(info) => {
+                        write.descriptor_count = info.len() as u32;
+                        write.p_texel_buffer_view = info.as_ptr();
+                    }
+                }
+            }
+        }
+
+        unsafe {
+            let fns = layout.device().fns();
+
+            (fns.v1_0.update_descriptor_sets)(
+                layout.device().handle(),
+                write_descriptor_set.len() as u32,
+                write_descriptor_set.as_ptr(),
+                0,
+                ptr::null(),
+            );
+        }
+
+        Ok(DescriptorSetInner {
+            layout,
+            variable_descriptor_count,
+            resources,
+        })
     }
 
-    #[inline]
-    fn layout(&self) -> &Arc<DescriptorSetLayout> {
-        (**self).layout()
+    pub(crate) fn layout(&self) -> &Arc<DescriptorSetLayout> {
+        &self.layout
     }
 
-    #[inline]
-    fn num_buffers(&self) -> usize {
-        (**self).num_buffers()
+    pub(crate) fn resources(&self) -> &DescriptorSetResources {
+        &self.resources
     }
+}
 
+/// The resources that are bound to a descriptor set.
+#[derive(Clone)]
+pub struct DescriptorSetResources {
+    binding_resources: HashMap<u32, DescriptorBindingResources>,
+}
+
+impl DescriptorSetResources {
+    /// Creates a new `DescriptorSetResources` matching the provided descriptor set layout, and
+    /// all descriptors set to `None`.
     #[inline]
-    fn buffer(&self, index: usize) -> Option<(&dyn BufferAccess, u32)> {
-        (**self).buffer(index)
+    pub fn new(layout: &DescriptorSetLayout, variable_descriptor_count: u32) -> Self {
+        assert!(variable_descriptor_count <= layout.variable_descriptor_count());
+
+        let binding_resources = layout
+            .bindings()
+            .iter()
+            .map(|(&binding_num, binding)| {
+                let count = if binding.variable_descriptor_count {
+                    variable_descriptor_count
+                } else {
+                    binding.descriptor_count
+                } as usize;
+
+                let binding_resources = match binding.descriptor_type {
+                    DescriptorType::UniformBuffer
+                    | DescriptorType::StorageBuffer
+                    | DescriptorType::UniformBufferDynamic
+                    | DescriptorType::StorageBufferDynamic => {
+                        DescriptorBindingResources::Buffer(smallvec![None; count])
+                    }
+                    DescriptorType::UniformTexelBuffer | DescriptorType::StorageTexelBuffer => {
+                        DescriptorBindingResources::BufferView(smallvec![None; count])
+                    }
+                    DescriptorType::SampledImage
+                    | DescriptorType::StorageImage
+                    | DescriptorType::InputAttachment => {
+                        DescriptorBindingResources::ImageView(smallvec![None; count])
+                    }
+                    DescriptorType::CombinedImageSampler => {
+                        if binding.immutable_samplers.is_empty() {
+                            DescriptorBindingResources::ImageViewSampler(smallvec![None; count])
+                        } else {
+                            DescriptorBindingResources::ImageView(smallvec![None; count])
+                        }
+                    }
+                    DescriptorType::Sampler => {
+                        if binding.immutable_samplers.is_empty() {
+                            DescriptorBindingResources::Sampler(smallvec![None; count])
+                        } else if layout.push_descriptor() {
+                            // For push descriptors, no resource is written by default, this needs
+                            // to be done explicitly via a dummy write.
+                            DescriptorBindingResources::None(smallvec![None; count])
+                        } else {
+                            // For regular descriptor sets, all descriptors are considered valid
+                            // from the start.
+                            DescriptorBindingResources::None(smallvec![Some(()); count])
+                        }
+                    }
+                };
+                (binding_num, binding_resources)
+            })
+            .collect();
+
+        Self { binding_resources }
     }
 
+    /// Applies a descriptor write to the resources.
+    ///
+    /// # Panics
+    ///
+    /// - Panics if the binding number of a write does not exist in the resources.
+    /// - See also [`DescriptorBindingResources::update`].
     #[inline]
-    fn num_images(&self) -> usize {
-        (**self).num_images()
+    pub fn update(&mut self, write: &WriteDescriptorSet) {
+        self.binding_resources
+            .get_mut(&write.binding())
+            .expect("descriptor write has invalid binding number")
+            .update(write)
     }
 
+    /// Returns a reference to the bound resources for `binding`. Returns `None` if the binding
+    /// doesn't exist.
     #[inline]
-    fn image(&self, index: usize) -> Option<(&dyn ImageViewAbstract, u32)> {
-        (**self).image(index)
+    pub fn binding(&self, binding: u32) -> Option<&DescriptorBindingResources> {
+        self.binding_resources.get(&binding)
     }
 }
 
-impl PartialEq for dyn DescriptorSet + Send + Sync {
-    #[inline]
-    fn eq(&self, other: &Self) -> bool {
-        self.inner().internal_object() == other.inner().internal_object()
-            && self.device() == other.device()
-    }
+/// The resources that are bound to a single descriptor set binding.
+#[derive(Clone)]
+pub enum DescriptorBindingResources {
+    None(Elements<()>),
+    Buffer(Elements<(Subbuffer<[u8]>, Range<DeviceSize>)>),
+    BufferView(Elements<Arc<BufferView>>),
+    ImageView(Elements<Arc<dyn ImageViewAbstract>>),
+    ImageViewSampler(Elements<(Arc<dyn ImageViewAbstract>, Arc<Sampler>)>),
+    Sampler(Elements<Arc<Sampler>>),
 }
 
-impl Eq for dyn DescriptorSet + Send + Sync {}
+type Elements<T> = SmallVec<[Option<T>; 1]>;
 
-impl Hash for dyn DescriptorSet + Send + Sync {
+impl DescriptorBindingResources {
+    /// Applies a descriptor write to the resources.
+    ///
+    /// # Panics
+    ///
+    /// - Panics if the resource types do not match.
+    /// - Panics if the write goes out of bounds.
     #[inline]
-    fn hash<H: Hasher>(&self, state: &mut H) {
-        self.inner().internal_object().hash(state);
-        self.device().hash(state);
+    pub fn update(&mut self, write: &WriteDescriptorSet) {
+        fn write_resources<T: Clone>(first: usize, resources: &mut [Option<T>], elements: &[T]) {
+            resources
+                .get_mut(first..first + elements.len())
+                .expect("descriptor write for binding out of bounds")
+                .iter_mut()
+                .zip(elements)
+                .for_each(|(resource, element)| {
+                    *resource = Some(element.clone());
+                });
+        }
+
+        let first = write.first_array_element() as usize;
+
+        match (self, write.elements()) {
+            (
+                DescriptorBindingResources::None(resources),
+                WriteDescriptorSetElements::None(num_elements),
+            ) => {
+                resources
+                    .get_mut(first..first + *num_elements as usize)
+                    .expect("descriptor write for binding out of bounds")
+                    .iter_mut()
+                    .for_each(|resource| {
+                        *resource = Some(());
+                    });
+            }
+            (
+                DescriptorBindingResources::Buffer(resources),
+                WriteDescriptorSetElements::Buffer(elements),
+            ) => write_resources(first, resources, elements),
+            (
+                DescriptorBindingResources::BufferView(resources),
+                WriteDescriptorSetElements::BufferView(elements),
+            ) => write_resources(first, resources, elements),
+            (
+                DescriptorBindingResources::ImageView(resources),
+                WriteDescriptorSetElements::ImageView(elements),
+            ) => write_resources(first, resources, elements),
+            (
+                DescriptorBindingResources::ImageViewSampler(resources),
+                WriteDescriptorSetElements::ImageViewSampler(elements),
+            ) => write_resources(first, resources, elements),
+            (
+                DescriptorBindingResources::Sampler(resources),
+                WriteDescriptorSetElements::Sampler(elements),
+            ) => write_resources(first, resources, elements),
+            _ => panic!(
+                "descriptor write for binding {} has wrong resource type",
+                write.binding(),
+            ),
+        }
     }
 }
 
+#[derive(Clone)]
 pub struct DescriptorSetWithOffsets {
-    descriptor_set: Box<dyn DescriptorSet + Send + Sync>,
+    descriptor_set: Arc<dyn DescriptorSet>,
     dynamic_offsets: SmallVec<[u32; 4]>,
 }
 
 impl DescriptorSetWithOffsets {
-    #[inline]
-    pub fn new<S, O>(descriptor_set: S, dynamic_offsets: O) -> Self
-    where
-        S: DescriptorSet + Send + Sync + 'static,
-        O: IntoIterator<Item = u32>,
-    {
-        let dynamic_offsets: SmallVec<_> = dynamic_offsets.into_iter().collect();
-        let layout = descriptor_set.layout();
-        let properties = layout.device().physical_device().properties();
-        let min_uniform_off_align = properties.min_uniform_buffer_offset_alignment as u32;
-        let min_storage_off_align = properties.min_storage_buffer_offset_alignment as u32;
-        let mut dynamic_offset_index = 0;
-
-        // Ensure that the number of dynamic_offsets is correct and that each
-        // dynamic offset is a multiple of the minimum offset alignment specified
-        // by the physical device.
-        for desc in layout.desc().bindings() {
-            let desc = desc.as_ref().unwrap();
-            if let DescriptorDescTy::Buffer(DescriptorBufferDesc {
-                dynamic: Some(true),
-                storage,
-            }) = desc.ty
-            {
-                // Don't check alignment if there are not enough offsets anyway
-                if dynamic_offsets.len() > dynamic_offset_index {
-                    if storage {
-                        assert!(
-                            dynamic_offsets[dynamic_offset_index] % min_storage_off_align == 0,
-                            "Dynamic storage buffer offset must be a multiple of min_storage_buffer_offset_alignment: got {}, expected a multiple of {}",
-                            dynamic_offsets[dynamic_offset_index],
-                            min_storage_off_align
-                        );
-                    } else {
-                        assert!(
-                            dynamic_offsets[dynamic_offset_index] % min_uniform_off_align == 0,
-                            "Dynamic uniform buffer offset must be a multiple of min_uniform_buffer_offset_alignment: got {}, expected a multiple of {}",
-                            dynamic_offsets[dynamic_offset_index],
-                            min_uniform_off_align
-                        );
-                    }
-                }
-                dynamic_offset_index += 1;
-            }
-        }
-
-        assert!(
-            !(dynamic_offsets.len() < dynamic_offset_index),
-            "Too few dynamic offsets: got {}, expected {}",
-            dynamic_offsets.len(),
-            dynamic_offset_index
-        );
-        assert!(
-            !(dynamic_offsets.len() > dynamic_offset_index),
-            "Too many dynamic offsets: got {}, expected {}",
-            dynamic_offsets.len(),
-            dynamic_offset_index
-        );
-
-        DescriptorSetWithOffsets {
-            descriptor_set: Box::new(descriptor_set),
-            dynamic_offsets,
+    pub fn new(
+        descriptor_set: Arc<dyn DescriptorSet>,
+        dynamic_offsets: impl IntoIterator<Item = u32>,
+    ) -> Self {
+        Self {
+            descriptor_set,
+            dynamic_offsets: dynamic_offsets.into_iter().collect(),
         }
     }
 
     #[inline]
-    pub fn as_ref(&self) -> (&dyn DescriptorSet, &[u32]) {
+    pub fn as_ref(&self) -> (&Arc<dyn DescriptorSet>, &[u32]) {
         (&self.descriptor_set, &self.dynamic_offsets)
     }
 
     #[inline]
-    pub fn into_tuple(
-        self,
-    ) -> (
-        Box<dyn DescriptorSet + Send + Sync>,
-        impl ExactSizeIterator<Item = u32>,
-    ) {
+    pub fn into_tuple(self) -> (Arc<dyn DescriptorSet>, impl ExactSizeIterator<Item = u32>) {
         (self.descriptor_set, self.dynamic_offsets.into_iter())
     }
 }
 
-impl<S> From<S> for DescriptorSetWithOffsets
+impl<S> From<Arc<S>> for DescriptorSetWithOffsets
 where
-    S: DescriptorSet + Send + Sync + 'static,
+    S: DescriptorSet + 'static,
 {
-    #[inline]
-    fn from(descriptor_set: S) -> Self {
-        Self::new(descriptor_set, std::iter::empty())
+    fn from(descriptor_set: Arc<S>) -> Self {
+        DescriptorSetWithOffsets::new(descriptor_set, std::iter::empty())
+    }
+}
+
+#[derive(Clone, Copy, Debug)]
+pub enum DescriptorSetCreationError {
+    DescriptorSetUpdateError(DescriptorSetUpdateError),
+    OomError(OomError),
+}
+
+impl Error for DescriptorSetCreationError {
+    fn source(&self) -> Option<&(dyn Error + 'static)> {
+        match self {
+            Self::DescriptorSetUpdateError(err) => Some(err),
+            Self::OomError(err) => Some(err),
+        }
+    }
+}
+
+impl Display for DescriptorSetCreationError {
+    fn fmt(&self, f: &mut Formatter<'_>) -> Result<(), FmtError> {
+        match self {
+            Self::DescriptorSetUpdateError(_) => {
+                write!(f, "an error occurred while updating the descriptor set")
+            }
+            Self::OomError(_) => write!(f, "out of memory"),
+        }
+    }
+}
+
+impl From<DescriptorSetUpdateError> for DescriptorSetCreationError {
+    fn from(err: DescriptorSetUpdateError) -> Self {
+        Self::DescriptorSetUpdateError(err)
+    }
+}
+
+impl From<OomError> for DescriptorSetCreationError {
+    fn from(err: OomError) -> Self {
+        Self::OomError(err)
     }
 }