koi_gc/

lib.rs

use std::{
    cell::{Cell, RefCell},
    collections::{HashMap, HashSet},
    marker::PhantomData,
    ops::Deref,
    ptr::NonNull,
};

mod std_impls;

/// The data size type used by is the Gc. This is the hard limit to the number of objects in the Gc and the
/// the number of clones of each object. Thus, the total number of objects is this value squared.
/// This is set by the following feature flags:
///
/// - default (unset): u16 (65,535)
/// - small: u8 (255)
/// - medium: u16 (65,535)
/// - large: usize (32 bit: 4,294,967,295 , 64 bit: 18,446,744,073,709,551,615)
#[cfg(all(
    not(feature = "large"),
    not(feature = "medium"),
    not(feature = "small")
))]
pub type GcSize = u16; // Default
#[cfg(all(not(feature = "large"), not(feature = "medium"), feature = "small"))]
pub type GcSize = u8; // Small
#[cfg(all(not(feature = "large"), feature = "medium"))]
pub type GcSize = u16; // Medium
#[cfg(all(feature = "large"))]
pub type GcSize = usize; // Large

// todo replace RefCell with UnsafeCell in release mode
thread_local! {
    static CYCLE_CHECKS: Cell<bool> = Cell::new(true);
    static GC_REGISTRY: RefCell<HashMap<GcSize, NonNull<dyn GcInterface + 'static>>> = RefCell::new(HashMap::with_capacity(AVAILABLE_ID_INITIAL_CAPACITY));
    static AVAILABLE_IDS: RefCell<Vec<GcSize>> = RefCell::new(Vec::with_capacity(AVAILABLE_ID_INITIAL_CAPACITY));
    static GC_ID_COUNTER: Cell<GcSize> = Cell::new(0);
    // Track allocations since last collection
    static ALLOCATIONS_SINCE_LAST_CYCLE_CHECK: Cell<usize> = Cell::new(0);
    static REMAINING_ALLOCATIONS_AFTER_LAST_CYCLE_CHECK: Cell<usize> = Cell::new(0);
}

#[cfg(all(
    not(feature = "large"),
    not(feature = "medium"),
    not(feature = "small")
))]
const AVAILABLE_ID_INITIAL_CAPACITY: usize = 256; // Default
#[cfg(all(not(feature = "large"), not(feature = "medium"), feature = "small"))]
const AVAILABLE_ID_INITIAL_CAPACITY: usize = 32; // Small
#[cfg(all(not(feature = "large"), feature = "medium"))]
const AVAILABLE_ID_INITIAL_CAPACITY: usize = 256; // Medium
#[cfg(all(feature = "large"))]
const AVAILABLE_ID_INITIAL_CAPACITY: usize = 1024; // Large

/// The threshold number of allocations to consider performing a cycle collection
const COLLECTION_SOFT_THRESHOLD: usize = 100;
/// The threshold number of allocations to perform a cycle collection
const COLLECTION_HARD_THRESHOLD: usize = 10000;

pub struct Gc<T>
where
    T: VisitGc,
{
    ptr: NonNull<GcInner<T>>,
    phantom: PhantomData<GcInner<T>>,
}

struct GcInner<T>
where
    T: VisitGc,
{
    id: GcSize,
    strong: Cell<GcSize>,
    data: T,
}

struct GcInnerCheck {
    /// Temporary ref count used during cycle detection
    gc_refs: GcSize,
    /// Mark for reachability during cycle detection
    marked: bool,
}

impl<T> Gc<T>
where
    T: VisitGc,
{
    pub fn new(data: T) -> Self {
        if CYCLE_CHECKS.get() {
            Gc::maybe_collect_cycles();
        }
        ALLOCATIONS_SINCE_LAST_CYCLE_CHECK.with(|e| e.update(|e| e + 1));
        let id = AVAILABLE_IDS.with(|available_ids| {
            let mut available_ids = available_ids.borrow_mut();
            let next_id = if let Some(free_id) = available_ids.pop() {
                free_id
            } else {
                #[cfg(all(
                    not(feature = "large"),
                    not(feature = "medium"),
                    not(feature = "small")
                ))]
                const EXPECT_MSG: &str = "More than GcSize::MAX objects created and alive. Try setting feature flag 'large'"; // Default
                #[cfg(all(not(feature = "large"), not(feature = "medium"), feature = "small"))]
                const EXPECT_MSG: &str = "More than GcSize::MAX objects created and alive. Try setting feature flag 'medium' or 'large'"; // Small
                #[cfg(all(not(feature = "large"), feature = "medium"))]
                const EXPECT_MSG: &str = "More than GcSize::MAX objects created and alive. Try setting feature flag 'large'"; // Medium
                #[cfg(all(feature = "large"))]
                const EXPECT_MSG: &str = "More than GcSize::MAX objects created and alive."; // Large
                let next_id = GC_ID_COUNTER.get().checked_add(1).expect(EXPECT_MSG);
                GC_ID_COUNTER.set(next_id);
                next_id
            };
            next_id
        });
        let inner = Box::leak(Box::new(GcInner {
            id,
            strong: Cell::new(1),
            data,
        }));
        let ptr = NonNull::from(inner);
        GC_REGISTRY.with(|registry| {
            let mut registry = registry.borrow_mut();
            let current = registry.insert(id, ptr);
            debug_assert!(current.is_none());
        });
        Gc {
            ptr,
            phantom: PhantomData,
        }
    }

    fn inner(&self) -> &GcInner<T> {
        unsafe { self.ptr.as_ref() }
    }
}

impl Gc<()> {
    pub fn collect_cycles() -> usize {
        let mut ids_to_reclaim = Vec::new();
        GC_REGISTRY.with(|registry| {
            let registry = registry.borrow_mut();
            let mut gc_cycle_check = HashMap::with_capacity(registry.len());
            // Copy reference counts to gc_refs and clear marks
            for obj_ptr in registry.values() {
                let obj = unsafe { obj_ptr.as_ref() };
                let id = obj.id();
                let gc_refs = obj.strong();
                let obj_check = GcInnerCheck {
                    gc_refs,
                    marked: false,
                };
                gc_cycle_check.insert(id, obj_check);
            }

            for obj_ptr in registry.values() {
                let obj = unsafe { obj_ptr.as_ref() };
                obj.visit_edges(&mut |edge| {
                    let id = edge.id();
                    let obj_check_edge = gc_cycle_check.get_mut(&id).unwrap();
                    let current = obj_check_edge.gc_refs;
                    if current > 0 {
                        obj_check_edge.gc_refs = current - 1;
                    }
                });
            }

            let unaccounted_refs = &mut ids_to_reclaim; // reuse Vec allocation
            for obj_ptr in registry.values() {
                let obj = unsafe { obj_ptr.as_ref() };
                let id = obj.id();
                let obj_check = gc_cycle_check.get_mut(&id).unwrap();
                if obj_check.gc_refs > 0 {
                    unaccounted_refs.push(id);
                }
            }

            let to_visit = unaccounted_refs;
            let mut visited = HashSet::new();

            while let Some(id) = to_visit.pop() {
                if !visited.insert(id) {
                    continue; // Already visited
                }

                if let Some(obj_ptr) = registry.get(&id) {
                    let obj = unsafe { obj_ptr.as_ref() };
                    let id = obj.id();
                    let obj_check = gc_cycle_check.get_mut(&id).unwrap();
                    obj_check.marked = true;

                    // Add all edges to visit queue
                    obj.visit_edges(&mut |edge| {
                        let id = edge.id();
                        if !visited.contains(&id) {
                            to_visit.push(id);
                        }
                    });
                }
            }

            ids_to_reclaim.clear(); // reuse Vec allocation
            for (_, obj_ptr) in registry.iter() {
                let obj = unsafe { obj_ptr.as_ref() };
                let id = obj.id();
                let obj_check = gc_cycle_check.get_mut(&id).unwrap();
                if obj_check.marked {
                    obj_check.marked = false;
                    continue;
                }
                ids_to_reclaim.push(id);
            }
        });

        let collected = ids_to_reclaim.len();
        let num_of_allocations_remaining = GC_REGISTRY.with(|registry| {
            AVAILABLE_IDS.with(|ids| {
                for id in ids_to_reclaim {
                    let removed = registry.borrow_mut().remove(&id);
                    if let Some(removed) = removed {
                        ids.borrow_mut().push(id);
                        let temp = unsafe { removed.as_ref() };
                        // Set mark flag so does not try to remove from registry for within drop
                        temp.mark_gc_deleted();
                        unsafe {
                            drop(Box::from_raw(removed.as_ptr()));
                        }
                    }
                }
            });
            registry.borrow().len()
        });

        // Reset allocation counter
        ALLOCATIONS_SINCE_LAST_CYCLE_CHECK.set(0);
        REMAINING_ALLOCATIONS_AFTER_LAST_CYCLE_CHECK.set(num_of_allocations_remaining);

        collected
    }

    /// Check if [`Gc::collect_cycles`] should run based on the internal heuristic and if so runs it.
    /// This should only really ever be called externally if [`Gc::enable_acc`] was just called.
    fn maybe_collect_cycles() -> Option<usize> {
        let allocations_since_last_collection = ALLOCATIONS_SINCE_LAST_CYCLE_CHECK.get();
        if allocations_since_last_collection > COLLECTION_HARD_THRESHOLD {
            return Some(Gc::collect_cycles());
        } else if allocations_since_last_collection > COLLECTION_SOFT_THRESHOLD {
            // If allocations are exploding, do a cycle collection, otherwise things are likely being reclaimed as normal
            let current_allocations = GC_REGISTRY.with(|e| e.borrow().len());
            let has_allocations_quadrupled =
                current_allocations > 4 * REMAINING_ALLOCATIONS_AFTER_LAST_CYCLE_CHECK.get();
            if has_allocations_quadrupled {
                return Some(Gc::collect_cycles());
            }
        }
        None
    }

    /// Cycle collection is disabled by default. If called, enables automatic cycle collection -
    /// [`Gc::collect_cycles`] is automatically called based the
    /// internal allocation heuristic.
    /// Consider pairing this with [`Gc::maybe_collect_cycles`].
    pub fn enable_acc(&self) {
        CYCLE_CHECKS.set(true);
    }

    /// Cycle collection is disabled by default. If called, disables automatic cycle collection -
    /// and [`Gc::collect_cycles`] needs to be called manually
    pub fn disable_acc(&self) {
        CYCLE_CHECKS.set(false);
    }
}

impl<T> Clone for Gc<T>
where
    T: VisitGc,
{
    fn clone(&self) -> Self {
        let inner = self.inner();
        let current = inner.strong.get();
        inner.strong.set(current + 1);

        Gc {
            ptr: self.ptr,
            phantom: PhantomData,
        }
    }
}

impl<T> Drop for Gc<T>
where
    T: VisitGc,
{
    fn drop(&mut self) {
        let inner = self.inner();
        let current = inner.strong.get();
        // Deleted by the gc and already handled. Initial value of this type is always 1 so this is a special flag.
        if current == 0 {
            return;
        }
        if current == 1 {
            GC_REGISTRY.with(|registry| {
                let mut registry = registry.borrow_mut();
                let removed = registry.remove(&inner.id);
                debug_assert!(removed.is_some());
            });
            AVAILABLE_IDS.with(|ids| {
                ids.borrow_mut().push(inner.id);
            });
            // Deallocate the inner box
            unsafe {
                drop(Box::from_raw(self.ptr.as_ptr()));
            }
        } else {
            inner.strong.set(current - 1);
        }
    }
}

impl<T> Deref for Gc<T>
where
    T: VisitGc,
{
    type Target = T;

    fn deref(&self) -> &Self::Target {
        &self.inner().data
    }
}

/// Dyn interface for Gc. Marked as unsafe since should not be implemented.
pub unsafe trait GcInterface {
    fn id(&self) -> GcSize;
    fn strong(&self) -> GcSize;
    fn mark_gc_deleted(&self);
    fn visit_edges(&self, visitor: &mut dyn FnMut(&dyn GcInterface));
}

unsafe impl<T> GcInterface for GcInner<T>
where
    T: VisitGc,
{
    fn id(&self) -> GcSize {
        self.id
    }

    fn strong(&self) -> GcSize {
        self.strong.get()
    }

    fn mark_gc_deleted(&self) {
        self.strong.set(0);
    }

    fn visit_edges(&self, visitor: &mut dyn FnMut(&dyn GcInterface)) {
        self.data.visit(visitor);
    }
}

// todo verify this
/// SAFETY: Implementing this trait is marked as `unsafe` since it is possible incorrectly implement this
/// trait can cause the garbage collector to free data that is being pointed at by accessible
/// data
pub unsafe trait VisitGc: 'static {
    fn visit(&self, visitor: &mut dyn FnMut(&dyn GcInterface));
}

unsafe impl<T> VisitGc for Gc<T>
where
    T: VisitGc,
{
    fn visit(&self, visitor: &mut dyn FnMut(&dyn GcInterface)) {
        visitor(self.inner());
    }
}

#[cfg(test)]
mod tests {
    use std::rc::Rc;

    use super::*;

    #[derive(Clone)]
    struct GcNode {
        value: i32,
        next: Option<Gc<GcNode>>,
    }

    unsafe impl VisitGc for GcNode {
        fn visit(&self, visitor: &mut dyn FnMut(&dyn GcInterface)) {
            self.value.visit(visitor);
            self.next.visit(visitor);
        }
    }

    #[test]
    fn test_cycle_collection() {
        // Create a cycle: A -> B -> A
        let a = Gc::new(GcNode {
            value: 1,
            next: None,
        });

        let b = Gc::new(GcNode {
            value: 2,
            next: Some(a.clone()),
        });

        // complete the cycle
        unsafe {
            let a_inner = a.ptr.as_ref();
            let data_ptr = &a_inner.data as *const GcNode as *mut GcNode;
            (*data_ptr).next = Some(b.clone());
        }

        // no cycle
        let c = Gc::new(GcNode {
            value: 3,
            next: None,
        });

        let d = Gc::new(GcNode {
            value: 4,
            next: Some(c.clone()),
        });

        drop(a);
        drop(b);
        drop(c);
        drop(d);

        // Run collection - should collect the cycle
        let collected = Gc::collect_cycles();
        assert_eq!(collected, 2, "Should collect the 2 nodes in the cycle");
    }

    struct Leak<T> {
        cycle: RefCell<Option<Rc<Rc<Leak<T>>>>>,
        data: T,
    }

    fn rc_cycle_forget<T>(data: T) {
        let e = Rc::new(Leak {
            cycle: RefCell::new(None),
            data: data,
        });
        *e.cycle.borrow_mut() = Some(Rc::new(e.clone()));
    }

    struct GcLeak<T>
    where
        T: VisitGc,
    {
        cycle: RefCell<Option<Gc<Gc<GcLeak<T>>>>>,
        data: T,
    }

    unsafe impl<T> VisitGc for GcLeak<T>
    where
        T: VisitGc,
    {
        fn visit(&self, visitor: &mut dyn FnMut(&dyn GcInterface)) {
            self.cycle.visit(visitor);
            self.data.visit(visitor);
        }
    }

    fn gc_cycle_forget<T: VisitGc>(data: T) {
        let e = Gc::new(GcLeak {
            cycle: RefCell::new(None),
            data: data,
        });
        *e.cycle.borrow_mut() = Some(Gc::new(e.clone()));
    }

    #[test]
    fn not_leaked() {
        thread_local! {
            static INCREMENT_ON_DROP: Cell<usize> = Cell::new(0);
        }

        struct TestData {}

        unsafe impl VisitGc for TestData {
            fn visit(&self, visitor: &mut dyn FnMut(&dyn GcInterface)) {}
        }

        impl Drop for TestData {
            fn drop(&mut self) {
                INCREMENT_ON_DROP.set(INCREMENT_ON_DROP.get() + 1);
            }
        }

        let temp = Gc::new(TestData {});
        drop(temp);
        assert_eq!(INCREMENT_ON_DROP.get(), 1);
        let temp = Gc::new(TestData {});
        rc_cycle_forget(temp);
        Gc::collect_cycles();
        assert_eq!(
            INCREMENT_ON_DROP.get(),
            1,
            "Gc::collect_cycles should not nothing since it does not know about rc cycles"
        );
        let temp = Gc::new(TestData {});
        gc_cycle_forget(temp);
        let temp = Gc::new(TestData {});
        assert_eq!(
            INCREMENT_ON_DROP.get(),
            1,
            "Gc::collect_cycles has not ran yet"
        );
        Gc::collect_cycles();
        assert_eq!(
            INCREMENT_ON_DROP.get(),
            2,
            "Gc::collect_cycles should nothing detect the cycle and only remove one"
        );
        drop(temp);
        assert_eq!(INCREMENT_ON_DROP.get(), 3);
        let temp = Gc::new(TestData {});
        std::mem::forget(temp);
        assert_eq!(INCREMENT_ON_DROP.get(), 3);
        Gc::collect_cycles();
        assert_eq!(
            INCREMENT_ON_DROP.get(),
            3,
            "Gc::collect_cycles only works for cycles. mem::forget is unfortunately marked as safe"
        );
    }
}