Simple GC in Rust

gc-example.rs

use std::cell::RefCell;

use crate::gc::{collect, Gc, GcCtx, Trace};

#[derive(Clone)]
struct Node {
    id: usize,
    next: RefCell<Option<Gc<Node>>>,
}

impl Drop for Node {
    fn drop(&mut self) {
        println!("drop {}", self.id);
    }
}

impl Trace for Node {
    fn trace(&self, ctx: &GcCtx) {
        if let Some(next) = self.next.borrow().as_ref() {
            ctx.mark(next)
        }
    }
}

fn main() {
    {
        // foo1 -> foo2 -> foo1
        let foo1 = Gc::new(Node {
            id: 0,
            next: RefCell::new(None),
        });
        let foo1_clone = foo1.clone();
        foo1_clone.unset_root();
        let foo2 = Gc::new(Node {
            id: 1,
            next: RefCell::new(Some(foo1_clone)),
        });
        foo2.unset_root();
        *foo1.next.borrow_mut() = Some(foo2);

        let _ = Gc::new(Node {
            id: 2,
            next: RefCell::new(None),
        });

        collect();
    }
    println!("---");
    collect();
}

gc.rs

use std::cell::{Cell, RefCell};
use std::marker::PhantomData;
use std::ops::Deref;
use std::ptr::NonNull;
use std::rc::Rc;

pub trait Trace {
    fn trace(&self, ctx: &GcCtx);
}

struct GlobalGc {
    objects: Vec<NonNull<GcInner<dyn Trace>>>,
    color: bool,
}

impl GlobalGc {
    fn push(&mut self, object: NonNull<GcInner<dyn Trace>>) {
        self.objects.push(object);
    }

    fn collect(&mut self) {
        self.color = !self.color;
        let ctx = GcCtx::new(self.color);

        for object in &mut self.objects {
            unsafe {
                if object.as_ref().count > 0 {
                    object.as_mut().mark = self.color;
                    object.as_ref().value.trace(&ctx);
                }
            }
        }

        let mut i = 0;
        while i < self.objects.len() {
            let object = self.objects[i];
            unsafe {
                if object.as_ref().count == 0 && object.as_ref().mark != self.color {
                    let t = self.objects.swap_remove(i);
                    drop(Box::from_raw(t.as_ptr()));
                } else {
                    i += 1;
                }
            }
        }
    }
}

pub struct GcCtx {
    color: bool,
}

impl GcCtx {
    fn new(color: bool) -> GcCtx {
        GcCtx { color }
    }

    pub fn mark<T: Trace + 'static>(&self, object: &Gc<T>) {
        let mut inner = object.inner.get();
        unsafe {
            if inner.as_ref().mark == self.color {
                return;
            }
            inner.as_mut().mark = self.color;
            inner.as_ref().value.trace(self);
        }
    }
}

thread_local! {
    static LOCAL_GC: RefCell<GlobalGc> = RefCell::new(GlobalGc {
        objects: Vec::new(),
        color: false,
    });
}

pub fn collect() {
    LOCAL_GC.with(|gc| {
        let mut gc = gc.borrow_mut();
        gc.collect();
    })
}

struct GcInner<T: ?Sized + 'static> {
    mark: bool,
    count: usize,
    value: T,
}

impl<T> GcInner<T> {
    fn new(value: T, mark: bool) -> GcInner<T> {
        GcInner {
            value,
            mark,
            count: 1,
        }
    }
}

pub struct Gc<T: 'static> {
    inner: Cell<NonNull<GcInner<T>>>,
    root: Cell<bool>,
    marker: PhantomData<Rc<T>>,
}

impl<T: Trace> Gc<T> {
    pub fn new(value: T) -> Gc<T> {
        LOCAL_GC.with(|gc| {
            let mut gc = gc.borrow_mut();
            let inner = NonNull::from(Box::leak(Box::new(GcInner::new(value, gc.color))));
            gc.push(inner);
            Gc {
                inner: Cell::new(inner),
                root: Cell::new(true),
                marker: PhantomData,
            }
        })
    }

    pub fn set_root(&self) {
        if self.root.get() {
            return;
        }
        self.root.set(true);
        let mut inner = self.inner.get();
        unsafe {
            inner.as_mut().count += 1;
        }
    }

    pub fn unset_root(&self) {
        if !self.root.get() {
            return;
        }
        self.root.set(false);
        let mut inner = self.inner.get();
        unsafe {
            inner.as_mut().count -= 1;
        }
    }
}

impl<T: Trace + Clone> Clone for Gc<T> {
    fn clone(&self) -> Gc<T> {
        let mut inner = self.inner.get();
        unsafe {
            inner.as_mut().count += 1;
        }
        Gc {
            inner: Cell::new(inner),
            root: Cell::new(true),
            marker: PhantomData,
        }
    }
}

impl<T: Trace> Deref for Gc<T> {
    type Target = T;

    fn deref(&self) -> &T {
        let inner = self.inner.get();
        unsafe { &inner.as_ref().value }
    }
}

impl<T> Drop for Gc<T> {
    fn drop(&mut self) {
        if self.root.get() {
            let mut inner = self.inner.get();
            unsafe {
                inner.as_mut().count -= 1;
            }
        }
    }
}