From b57d86c1bcc638a5148c1ec5bf886a7cbf838754 Mon Sep 17 00:00:00 2001
From: Julian Orth <ju.orth@gmail.com>
Date: Sat, 7 Sep 2024 17:42:29 +0200
Subject: [PATCH] vulkan: add async staging buffer allocation

---
 src/gfx_apis/vulkan/allocator.rs |  1 -
 src/gfx_apis/vulkan/renderer.rs  |  1 -
 src/gfx_apis/vulkan/staging.rs   | 91 +++++++++++++++++++++++++-------
 src/utils/on_drop.rs             |  2 -
 4 files changed, 72 insertions(+), 23 deletions(-)

diff --git a/src/gfx_apis/vulkan/allocator.rs b/src/gfx_apis/vulkan/allocator.rs
index 77ffced0..6cff0834 100644
--- a/src/gfx_apis/vulkan/allocator.rs
+++ b/src/gfx_apis/vulkan/allocator.rs
@@ -199,7 +199,6 @@ impl VulkanAllocator {
 }
 
 impl VulkanThreadedAllocator {
-    #[expect(dead_code)]
     pub fn async_alloc(
         self: &Rc<Self>,
         renderer: &Rc<VulkanRenderer>,
diff --git a/src/gfx_apis/vulkan/renderer.rs b/src/gfx_apis/vulkan/renderer.rs
index 05982da7..71877a7f 100644
--- a/src/gfx_apis/vulkan/renderer.rs
+++ b/src/gfx_apis/vulkan/renderer.rs
@@ -81,7 +81,6 @@ pub struct VulkanRenderer {
     pub(super) tex_descriptor_set_layout: Rc<VulkanDescriptorSetLayout>,
     pub(super) defunct: Cell<bool>,
     pub(super) pending_cpu_jobs: CopyHashMap<u64, PendingJob>,
-    #[expect(dead_code)]
     pub(super) shm_allocator: Rc<VulkanThreadedAllocator>,
 }
 
diff --git a/src/gfx_apis/vulkan/staging.rs b/src/gfx_apis/vulkan/staging.rs
index d09d705a..e2d6b39f 100644
--- a/src/gfx_apis/vulkan/staging.rs
+++ b/src/gfx_apis/vulkan/staging.rs
@@ -1,11 +1,13 @@
 use {
     crate::{
+        cpu_worker::CpuWorker,
         gfx_apis::vulkan::{
             allocator::{VulkanAllocation, VulkanAllocator},
             device::VulkanDevice,
+            renderer::VulkanRenderer,
             VulkanError,
         },
-        utils::on_drop::OnDrop,
+        utils::on_drop::{OnDrop, OnDrop2},
     },
     ash::vk::{Buffer, BufferCreateInfo, BufferUsageFlags, MappedMemoryRange},
     gpu_alloc::UsageFlags,
@@ -28,24 +30,8 @@ impl VulkanDevice {
         download: bool,
         transient: bool,
     ) -> Result<VulkanStagingBuffer, VulkanError> {
-        let mut vk_usage = BufferUsageFlags::empty();
-        let mut usage = UsageFlags::empty();
-        if upload {
-            vk_usage |= BufferUsageFlags::TRANSFER_SRC;
-            usage |= UsageFlags::UPLOAD;
-        }
-        if download {
-            vk_usage |= BufferUsageFlags::TRANSFER_DST;
-            usage |= UsageFlags::DOWNLOAD;
-        }
-        if transient {
-            usage |= UsageFlags::TRANSIENT;
-        }
-        let buffer = {
-            let create_info = BufferCreateInfo::default().size(size).usage(vk_usage);
-            let buffer = unsafe { self.device.create_buffer(&create_info, None) };
-            buffer.map_err(VulkanError::CreateBuffer)?
-        };
+        let (vk_usage, usage) = get_usage(upload, download, transient);
+        let buffer = self.create_buffer(size, vk_usage)?;
         let destroy_buffer = OnDrop(|| unsafe { self.device.destroy_buffer(buffer, None) });
         let memory_requirements = unsafe { self.device.get_buffer_memory_requirements(buffer) };
         let allocation = allocator.alloc(&memory_requirements, usage, true)?;
@@ -64,6 +50,73 @@ impl VulkanDevice {
             size,
         })
     }
+
+    #[expect(dead_code)]
+    pub(super) fn create_shm_staging(
+        self: &Rc<Self>,
+        renderer: &Rc<VulkanRenderer>,
+        cpu: &Rc<CpuWorker>,
+        size: u64,
+        upload: bool,
+        download: bool,
+        cb: impl FnOnce(Result<VulkanStagingBuffer, VulkanError>) + 'static,
+    ) -> Result<(), VulkanError> {
+        let (vk_usage, usage) = get_usage(upload, download, false);
+        let buffer = self.create_buffer(size, vk_usage)?;
+        let memory_requirements = unsafe { self.device.get_buffer_memory_requirements(buffer) };
+        let slf = self.clone();
+        let destroy_buffer =
+            OnDrop2::new(move || unsafe { slf.device.destroy_buffer(buffer, None) });
+        let slf = self.clone();
+        let finish_allocation = move |res| {
+            let allocation: VulkanAllocation = res?;
+            {
+                let res = unsafe {
+                    slf.device
+                        .bind_buffer_memory(buffer, allocation.memory, allocation.offset)
+                };
+                res.map_err(VulkanError::BindBufferMemory)?;
+            }
+            destroy_buffer.forget();
+            Ok(VulkanStagingBuffer {
+                device: slf.clone(),
+                allocation,
+                buffer,
+                size,
+            })
+        };
+        renderer.shm_allocator.async_alloc(
+            renderer,
+            cpu,
+            memory_requirements,
+            usage,
+            true,
+            move |res| cb(finish_allocation(res)),
+        )
+    }
+
+    fn create_buffer(&self, size: u64, usage: BufferUsageFlags) -> Result<Buffer, VulkanError> {
+        let create_info = BufferCreateInfo::default().size(size).usage(usage);
+        let buffer = unsafe { self.device.create_buffer(&create_info, None) };
+        buffer.map_err(VulkanError::CreateBuffer)
+    }
+}
+
+fn get_usage(upload: bool, download: bool, transient: bool) -> (BufferUsageFlags, UsageFlags) {
+    let mut vk_usage = BufferUsageFlags::empty();
+    let mut usage = UsageFlags::empty();
+    if upload {
+        vk_usage |= BufferUsageFlags::TRANSFER_SRC;
+        usage |= UsageFlags::UPLOAD;
+    }
+    if download {
+        vk_usage |= BufferUsageFlags::TRANSFER_DST;
+        usage |= UsageFlags::DOWNLOAD;
+    }
+    if transient {
+        usage |= UsageFlags::TRANSIENT;
+    }
+    (vk_usage, usage)
 }
 
 impl VulkanStagingBuffer {
diff --git a/src/utils/on_drop.rs b/src/utils/on_drop.rs
index 15557f0a..f7d7644d 100644
--- a/src/utils/on_drop.rs
+++ b/src/utils/on_drop.rs
@@ -24,14 +24,12 @@ where
 }
 
 impl<F: FnOnce()> OnDrop2<F> {
-    #[expect(dead_code)]
     pub fn new(f: F) -> Self {
         Self {
             f: ManuallyDrop::new(f),
         }
     }
 
-    #[expect(dead_code)]
     pub fn forget(mut self) {
         unsafe {
             ManuallyDrop::drop(&mut self.f);