From a334604ac48ea12aaa362d139f7a5ced8f36bcba Mon Sep 17 00:00:00 2001
From: John Wells <john@attackgoat.com>
Date: Sat, 27 Jan 2024 23:12:39 -0500
Subject: [PATCH] Add basic pool memory management functions

---
 src/pool/hash.rs | 55 +++++++++++++++++++++++++++++++++++++++++++++++-
 src/pool/lazy.rs | 52 ++++++++++++++++++++++++++++++++++++++++-----
 2 files changed, 101 insertions(+), 6 deletions(-)

diff --git a/src/pool/hash.rs b/src/pool/hash.rs
index e828017..401257a 100644
--- a/src/pool/hash.rs
+++ b/src/pool/hash.rs
@@ -16,6 +16,7 @@ use {
         RenderPass, RenderPassInfo,
     },
     parking_lot::Mutex,
+    paste::paste,
     std::{
         collections::{HashMap, VecDeque},
         sync::Arc,
@@ -34,7 +35,6 @@ pub struct HashPool {
     render_pass_cache: HashMap<RenderPassInfo, Cache<RenderPass>>,
 }
 
-// TODO: Add some sort of manager features (like, I dunno, "Clear Some Memory For me")
 impl HashPool {
     /// Constructs a new `HashPool`.
     pub fn new(device: &Arc<Device>) -> Self {
@@ -50,8 +50,61 @@ impl HashPool {
             render_pass_cache: Default::default(),
         }
     }
+
+    /// Clears the pool, removing all resources.
+    pub fn clear(&mut self) {
+        self.clear_accel_structs();
+        self.clear_buffers();
+        self.clear_images();
+    }
 }
 
+macro_rules! resource_mgmt_fns {
+    ($fn_plural:literal, $doc_plural:literal, $ty:ty, $field:ident) => {
+        paste! {
+            impl HashPool {
+                #[doc = "Clears the pool of " $doc_plural ", removing all resources."]
+                pub fn [<clear_ $fn_plural>](&mut self) {
+                    self.$field.clear();
+                }
+
+                #[doc = "Clears the pool of " $doc_plural ", removing resources matching the
+given information."]
+                pub fn [<clear_ $fn_plural _by_info>](
+                    &mut self,
+                    info: impl Into<$ty>,
+                ) {
+                    self.$field.remove(&info.into());
+                }
+
+                #[doc = "Retains only the " $doc_plural " specified by the predicate.\n\nIn other
+words, remove all resources for which `f(&" $ty ")` returns `false`.\n\n"]
+                /// The elements are visited in unsorted (and unspecified) order.
+                ///
+                /// # Performance
+                ///
+                /// Provides the same performance guarantees as
+                /// [`HashMap::retain`](HashMap::retain).
+                pub fn [<retain_ $fn_plural>]<F>(&mut self, mut f: F)
+                where
+                    F: FnMut(&$ty) -> bool,
+                {
+                    self.$field.retain(|info, _| f(info))
+                }
+            }
+        }
+    };
+}
+
+resource_mgmt_fns!(
+    "accel_structs",
+    "acceleration structures",
+    AccelerationStructureInfo,
+    acceleration_structure_cache
+);
+resource_mgmt_fns!("buffers", "buffers", BufferInfo, buffer_cache);
+resource_mgmt_fns!("images", "images", ImageInfo, image_cache);
+
 impl Pool<CommandBufferInfo, CommandBuffer> for HashPool {
     #[profiling::function]
     fn lease(&mut self, info: CommandBufferInfo) -> Result<Lease<CommandBuffer>, DriverError> {
diff --git a/src/pool/lazy.rs b/src/pool/lazy.rs
index 1e8ab36..8c6752e 100644
--- a/src/pool/lazy.rs
+++ b/src/pool/lazy.rs
@@ -50,7 +50,6 @@ pub struct LazyPool {
     render_pass_cache: HashMap<RenderPassInfo, Cache<RenderPass>>,
 }
 
-// TODO: Add some sort of manager features (like, I dunno, "Clear Some Memory For me")
 impl LazyPool {
     /// Constructs a new `LazyPool`.
     pub fn new(device: &Arc<Device>) -> Self {
@@ -66,6 +65,52 @@ impl LazyPool {
             render_pass_cache: Default::default(),
         }
     }
+
+    /// Clears the pool, removing all resources.
+    pub fn clear(&mut self) {
+        self.clear_accel_structs();
+        self.clear_buffers();
+        self.clear_images();
+    }
+
+    /// Clears the pool of acceleration structures, removing all resources.
+    pub fn clear_accel_structs(&mut self) {
+        self.acceleration_structure_cache.clear();
+    }
+
+    /// Clears the pool of acceleration structures, removing resources matching the given
+    /// type.
+    pub fn clear_accel_structs_by_ty(&mut self, ty: vk::AccelerationStructureTypeKHR) {
+        self.acceleration_structure_cache.remove(&ty);
+    }
+
+    /// Clears the pool of buffers, removing all resources.
+    pub fn clear_buffers(&mut self) {
+        self.buffer_cache.clear();
+    }
+
+    /// Clears the pool of images, removing all resources.
+    pub fn clear_images(&mut self) {
+        self.image_cache.clear();
+    }
+
+    /// Retains only the acceleration structures specified by the predicate.
+    ///
+    /// In other words, remove all resources for which `f(vk::AccelerationStructureTypeKHR)` returns
+    /// `false`.
+    ///
+    /// The elements are visited in unsorted (and unspecified) order.
+    ///
+    /// # Performance
+    ///
+    /// Provides the same performance guarantees as
+    /// [`HashMap::retain`](HashMap::retain).
+    pub fn retain_accel_structs<F>(&mut self, mut f: F)
+    where
+        F: FnMut(vk::AccelerationStructureTypeKHR) -> bool,
+    {
+        self.acceleration_structure_cache.retain(|&ty, _| f(ty))
+    }
 }
 
 impl Pool<AccelerationStructureInfo, AccelerationStructure> for LazyPool {
@@ -123,10 +168,7 @@ impl Pool<BufferInfo, Buffer> for LazyPool {
             // Look for a compatible buffer (same mapping mode, big enough, superset of usage flags)
             for idx in 0..cache.len() {
                 let item = &cache[idx];
-                if item.info.can_map == info.can_map
-                    && item.info.size >= info.size
-                    && item.info.usage.contains(info.usage)
-                {
+                if item.info.size >= info.size && item.info.usage.contains(info.usage) {
                     let item = cache.remove(idx).unwrap();
 
                     return Ok(Lease::new(cache_ref, item));