From 7769fb15ee2a64024edb9168be6cd30190d93502 Mon Sep 17 00:00:00 2001
From: antonino <antomani103@gmail.com>
Date: Sun, 30 Jul 2023 16:12:14 +0200
Subject: [PATCH] wip: add: Abstraction for bindless resources

---
 src/resource/bindless.rs | 166 +++++++++++++++++++++++++++++++++++++++
 src/resource/mod.rs      |   1 +
 2 files changed, 167 insertions(+)
 create mode 100644 src/resource/bindless.rs

diff --git a/src/resource/bindless.rs b/src/resource/bindless.rs
new file mode 100644
index 0000000..9a3d785
--- /dev/null
+++ b/src/resource/bindless.rs
@@ -0,0 +1,166 @@
+use std::sync::{Arc, Mutex};
+
+use anyhow::Result;
+
+use ash::vk;
+
+use crate::{pipeline::set_layout::DescriptorSetLayoutCreateInfo, util::cache::Resource};
+
+static MAX_BINDLESS_COUNT: u32 = 4096;
+
+pub trait Bindlessable {
+    fn descriptor_type() -> vk::DescriptorType;
+    fn resource_binding(binding: u32) -> vk::DescriptorSetLayoutBinding {
+        vk::DescriptorSetLayoutBinding {
+            binding,
+            descriptor_type: Self::descriptor_type(),
+            descriptor_count: MAX_BINDLESS_COUNT,
+            stage_flags: vk::ShaderStageFlags::ALL,
+            p_immutable_samplers: std::ptr::null(),
+        }
+    }
+    fn descriptor_info(&self) -> vk::DescriptorImageInfo;
+
+    fn into_bindless(self, pool: &mut BindlessPool<Self>) -> BindlessResource<Self>
+    where
+        Self: Sized
+    {
+        pool.alloc(self)
+    }
+}
+
+impl Bindlessable for crate::image::Image {
+    fn descriptor_type() -> vk::DescriptorType {
+        vk::DescriptorType::STORAGE_IMAGE
+    }
+
+    fn descriptor_info(&self) -> vk::DescriptorImageInfo {
+        todo!()
+    }
+}
+
+pub struct CombinedImageSampler {
+    sampler: crate::sampler::Sampler,
+    image_view: crate::image::ImageView,
+    image_layout: Option<vk::ImageLayout>,
+}
+
+impl Bindlessable for CombinedImageSampler {
+    fn descriptor_type() -> vk::DescriptorType {
+        vk::DescriptorType::COMBINED_IMAGE_SAMPLER
+    }
+
+    fn descriptor_info(&self) -> vk::DescriptorImageInfo {
+        vk::DescriptorImageInfo {
+            sampler: unsafe { self.sampler.handle() },
+            image_view: unsafe { self.image_view.handle() },
+            image_layout: self.image_layout.unwrap_or(vk::ImageLayout::SHADER_READ_ONLY_OPTIMAL),
+        }
+    }
+}
+
+pub struct BindlessResource<R: Bindlessable> {
+    key: u32,
+    pool: BindlessPool<R>,
+}
+
+impl<R: Bindlessable> Drop for BindlessResource<R> {
+    fn drop(&mut self) {
+        self.pool.with(|p| p.take(self.key));
+    }
+}
+
+impl<R: Bindlessable> BindlessResource<R> {
+    pub fn index(&self) -> u32 {
+        self.key
+    }
+}
+
+struct BindlessPoolInner<R> {
+    items: Vec<Option<R>>,
+    free: Vec<u32>,
+    descriptor_set: crate::DescriptorSet,
+}
+
+impl<R: Bindlessable> BindlessPoolInner<R> {
+    fn update_descriptor_set(&mut self, i: u32, r: &R) {
+        let vk_writes = vk::WriteDescriptorSet {
+                s_type: vk::StructureType::WRITE_DESCRIPTOR_SET,
+                p_next: std::ptr::null(),
+                dst_set: self.descriptor_set.handle,
+                dst_binding: 0,
+                dst_array_element: i,
+                descriptor_count: 1,
+                descriptor_type: R::descriptor_type(),
+                p_image_info: &r.descriptor_info() as *const _,
+                p_buffer_info: std::ptr::null(),
+                p_texel_buffer_view: std::ptr::null(),
+            };
+
+        unsafe {
+            self.descriptor_set.device.update_descriptor_sets(&[vk_writes], &[]);
+        }
+    }
+
+
+    fn take(&mut self, key: u32) -> Option<R> {
+        if key <= self.items.len() as _ {
+            None
+        } else {
+            self.items[key as usize].take().and_then(|ob| {
+                self.free.push(key);
+                Some(ob)
+            })
+        }
+    }
+}
+
+pub struct BindlessPool<R> {
+    inner: Arc<Mutex<BindlessPoolInner<R>>>,
+}
+
+impl<P: Bindlessable> BindlessPool<P> {
+    fn with<F: FnOnce(&mut BindlessPoolInner<P>) -> R, R>(&self, f: F) -> R {
+        let mut inner = self.inner.lock().unwrap();
+        f(&mut inner)
+    }
+
+    pub fn alloc(&mut self, item: P) -> BindlessResource<P> {
+        self.with(|p| {
+            let key = p.free
+                .pop()
+                .unwrap_or_else(|| {
+                    p.items.push(None);
+                    p.items.len() as u32 - 1
+                });
+            p.update_descriptor_set(key, &item);
+            p.items[key as usize] = Some(item);
+            BindlessResource {
+                key,
+                pool: Self { inner: self.inner.clone() }
+            }
+        })
+    }
+}
+
+impl<P: Bindlessable> BindlessPool<P> {
+    fn new(device: crate::Device, pool: vk::DescriptorPool) -> Result<Self> {
+        let dsl_info = DescriptorSetLayoutCreateInfo {
+            bindings: vec![
+                P::resource_binding(0)
+            ],
+            persistent: true,
+            flags: vec![vk::DescriptorBindingFlags::UPDATE_AFTER_BIND, vk::DescriptorBindingFlags::PARTIALLY_BOUND],
+        };
+        let dsl = crate::pipeline::set_layout::DescriptorSetLayout::create(device.clone(), &dsl_info, ())?;
+        let descriptor_set = unsafe { crate::DescriptorSet::new_uninitialized(device, dsl.handle(), pool)? };
+
+        let inner = BindlessPoolInner {
+            items: vec![],
+            free: vec![],
+            descriptor_set,
+        };
+
+        Ok(Self { inner: Arc::new(Mutex::new(inner)) })
+    }
+}
diff --git a/src/resource/mod.rs b/src/resource/mod.rs
index 0aae812..786dc89 100644
--- a/src/resource/mod.rs
+++ b/src/resource/mod.rs
@@ -6,3 +6,4 @@ pub mod pool;
 pub mod query_pool;
 pub mod raytracing;
 pub mod sampler;
+pub mod bindless;