gladiator 0.0.0-alpha

//! A precise, concurrent garbage collector written in Rust.

use std::{sync::{RwLock, RwLockWriteGuard, RwLockReadGuard}, marker::PhantomData, ptr::NonNull, cell::Cell, ops::{Deref, DerefMut}, fmt::{Debug, Formatter, Display}, mem::size_of};

use trace::{Trace, Signal};

pub mod trace;

struct Allocation<'arena, T: Trace<'arena, T>> {
    safety: PhantomData<&'arena mut ()>,

    /// Prevents multiple threads from using this allocation at once.
    pub guard: RwLock<()>,

    /// The number of rooted references to this allocation.
    pub roots: Cell<usize>,

    /// Whether or not the allocation has been rooted.
    pub mark: Cell<bool>,

    /// The list allocation in the linked list of allocations, if any.
    pub next: Option<NonNull<Allocation<'arena, T>>>,

    /// The data wrapped by the allocation, of course.
    pub data: T,
}

impl<'arena, T: Trace<'arena, T>> Allocation<'arena, T> {
    unsafe fn mark(&self, signal: &Signal<'arena, T>) {
        if !self.mark.get() {
            self.mark.set(true);
            self.data.send(signal);
        }
    }

    unsafe fn root(&self) {
        self.roots.set(self.roots.get() + 1);
    }

    unsafe fn unroot(&self) {
        self.roots.set(self.roots.get() - 1);
    }
}

/// A value mutably borrowed from a [`Gc`] smart pointer.
/// 
/// It is impossible for any other borrows to exist at the same time as a mutable borrow.
pub struct BorrowMut<'guard, 'arena, T: Trace<'arena, T>> {
    _guard: RwLockWriteGuard<'guard, ()>,
    inner: NonNull<Allocation<'arena, T>>,
}

impl<'guard, 'arena: 'guard, T: Trace<'arena, T>> BorrowMut<'guard, 'arena, T> {
    #[inline]
    unsafe fn new(inner: NonNull<Allocation<'arena, T>>) -> Self {
        BorrowMut { 
            _guard: inner.as_ref().guard.write().expect("a garbage collected value cannot be poisoned"),
            inner,
        }
    }

    #[inline]
    unsafe fn try_new(inner: NonNull<Allocation<'arena, T>>) -> Option<Self> {
        Some(BorrowMut { 
            _guard: match inner.as_ref().guard.try_write() {
                Ok(guard) => guard,
                Err(_) => return None,
            },
            inner,
        })
    }
}

impl<'guard, 'arena: 'guard, T: Trace<'arena, T>> AsRef<T> for BorrowMut<'guard, 'arena, T> {
    fn as_ref(&self) -> &T {
        unsafe { &self.inner.as_ref().data }
    }
}

impl<'guard, 'arena: 'guard, T: Trace<'arena, T>> AsMut<T> for BorrowMut<'guard, 'arena, T> {
    fn as_mut(&mut self) -> &mut T {
        unsafe { &mut self.inner.as_mut().data }
    }
}

impl<'guard, 'arena: 'guard, T: Trace<'arena, T>> Deref for BorrowMut<'guard, 'arena, T> {
    type Target = T;

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

impl<'guard, 'arena: 'guard, T: Trace<'arena, T>> DerefMut for BorrowMut<'guard, 'arena, T> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        unsafe { &mut self.inner.as_mut().data }
    }
}

/// A value borrowed from a [`Gc`] smart pointer.
/// 
/// It is impossible for any mutable borrows to exist at the same time as immutable borrows.
pub struct Borrow<'guard, 'arena, T: Trace<'arena, T> + Sized> {
    _guard: RwLockReadGuard<'guard, ()>,
    inner: NonNull<Allocation<'arena, T>>,
}

impl<'guard, 'arena: 'guard, T: Trace<'arena, T>> Borrow<'guard, 'arena, T> {
    #[inline]
    unsafe fn new(inner: NonNull<Allocation<'arena, T>>) -> Self {
        Borrow { 
            _guard: inner.as_ref().guard.read().expect("a garbage collected value cannot be poisoned"),
            inner,
        }
    }

    #[inline]
    unsafe fn try_new(inner: NonNull<Allocation<'arena, T>>) -> Option<Self> {
        Some(Borrow { 
            _guard: match inner.as_ref().guard.try_read() {
                Ok(guard) => guard,
                Err(_) => return None,
            },
            inner,
        })
    }
}

impl<'guard, 'arena: 'guard, T: Trace<'arena, T>> AsRef<T> for Borrow<'guard, 'arena, T> {
    fn as_ref(&self) -> &T {
        unsafe { &self.inner.as_ref().data }
    }
}

impl<'guard, 'arena: 'guard, T: Trace<'arena, T>> Deref for Borrow<'guard, 'arena, T> {
    type Target = T;

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

/// A smart pointer to a garbage collected value in an [`Arena`].
pub struct Gc<'arena, T: Trace<'arena, T>> {
    weak: Cell<bool>,
    ptr: Cell<NonNull<Allocation<'arena, T>>>,
}

impl<'arena, T: Trace<'arena, T>> Gc<'arena, T> {
    /// Borrows the value that this [`Gc`] points to.  If the value is currently locked by another thread, it will block the current thread until the value becomes
    /// available.
    /// 
    /// # Panics
    /// - If the value that this [`Gc`] references is already borrowed by this thread.
    /// - If another thread panicked while a reference open to the [`Gc`]'s value.
    pub fn borrow(&self) -> Borrow<'_, 'arena, T> {
        unsafe { Borrow::new(self.ptr.get()) }
    }

    /// Attempts to borrow the value that this [`Gc`] points to without blocking the current thread.
    pub fn try_borrow(&self) -> Option<Borrow<'_, 'arena, T>> {
        unsafe { Borrow::try_new(self.ptr.get()) }
    }

    /// Borrows the value that this [`Gc`] points to.  If the value is currently locked by another thread, it will block the current thread until the value becomes
    /// available.
    /// 
    /// # Panics
    /// - If the value that this [`Gc`] references is already borrowed by this thread.
    /// - If another thread panicked while a reference open to the [`Gc`]'s value.
    pub fn borrow_mut(&self) -> BorrowMut<'_, 'arena, T> {
        unsafe { BorrowMut::new(self.ptr.get()) }
    }

    /// Attempts to borrow the value that this [`Gc`] points to without blocking the current thread.
    pub fn try_borrow_mut(&self) -> Option<BorrowMut<'_, 'arena, T>> {
        unsafe { BorrowMut::try_new(self.ptr.get()) }
    }

    /// Returns whether or not this [`Gc`] is a weak reference.
    pub fn is_weak(&self) -> bool {
        self.weak.get()
    }

    /// Returns a weak [`Gc`] which points to the same value as this [`Gc`].
    /// 
    /// Weak [`Gc`]s are references that do not count towards ownership of a garbage collected value.  
    pub fn as_weak(&self) -> Self {
        Self { weak: Cell::new(true), ptr: self.ptr.clone() }
    }

    /// Converts this [`Gc`] into a weak [`Gc`] reference.  If the [`Gc`] is already weak, this just returns itself then.
    pub fn downgrade(self) -> Self {
        Self { weak: Cell::new(true), ptr: self.ptr.clone() }
        // old one implicitly dropped, decrementing root count of the value
    }

    /// Ends the lifetime of this [`Gc`] before it would naturally be dropped.
    pub fn consume(self) {}
}

unsafe impl<'arena, T: Trace<'arena, T>> Send for Gc<'arena, T> {}

unsafe impl<'arena, T: Trace<'arena, T>> Trace<'arena, T> for Gc<'arena, T> {
    unsafe fn send(&self, signal: &trace::Signal<'arena, T>) {
        match signal.kind() {
            trace::SignalKind::Mark => {
                self.ptr.get().as_ref().mark(signal);
            },
            trace::SignalKind::Root => {
                assert!(self.weak.get(), "cannot double-root a Gc");

                self.borrow().inner.as_ref().root();
                self.weak.set(false);
            },
            trace::SignalKind::Unroot => {
                assert!(self.weak.get(), "cannot double-unroot a Gc");

                self.borrow().inner.as_ref().unroot();
                self.weak.set(true);
            },
        }
    }
}

impl<'arena, T: Trace<'arena, T>> Clone for Gc<'arena, T> {
    fn clone(&self) -> Self {
        unsafe { self.ptr.get().as_ref().root() }
        Self { weak: Cell::new(false), ptr: self.ptr.clone() }
    }
}

impl<'arena, T: Trace<'arena, T>> Drop for Gc<'arena, T> {
    fn drop(&mut self) {
        if !self.is_weak() {
            unsafe { self.ptr.get().as_ref().unroot() }
        }
    }
}

impl<'arena, T: Debug + Trace<'arena, T>> Debug for Gc<'arena, T> {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        self.borrow().fmt(f)
    }
}

impl<'arena, T: Display + Trace<'arena, T>> Display for Gc<'arena, T> {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        self.borrow().fmt(f)
    }
}

impl<'arena, T: Eq + Trace<'arena, T>> Eq for Gc<'arena, T> {}

impl<'arena, T: Ord + Trace<'arena, T>> Ord for Gc<'arena, T> {
    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
        self.borrow().cmp(&other.borrow())
    }
}

impl<'arena, T: PartialEq + Trace<'arena, T>> PartialEq for Gc<'arena, T> {
    fn eq(&self, other: &Self) -> bool {
        *self.borrow() == *other.borrow()
    }
}

impl<'arena, T: PartialOrd + Trace<'arena, T>> PartialOrd for Gc<'arena, T> {
    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
        self.borrow().partial_cmp(&other.borrow())
    }
}

struct ArenaState<'arena, T: Trace<'arena, T>> {
    boxes: Option<NonNull<Allocation<'arena, T>>>,
    allocated_bytes: usize,
}

/// A garbage collected arena of memory.
pub struct Arena<'arena, T: Trace<'arena, T> + Sized> {
    state: RwLock<ArenaState<'arena, T>>,
}

unsafe impl<'arena, T: Trace<'arena, T>> Send for Arena<'arena, T> {}
unsafe impl<'arena, T: Trace<'arena, T>> Sync for Arena<'arena, T> {}

impl<'arena, T: Trace<'arena, T> + Sized> Arena<'arena, T> {
    /// Creates an empty [`Arena`].
    pub fn new() -> Self {
        Self { state: RwLock::new(ArenaState { boxes: None, allocated_bytes: 0 }) }
    }

    /// Stores a value in this [`Arena`] and returns a rooted reference to it.
    pub fn host(&self, data: T) -> Gc<'arena, T> {
        let mut state = self.state.write().expect("an Arena's state cannot be locked");

        unsafe {
            // // it's being stored on the heap, it cannot be rooted.
            // data.send(&Signal::new(trace::SignalKind::Unroot));
            let allocation = Allocation {
                safety: PhantomData,
                guard: RwLock::new(()),
                roots: Cell::new(1),
                mark: Cell::new(false),
                next: state.boxes,
                data
            };
            let ptr = NonNull::new_unchecked(Box::into_raw(Box::new(allocation)));
            state.boxes = Some(ptr);
            state.allocated_bytes += size_of::<Allocation<T>>();
            Gc { weak: Cell::new(false), ptr: Cell::new(ptr) }
        }
    }

    /// Performs a garbage collection sweep.
    pub fn collect(&self) {
        let mut state = self.state.write().expect("an Arena's state cannot be locked");

        unsafe {
            // Mark
            {
                let mut next_allocation = state.boxes;
                while let Some(allocation) = next_allocation {
                    let item = allocation.as_ref();

                    if item.roots.get() > 0 {
                        item.mark.set(true);
                        item.data.send(&Signal::new(trace::SignalKind::Mark));
                    }

                    next_allocation = item.next;
                }
            }

            // Sweep
            {
                let mut next_allocation = &mut state.boxes;
                while let Some(mut allocation) = *next_allocation {
                    let item = allocation.as_mut();

                    if !item.mark.get() {
                        *next_allocation = item.next.clone();
                        item.data.finalize();
                        Box::from_raw(allocation.as_mut()); // implicitly free
                    } else {
                        item.mark.set(false);
                        next_allocation = &mut item.next;
                    }
                }
            }
        }
    }
}

impl<'arena, T: Trace<'arena, T> + Sized> Drop for Arena<'arena, T> {
    fn drop(&mut self) {
        let state = self.state.write().expect("an Arena's state cannot be locked");
        let mut next = state.boxes;

        unsafe {
            while let Some(mut allocation) = next {
                // assume that the data doesn't need to be locked as the lifetime of the data at most as long as the Arena's.
                allocation.as_mut().data.finalize();
                next = allocation.as_ref().next;
                Box::from_raw(allocation.as_ptr()); // implicitly free
            }
        }
    }
}