From a084afc7ae001bd7d08d0a9fd7a6272f4f56d750 Mon Sep 17 00:00:00 2001
From: Josh Simmons <josh@nega.tv>
Date: Sun, 9 Jun 2024 23:04:17 +0200
Subject: [PATCH] narcissus-gpu: Expose buffer device addresses

---
 .../narcissus-gpu/src/backend/vulkan/mod.rs   | 44 +++++++++++++++++--
 engine/narcissus-gpu/src/lib.rs               |  3 ++
 engine/narcissus-gpu/src/mapped_buffer.rs     |  1 +
 3 files changed, 45 insertions(+), 3 deletions(-)

diff --git a/engine/narcissus-gpu/src/backend/vulkan/mod.rs b/engine/narcissus-gpu/src/backend/vulkan/mod.rs
index f0944c6..01df29a 100644
--- a/engine/narcissus-gpu/src/backend/vulkan/mod.rs
+++ b/engine/narcissus-gpu/src/backend/vulkan/mod.rs
@@ -184,7 +184,7 @@ fn vulkan_shader_module(
 struct VulkanBuffer {
     memory: VulkanMemory,
     buffer: vk::Buffer,
-    _address: vk::DeviceAddress,
+    address: vk::DeviceAddress,
     map_count: u64,
 }
 
@@ -260,6 +260,7 @@ struct VulkanBoundPipeline {
 #[derive(Clone)]
 struct VulkanTransientBuffer {
     buffer: vk::Buffer,
+    address: u64,
     memory: VulkanMemory,
 }
 
@@ -986,7 +987,7 @@ impl Device for VulkanDevice {
         let handle = self.buffer_pool.lock().insert(VulkanBuffer {
             memory,
             buffer,
-            _address: address,
+            address,
             map_count: 0,
         });
 
@@ -2693,6 +2694,17 @@ impl Device for VulkanDevice {
         self.wsi_end_frame(self.frame_mut(&mut frame));
         self.frame_counter.release(frame);
     }
+
+    fn get_buffer_address(&self, buffer: BufferArg) -> u64 {
+        let buffer = match buffer {
+            BufferArg::Unmanaged(buffer) => buffer.0,
+            BufferArg::Persistent(buffer) => buffer.buffer.0,
+            BufferArg::Transient(buffer) => return buffer.address,
+        };
+        let buffer_pool = self.buffer_pool.lock();
+        let buffer = buffer_pool.get(buffer).unwrap();
+        buffer.address
+    }
 }
 
 impl VulkanDevice {
@@ -2740,6 +2752,16 @@ impl VulkanDevice {
                 )
             };
 
+            let address = unsafe {
+                self.device_fn.get_buffer_device_address(
+                    self.device,
+                    &vk::BufferDeviceAddressInfo {
+                        buffer,
+                        ..default()
+                    },
+                )
+            };
+
             let ptr = NonNull::new(memory.mapped_ptr()).unwrap();
 
             frame.destroyed_buffers.lock().push_back(buffer);
@@ -2749,6 +2771,7 @@ impl VulkanDevice {
                 ptr,
                 len: size as usize,
                 buffer: buffer.as_raw(),
+                address,
                 offset: 0,
                 phantom: PhantomData,
             };
@@ -2816,6 +2839,7 @@ impl VulkanDevice {
             .unwrap(),
             len: size as usize,
             buffer: current.buffer.as_raw(),
+            address: current.address + allocator.offset,
             offset: allocator.offset,
             phantom: PhantomData,
         }
@@ -2874,7 +2898,21 @@ impl VulkanDevice {
             )
         };
 
-        VulkanTransientBuffer { buffer, memory }
+        let address = unsafe {
+            self.device_fn.get_buffer_device_address(
+                self.device,
+                &vk::BufferDeviceAddressInfo {
+                    buffer,
+                    ..default()
+                },
+            )
+        };
+
+        VulkanTransientBuffer {
+            buffer,
+            address,
+            memory,
+        }
     }
 
     fn unwrap_buffer_arg(&self, buffer_arg: &BufferArg) -> (vk::Buffer, u64, u64) {
diff --git a/engine/narcissus-gpu/src/lib.rs b/engine/narcissus-gpu/src/lib.rs
index cc73295..9b0d2ee 100644
--- a/engine/narcissus-gpu/src/lib.rs
+++ b/engine/narcissus-gpu/src/lib.rs
@@ -788,6 +788,9 @@ pub trait Device {
     fn debug_name_buffer(&self, buffer: BufferArg, label_name: &str);
     fn debug_name_image(&self, image: Image, label_name: &str);
 
+    // Danger Zone
+    fn get_buffer_address(&self, buffer: BufferArg) -> u64;
+
     fn destroy_buffer(&self, frame: &Frame, buffer: Buffer);
     fn destroy_persistent_buffer(&self, frame: &Frame, buffer: PersistentBuffer);
     fn destroy_image(&self, frame: &Frame, image: Image);
diff --git a/engine/narcissus-gpu/src/mapped_buffer.rs b/engine/narcissus-gpu/src/mapped_buffer.rs
index 21ba0d7..e3e8e84 100644
--- a/engine/narcissus-gpu/src/mapped_buffer.rs
+++ b/engine/narcissus-gpu/src/mapped_buffer.rs
@@ -85,6 +85,7 @@ pub struct TransientBuffer<'a> {
     pub(crate) offset: u64,
     pub(crate) len: usize,
     pub(crate) buffer: u64,
+    pub(crate) address: u64,
     pub(crate) phantom: PhantomData<&'a u8>,
 }
 
-- 
2.49.0