From e8f6ff5f8a54d14767bcd1e36ca71a539791f290 Mon Sep 17 00:00:00 2001
From: Jeff Muizelaar <jrmuizel@gmail.com>
Date: Mon, 9 Oct 2023 15:16:50 -0400
Subject: [PATCH] Add a function to free dead roots without doing a full cycle
 collection

---
 src/collect.rs | 20 ++++++++++++++++++++
 src/lib.rs     | 15 +++++++++++++++
 2 files changed, 35 insertions(+)

diff --git a/src/collect.rs b/src/collect.rs
index b121963..5e8dcd0 100644
--- a/src/collect.rs
+++ b/src/collect.rs
@@ -100,6 +100,26 @@ pub fn number_of_roots_buffered() -> usize {
     ROOTS.with(|r| r.borrow().len())
 }
 
+
+/// Free's all the of the roots that have a reference count of 0.
+/// This is much faster than doing a full cycle collection and will
+/// ensure that any memory that was never part of a cycle is freed.
+pub fn free_dead_roots() {
+    ROOTS.with(|r| {
+        let mut v = r.borrow_mut();
+        v.retain_mut(|root| {
+            let s: &dyn CcBoxPtr = unsafe { root.as_ref() };
+            if s.data().strong() == 0 {
+                s.data().buffered.set(false);
+                unsafe { free(*root) };
+                false
+            } else {
+                true
+            }
+        });
+    });
+}
+
 /// Invoke cycle collection for all `Cc<T>`s on this thread.
 ///
 /// You may wish to do this when the roots buffer reaches a certain size, when
diff --git a/src/lib.rs b/src/lib.rs
index a1a1805..11cf4e7 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -984,6 +984,8 @@ pub(crate) unsafe fn drop_value(mut ptr: NonNull<dyn CcBoxPtr>) {
 mod tests {
     use core::cell::RefCell;
 
+    use crate::{number_of_roots_buffered, collect::free_dead_roots};
+
     use super::{collect_cycles, Cc, Trace, Tracer, Weak};
 
     // Tests copied from `Rc<T>`.
@@ -1506,4 +1508,17 @@ mod tests {
         drop(b);
         collect_cycles();
     }
+
+
+    #[test]
+    fn freeing_dead_roots() {
+        assert_eq!(number_of_roots_buffered(), 0);
+        let a = Cc::new(1);
+        let b = a.clone();
+        drop(b);
+        assert_eq!(number_of_roots_buffered(), 1);
+        drop(a);
+        free_dead_roots();
+        assert_eq!(number_of_roots_buffered(), 0);
+    }
 }