Compiler Explorer

Source code

use std::cell::{Cell, RefCell, UnsafeCell};
use std::marker::PhantomData;
use std::mem;
use std::ptr::NonNull;

#[derive(Clone)]
pub struct ArenaParameters {
    pub(crate) pause_factor: f64,
    pub(crate) timing_factor: f64,
    pub(crate) min_sleep: usize,
}

pub(crate) struct GcBox<T: Collect + ?Sized> {
    pub(crate) flags: GcFlags,
    pub(crate) next: Cell<Option<NonNull<GcBox<dyn Collect>>>>,
    pub(crate) value: UnsafeCell<T>,
}

pub(crate) struct GcFlags(Cell<u8>);

impl GcFlags {
    #[inline]
    pub(crate) fn new() -> GcFlags {
        GcFlags(Cell::new(0))
    }

#[inline]
    pub(crate) fn needs_trace(&self) -> bool {
        self.0.get() & 0x4 != 0x0
    }

#[inline]
    pub(crate) fn set_needs_trace(&self, needs_trace: bool) {
        self.0
            .set((self.0.get() & !0x4) | if needs_trace { 0x4 } else { 0x0 });
    }
}

// Phantom type that holds a lifetime and ensures that it is invariant.
pub(crate) type Invariant<'a> = PhantomData<Cell<&'a ()>>;

#[derive(Copy, Clone, Eq, PartialEq)]
pub(crate) enum GcColor {
    White,
    Gray,
    Black,
}

pub unsafe trait Collect {
    #[inline]
    fn needs_trace() -> bool
    where
        Self: Sized,
    { true }
}

/// Handle value given by arena callbacks during construction and mutation.  Allows allocating new
/// `Gc` pointers and internally mutating values held by `Gc` pointers.
#[derive(Copy, Clone)]
pub struct MutationContext<'gc, 'context> {
    _invariant: Invariant<'gc>,
    context: &'context Context,
}

impl<'gc, 'context> MutationContext<'gc, 'context> {
    #[inline]
    pub(crate) unsafe fn allocate_old<T: 'gc + Collect>(self, t: T) -> NonNull<GcBox<T>> {
        self.context.allocate_old(t)
    }
    #[inline]
    pub(crate) unsafe fn allocate_new<T: 'gc + Collect>(self, t: T) -> NonNull<GcBox<T>> {
        self.context.allocate_new(t)
    }
}

/// Handle value given by arena callbacks during garbage collection, which must be passed through
/// `Collect::trace` implementations.
#[derive(Copy, Clone)]
pub struct CollectionContext<'context> {
    context: &'context Context,
}

// Main gc context type, public because it must be accessible from the `make_arena!` macro.
#[doc(hidden)]
pub struct Context {
    parameters: ArenaParameters,

phase: Cell<Phase>,
    total_allocated: Cell<usize>,
    remembered_size: Cell<usize>,
    wakeup_total: Cell<usize>,
    allocation_debt: Cell<f64>,

all: Cell<Option<NonNull<GcBox<dyn Collect>>>>,
    sweep: Cell<Option<NonNull<GcBox<dyn Collect>>>>,
    sweep_prev: Cell<Option<NonNull<GcBox<dyn Collect>>>>,

gray: RefCell<Vec<NonNull<GcBox<dyn Collect>>>>,
    gray_again: RefCell<Vec<NonNull<GcBox<dyn Collect>>>>,
}

impl Context {
    unsafe fn allocate_old<T: Collect>(&self, t: T) -> NonNull<GcBox<T>> {
        let alloc_size = mem::size_of::<GcBox<T>>();
        self.total_allocated
            .set(self.total_allocated.get() + alloc_size);
        if self.phase.get() == Phase::Sleep && self.total_allocated.get() > self.wakeup_total.get()
        {
            self.phase.set(Phase::Wake);
        }

if self.phase.get() != Phase::Sleep {
            self.allocation_debt.set(
                self.allocation_debt.get()
                    + alloc_size as f64
                    + alloc_size as f64 / self.parameters.timing_factor,
            );
        }

let gc_box = GcBox {
            flags: GcFlags::new(),
            next: Cell::new(self.all.get()),
            value: UnsafeCell::new(t),
        };
        gc_box.flags.set_needs_trace(T::needs_trace());
        let ptr = NonNull::new_unchecked(Box::into_raw(Box::new(gc_box)));
        self.all.set(Some(static_gc_box(ptr)));
        if self.phase.get() == Phase::Sweep && self.sweep_prev.get().is_none() {
            self.sweep_prev.set(self.all.get());
        }

ptr
    }

unsafe fn allocate_new<T: Collect>(&self, t: T) -> NonNull<GcBox<T>> {
        let alloc_size = mem::size_of::<GcBox<T>>();
        self.total_allocated
            .set(self.total_allocated.get() + alloc_size);
        if self.phase.get() == Phase::Sleep && self.total_allocated.get() > self.wakeup_total.get()
        {
            self.phase.set(Phase::Wake);
        }

let flags = GcFlags::new();
        flags.set_needs_trace(T::needs_trace());

let mut uninitialized = Box::new(mem::MaybeUninit::<GcBox<T>>::uninit());

std::ptr::write(&mut (*uninitialized.as_mut_ptr()).flags, flags);
        std::ptr::write(&mut (*uninitialized.as_mut_ptr()).next, Cell::new(self.all.get()));
        std::ptr::write(&mut (*uninitialized.as_mut_ptr()).value, UnsafeCell::new(t));

let ptr = NonNull::new_unchecked(Box::into_raw(uninitialized) as *mut GcBox<T>);
        self.all.set(Some(static_gc_box(ptr)));
        if self.phase.get() == Phase::Sweep && self.sweep_prev.get().is_none() {
            self.sweep_prev.set(self.all.get());
        }

ptr
    }

}

#[derive(Copy, Clone, Eq, PartialEq)]
enum Phase {
    Wake,
    Propagate,
    Sweep,
    Sleep,
}

#[inline]
unsafe fn static_gc_box<'gc>(
    ptr: NonNull<GcBox<dyn Collect + 'gc>>,
) -> NonNull<GcBox<dyn Collect>> {
    mem::transmute(ptr)
}

pub struct GcCell<'gc, T: 'gc + Collect>(Gc<'gc, GcRefCell<T>>);

impl<'gc, T: Collect + 'gc> Copy for GcCell<'gc, T> {}

impl<'gc, T: Collect + 'gc> Clone for GcCell<'gc, T> {
    fn clone(&self) -> GcCell<'gc, T> {
        *self
    }
}
unsafe impl<'gc, T: 'gc + Collect> Collect for GcCell<'gc, T> {}

impl<'gc, T: 'gc + Collect> GcCell<'gc, T> {
    pub fn allocate_old(mc: MutationContext<'gc, '_>, t: T) -> GcCell<'gc, T> {
        GcCell(Gc::allocate_old(
            mc,
            GcRefCell {
                cell: RefCell::new(t),
            },
        ))
    }
    pub fn allocate_new(mc: MutationContext<'gc, '_>, t: T) -> GcCell<'gc, T> {
        GcCell(Gc::allocate_new(
            mc,
            GcRefCell {
                cell: RefCell::new(t),
            },
        ))
    }
}
struct GcRefCell<T: Collect> {
    cell: RefCell<T>,
}

unsafe impl<'gc, T: Collect + 'gc> Collect for GcRefCell<T> {}

pub struct Gc<'gc, T: 'gc + Collect> {
    pub(crate) ptr: NonNull<GcBox<T>>,
    _invariant: Invariant<'gc>,
}

impl<'gc, T: Collect + 'gc> Copy for Gc<'gc, T> {}

impl<'gc, T: Collect + 'gc> Clone for Gc<'gc, T> {
    fn clone(&self) -> Gc<'gc, T> {
        *self
    }
}

unsafe impl<'gc, T: 'gc + Collect> Collect for Gc<'gc, T> {}

impl<'gc, T: 'gc + Collect> Gc<'gc, T> {
    pub fn allocate_old(mc: MutationContext<'gc, '_>, t: T) -> Gc<'gc, T> {
        Gc {
            ptr: unsafe { mc.allocate_old(t) },
            _invariant: PhantomData,
        }
    }
    pub fn allocate_new(mc: MutationContext<'gc, '_>, t: T) -> Gc<'gc, T> {
        Gc {
            ptr: unsafe { mc.allocate_new(t) },
            _invariant: PhantomData,
        }
    }
}

pub struct Data {
    x: [u8; 1000]
}

unsafe impl Collect for Data
where Data: 'static,
{
    #[inline]
    fn needs_trace() -> bool { false }
}

pub fn old<'gc>(gc_context: MutationContext<'gc, '_>) -> GcCell<'gc, Data> {
    GcCell::allocate_old(gc_context, Data { x: [0; 1000] })
}
pub fn new<'gc>(gc_context: MutationContext<'gc, '_>) -> GcCell<'gc, Data> {
    GcCell::allocate_new(gc_context, Data { x: [0; 1000] })
}