From 4f1638904c05243a1f12d9150a1df06b5880dbb8 Mon Sep 17 00:00:00 2001
From: Stephen Seo <seo.disparate@gmail.com>
Date: Mon, 18 Mar 2024 15:04:02 +0900
Subject: [PATCH] Use a macro to create cleanup struct helper

---
 src/helper.rs | 27 +++++++++++++++++++++++++++
 src/main.rs   | 31 +++++++++++++------------------
 2 files changed, 40 insertions(+), 18 deletions(-)
 create mode 100644 src/helper.rs

diff --git a/src/helper.rs b/src/helper.rs
new file mode 100644
index 0000000..a318f59
--- /dev/null
+++ b/src/helper.rs
@@ -0,0 +1,27 @@
+#[macro_export]
+macro_rules! cleanup_func {
+    (func: $cleanup_fn:expr, name: $name:ident, hold_name: $hold_name:ident) => {
+        struct $name <T: Fn() -> ()> {
+            func: T,
+        }
+        impl<T> Drop for $name <T> where T: Fn() -> () {
+            fn drop(&mut self) {
+                (self.func)();
+            }
+        }
+        impl<T> $name <T> where T: Fn() -> () {
+            fn new(func: T) -> Self where T: Fn() -> () {
+                Self {
+                    func,
+                }
+            }
+        }
+        $hold_name = $name::new($cleanup_fn);
+    };
+    (func: $cleanup_fn:expr, name: $name:ident, hold_name: $hold_name:ident, $(var_pair: $orig_var:expr, $new_var:ident),*) => {
+        {
+            $(let $new_var = $orig_var;)*
+            cleanup_func!(func: $cleanup_fn, name: $name, hold_name: $hold_name);
+        }
+    }
+}
diff --git a/src/main.rs b/src/main.rs
index 7a452e2..3f93fc7 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,4 +1,5 @@
 mod ffi;
+mod helper;
 mod math3d;
 
 use std::collections::HashSet;
@@ -1619,24 +1620,18 @@ impl VulkanApp {
                 | ffi::VkMemoryPropertyFlagBits_VK_MEMORY_PROPERTY_HOST_COHERENT_BIT,
         )?;
 
-        struct CleanupStaging {
-            device: ffi::VkDevice,
-            staging_buf: ffi::VkBuffer,
-            staging_buf_mem: ffi::VkDeviceMemory,
-        }
-        impl Drop for CleanupStaging {
-            fn drop(&mut self) {
-                unsafe {
-                    ffi::vkDestroyBuffer(self.device, self.staging_buf, std::ptr::null());
-                    ffi::vkFreeMemory(self.device, self.staging_buf_mem, std::ptr::null());
-                }
-            }
-        }
-        let _cleanup_staging = CleanupStaging {
-            device: self.device,
-            staging_buf: staging_buffer,
-            staging_buf_mem: staging_buffer_mem,
-        };
+        let _inst;
+        cleanup_func!(
+            func: move || unsafe {
+                ffi::vkDestroyBuffer(device_copy, staging_buf_copy, std::ptr::null());
+                ffi::vkFreeMemory(device_copy, staging_buf_mem_copy, std::ptr::null());
+            },
+            name: CleanupStaging,
+            hold_name: _inst,
+            var_pair: self.device, device_copy,
+            var_pair: staging_buffer, staging_buf_copy,
+            var_pair: staging_buffer_mem, staging_buf_mem_copy
+        );
 
         let mut data_ptr: *mut c_void = unsafe { std::mem::zeroed() };
         unsafe {