crftng-intrprtrs/gc/src/allocator.rs

use std::ptr::NonNull;

use super::gc_ref::GcRef;
use super::trace;
use crate::trace::GCTrace;

#[derive(Default)]
pub struct GCAllocator {
	allocations: Vec<Allocation>,
}

impl GCAllocator {
	#[inline(always)]
	/// Allocate a value on the heap
	pub fn alloc<T: trace::GCTrace>(&mut self, x: T) -> GcRef<T> {
		let alloc = Allocation::new(x);
		let ptr = alloc.ptr as *mut T;
		self.allocations.push(alloc);
		unsafe { GcRef::new(NonNull::new(ptr).unwrap()) }
	}

	/// # Safety
	/// Root needs to contain all the accessible gc allocated references
	/// In case root itself is gc allocated use [`gc_ref_root`]
	///
	/// [`gc_ref_root`]: GCAllocator::gc_ref_root
	pub unsafe fn gc<T: trace::GCTrace>(&mut self, root: &T) {
		// Mark
		let mut tracer = trace::GCTracer::with_capacity(self.allocations.len());
		root.trace(&mut tracer);

		// And sweep
		let inaccessible = self
			.allocations
			.drain_filter(|a| !tracer.is_accessible(a.ptr));

		// And sweep
		for mut to_free in inaccessible {
			to_free.drop();
		}
	}

	// Specialization when ;-;
	/// [`gc`] but for roots that are allocated on the heap
	/// # Safety
	/// See [`gc`]
	///
	/// [`gc`]: GCAllocator::gc
	pub unsafe fn gc_ref_root<T: trace::GCTrace>(&mut self, root: &GcRef<T>) {
		// Mark
		let mut tracer = trace::GCTracer::with_capacity(self.allocations.len());
		tracer.mark_reachable_rec(root);
		root.trace(&mut tracer);

		let inaccessible = self
			.allocations
			.drain_filter(|a| !tracer.is_accessible(a.ptr));

		// And sweep
		for mut to_free in inaccessible {
			to_free.drop();
		}
	}
}

struct Allocation {
	ptr: *mut (),
	drop: unsafe fn(*mut ()),
}

impl Allocation {
	fn new<T>(x: T) -> Self {
		let alloc = Box::new(x);
		let ptr = Box::into_raw(alloc) as *mut ();

		let drop = |ptr| unsafe {
			std::mem::drop(Box::from_raw(ptr as *mut T));
		};

		Self { ptr, drop }
	}

	fn drop(&mut self) {
		unsafe { (self.drop)(self.ptr) };
	}
}