cycle_ptr 0.1.0

//! Implement the strong reference counter for objects.
//!
//! The strong reference counter tracks the number of strong references,
//! as well as the [Colour] and [ObjectState] of the object.
use crate::generation::colour::{Colour, GCInitialColour};
use std::cell::RefCell;
use std::fmt::Debug;

#[cfg(feature = "weak_pointer")]
use crate::errors::{Error, ErrorEnum};
#[cfg(feature = "weak_pointer")]
use std::backtrace::Backtrace;
#[cfg(feature = "weak_pointer")]
use std::sync::Arc;
#[cfg(feature = "multi_thread")]
use std::sync::atomic::{AtomicUsize, Ordering};

/// Counters (`refcount_inc`/`refcount_dec`) must skip the lowest two bits.
/// So an increment of `1`, is really an increment of `1 << 2`.
/// This constant ensures the shift skips all the colour bits.
const COUNT_SHIFT: u32 = 2;
/// Bit-mask to extract the colour from a refcount number.
const COLOUR_MASK: usize = 0x3_usize;

/// Not actually a colour, but since expired elements don't have a meaningful colour, and because the colours only use 3 of the 4 values in the [COLOUR_MASK], we use the spare value as expired.
const EXPIRED_COLOUR: usize = 0x0_usize;
/// White colour indicates an element isn't reachable from a black coloured element.
const WHITE_COLOUR: usize = 0x1_usize;
/// Grey colour indicates an element is reachable.
const GREY_COLOUR: usize = 0x2_usize;
/// Black colour indicates an element is reachable, and none of the elements directly reachable from it, are white.
const BLACK_COLOUR: usize = 0x3_usize;

/// Outcome of a reference count decrement.
#[derive(Clone, Copy, PartialEq, Eq)]
pub(crate) enum RefCountResult {
    /// Indicates the element may have become unreachable.
    NeedGC,
    /// Indicates the element is reachable.
    NoNeed,
}

/// The liveness-state of an object.
#[derive(PartialEq, Eq, Debug)]
pub(crate) enum ObjectState {
    /// The object is valid and can safely be used.
    Live,
    /// The object is expired and is about-to-be (or already-has-been) destroyed.
    Expired,
}

/// Reference counter for control blocks.
pub(crate) struct RefCount {
    /// Internal value of the reference counter.
    ///
    /// Tracks reference count, and uses the lowest two bits to track colour and state.
    v: RefCell<usize>,
}

/// Reference counter for control blocks.
#[cfg(feature = "multi_thread")]
pub(crate) struct MTRefCount {
    /// Internal value of the reference counter.
    ///
    /// Tracks reference count, and uses the lowest two bits to track colour and state.
    v: AtomicUsize,
}

impl Default for RefCount {
    fn default() -> Self {
        RefCount {
            v: RefCell::new((1_usize << COUNT_SHIFT) | BLACK_COLOUR),
        }
    }
}

#[cfg(feature = "multi_thread")]
impl Default for MTRefCount {
    fn default() -> Self {
        MTRefCount {
            v: AtomicUsize::new((1_usize << COUNT_SHIFT) | BLACK_COLOUR),
        }
    }
}

impl Debug for RefCount {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> Result<(), std::fmt::Error> {
        let (refs, colour, state) = self.load();
        f.debug_struct("RefCount")
            .field("references", &refs)
            .field("colour", &colour)
            .field("state", &state)
            .finish()
    }
}

#[cfg(feature = "multi_thread")]
impl Debug for MTRefCount {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> Result<(), std::fmt::Error> {
        let (refs, colour, state) = self.load();
        f.debug_struct("MTRefCount")
            .field("references", &refs)
            .field("colour", &colour)
            .field("state", &state)
            .finish()
    }
}

impl RefCount {
    /// Create a new [RefCount] with initial refcount value.
    ///
    /// The returned [RefCount] will be [Colour::Black].
    #[inline]
    pub(super) const fn new(initial_refcount: usize) -> Self {
        RefCount {
            v: RefCell::new((initial_refcount << COUNT_SHIFT) | BLACK_COLOUR),
        }
    }

    /// Read the values in the reference counter.
    #[inline]
    pub(crate) fn load(&self) -> (usize, Colour, ObjectState) {
        let v = *self.v.borrow();
        let (colour, state) = match v & COLOUR_MASK {
            EXPIRED_COLOUR => (Colour::White, ObjectState::Expired),
            WHITE_COLOUR => (Colour::White, ObjectState::Live),
            GREY_COLOUR => (Colour::Grey, ObjectState::Live),
            BLACK_COLOUR => (Colour::Black, ObjectState::Live),
            _ => panic!("Colour constants are wrong."),
        };
        let refcount = v >> COUNT_SHIFT;
        (refcount, colour, state)
    }

    /// Reset the colour for the start of a GC cycle.
    ///
    /// Changes the colour to grey, if the object is reachable.
    /// Otherwise changes it to white.
    /// Returns the colour that was assigned.
    pub(super) fn gc_colour_reset(&self) -> GCInitialColour {
        let mut colour = GCInitialColour::White;
        self.v.replace_with(|v| {
            let mut v = *v;
            assert_ne!((v & COLOUR_MASK), EXPIRED_COLOUR);
            v &= !COLOUR_MASK;
            if (v & !COLOUR_MASK) == 0 {
                v |= WHITE_COLOUR;
                colour = GCInitialColour::White;
            } else {
                v |= GREY_COLOUR;
                colour = GCInitialColour::Grey;
            }
            v
        });
        colour
    }

    /// Mark an object as invalid.
    pub(super) fn invalidate(&self) {
        self.v.replace_with(|v| {
            let mut v = *v;
            assert_ne!(v & COLOUR_MASK, EXPIRED_COLOUR);
            v &= !COLOUR_MASK;
            v |= EXPIRED_COLOUR;
            v
        });
    }

    /// Attempt to mark an object as invalid.
    ///
    /// Returns a result indicating if we marked the object as invalid.
    pub(super) fn try_invalidate(&self) -> Result<(), ()> {
        let mut v = self.v.borrow_mut();
        match *v & COLOUR_MASK {
            EXPIRED_COLOUR => Err(()),
            _ => {
                *v &= !COLOUR_MASK;
                *v |= EXPIRED_COLOUR;
                Ok(())
            }
        }
    }

    /// Try to increment the reference counter.
    ///
    /// If the element is expired, the operation will fail.
    /// Ensures that if the reference counter is white, it'll be coloured grey.
    #[cfg(feature = "weak_pointer")]
    pub(super) fn try_inc(
        &self,
        has_data: impl FnOnce() -> bool,
        opt_backtrace: &Option<Arc<Backtrace>>,
    ) -> Result<(), Error> {
        let mut v: usize = *self.v.borrow();
        if (v & COLOUR_MASK) == EXPIRED_COLOUR {
            return Err(Error::new(ErrorEnum::Expired(opt_backtrace.clone())));
        }

        if !has_data() {
            return Err(Error::new(ErrorEnum::NotYetInitialized(
                opt_backtrace.clone(),
            )));
        }

        v = v.strict_add(1 << COUNT_SHIFT);
        if (v & COLOUR_MASK) == WHITE_COLOUR {
            v &= !COLOUR_MASK;
            v |= GREY_COLOUR;
        }
        *self.v.borrow_mut() = v;
        Ok(())
    }

    /// Increment the reference counter.
    ///
    /// Ensures that if the reference counter is white, it'll be coloured grey.
    pub(super) fn inc(&self) {
        self.v.replace_with(|v| {
            let mut v = (*v).strict_add(1 << COUNT_SHIFT);
            if (v & COLOUR_MASK) == WHITE_COLOUR {
                v &= !COLOUR_MASK;
                v |= GREY_COLOUR;
            }
            v
        });
    }

    /// Decrease the reference counter.
    ///
    /// Will return a [RefCountResult], which indicates if a GC is required.
    pub(super) fn dec(&self) -> RefCountResult {
        let mut result = RefCountResult::NoNeed;
        self.v.replace_with(|v| {
            let v = (*v).strict_sub(1 << COUNT_SHIFT);
            if (v & !COLOUR_MASK) == 0 {
                result = RefCountResult::NeedGC;
            }
            v
        });
        result
    }

    /// Decrease the reference counter, and ensure it's marked [Colour::Grey] if it is [Colour::White].
    pub(super) fn dec_reachable(&self) {
        self.v.replace_with(|v| {
            assert_ne!(*v & COLOUR_MASK, WHITE_COLOUR);
            assert_ne!(*v & COLOUR_MASK, EXPIRED_COLOUR);
            (*v).strict_sub(1 << COUNT_SHIFT)
        });
    }

    /// Mark an object as [Colour::Black].
    ///
    /// Requires the object is [Colour::Grey].
    pub(super) fn mark_black(&self) {
        self.v.replace_with(|v| {
            let mut v = *v;
            assert_eq!(v & COLOUR_MASK, GREY_COLOUR);
            v &= !COLOUR_MASK;
            v |= BLACK_COLOUR;
            v
        });
    }

    /// Attempt to mark an object as [Colour::Grey].
    ///
    /// Requires the object is valid.
    /// Only succeeds if the object is [Colour::White].
    pub(super) fn try_mark_grey(&self) -> Result<(), ()> {
        let mut v = *self.v.borrow();
        assert_ne!(v & COLOUR_MASK, EXPIRED_COLOUR);
        match v & COLOUR_MASK {
            WHITE_COLOUR => {
                v &= !COLOUR_MASK;
                v |= GREY_COLOUR;
                *self.v.borrow_mut() = v;
                Ok(())
            }
            _ => Err(()),
        }
    }
}

#[cfg(feature = "multi_thread")]
impl MTRefCount {
    /// Create a new [RefCount] with initial refcount value.
    ///
    /// The returned [RefCount] will be [Colour::Black].
    #[inline]
    pub(super) const fn new(initial_refcount: usize) -> Self {
        MTRefCount {
            v: AtomicUsize::new((initial_refcount << COUNT_SHIFT) | BLACK_COLOUR),
        }
    }

    /// Read the values in the reference counter.
    #[inline]
    pub(crate) fn load(&self) -> (usize, Colour, ObjectState) {
        let v = self.v.load(Ordering::Relaxed);
        let (colour, state) = match v & COLOUR_MASK {
            EXPIRED_COLOUR => (Colour::White, ObjectState::Expired),
            WHITE_COLOUR => (Colour::White, ObjectState::Live),
            GREY_COLOUR => (Colour::Grey, ObjectState::Live),
            BLACK_COLOUR => (Colour::Black, ObjectState::Live),
            _ => panic!("Colour constants are wrong."),
        };
        let refcount = v >> COUNT_SHIFT;
        (refcount, colour, state)
    }

    /// Reset the colour for the start of a GC cycle.
    ///
    /// Changes the colour to grey, if the object is reachable.
    /// Otherwise changes it to white.
    /// Returns the colour that was assigned.
    ///
    /// Note: we don't bother returning a [MTGeneration::gc_lists][super::MTGeneration::gc_lists],
    /// because this should always be called with [MTGeneration::gc_lists][super::MTGeneration::gc_lists] locked,
    /// so the caller can handle whatever by itself.
    pub(super) fn gc_colour_reset(&self) -> GCInitialColour {
        let mut old_v: usize = BLACK_COLOUR;
        loop {
            assert_ne!((old_v & COLOUR_MASK), EXPIRED_COLOUR);
            let mut new_v = old_v & !COLOUR_MASK;
            let colour = if (old_v & !COLOUR_MASK) == 0 {
                new_v |= WHITE_COLOUR;
                GCInitialColour::White
            } else {
                new_v |= GREY_COLOUR;
                GCInitialColour::Grey
            };
            match self
                .v
                .compare_exchange_weak(old_v, new_v, Ordering::Relaxed, Ordering::Relaxed)
            {
                Ok(_) => break colour,
                Err(x) => old_v = x,
            };
        }
    }

    /// Mark an object as invalid.
    pub(super) fn invalidate(&self) {
        let mut old_v = WHITE_COLOUR;
        loop {
            assert_ne!((old_v & COLOUR_MASK), EXPIRED_COLOUR);
            let new_v = (old_v & !COLOUR_MASK) | EXPIRED_COLOUR;
            match self
                .v
                .compare_exchange_weak(old_v, new_v, Ordering::Relaxed, Ordering::Relaxed)
            {
                Ok(_) => break,
                Err(x) => old_v = x,
            }
        }
    }

    /// Attempt to mark an object as invalid.
    ///
    /// Returns a result indicating if we marked the object as invalid.
    pub(super) fn try_invalidate(&self) -> Result<(), ()> {
        let mut old_v = self.v.load(Ordering::Relaxed);
        loop {
            if (old_v & COLOUR_MASK) == EXPIRED_COLOUR {
                return Err(());
            }

            let mut new_v = old_v;
            new_v &= !COLOUR_MASK;
            new_v |= EXPIRED_COLOUR;

            match self
                .v
                .compare_exchange(old_v, new_v, Ordering::Relaxed, Ordering::Relaxed)
            {
                Ok(_) => break,
                Err(x) => old_v = x,
            };
        }
        Ok(())
    }

    /// Try to increment the reference counter.
    ///
    /// If the element is expired, the operation will fail.
    /// Ensures that if the reference counter is white, it'll be coloured grey.
    #[cfg(feature = "weak_pointer")]
    pub(super) fn try_inc<R>(
        &self,
        has_data: impl FnOnce() -> bool,
        lock_fn: impl Fn() -> R,
        opt_backtrace: &Option<Arc<Backtrace>>,
    ) -> Result<Option<R>, Error> {
        let mut old_v = self.v.load(Ordering::Relaxed);
        let mut has_data = Some(has_data);
        let weak_lockout = loop {
            if (old_v & COLOUR_MASK) == EXPIRED_COLOUR {
                return Err(Error::new(ErrorEnum::Expired(opt_backtrace.clone())));
            }
            if has_data.take().map(|has_data| !has_data()).unwrap_or(false) {
                return Err(Error::new(ErrorEnum::NotYetInitialized(
                    opt_backtrace.clone(),
                )));
            }

            let mut weak_lockout = None;
            let mut new_v = old_v.strict_add(1 << COUNT_SHIFT);
            if (old_v & COLOUR_MASK) == WHITE_COLOUR {
                new_v &= !COLOUR_MASK;
                new_v |= GREY_COLOUR;
                weak_lockout = Some(lock_fn());
            }

            match self
                .v
                .compare_exchange(old_v, new_v, Ordering::Relaxed, Ordering::Relaxed)
            {
                Ok(_) => break weak_lockout,
                Err(x) => old_v = x,
            };
        };
        Ok(weak_lockout)
    }

    /// Increment the reference counter.
    ///
    /// Ensures that if the reference counter is white, it'll be coloured grey.
    pub(super) fn inc<R>(&self, lock_fn: impl Fn() -> R) -> Option<R> {
        let mut old_v = self.v.load(Ordering::Relaxed);
        loop {
            assert_ne!(old_v & COLOUR_MASK, EXPIRED_COLOUR);
            let mut new_v = old_v.strict_add(1 << COUNT_SHIFT);
            let mut r = None;
            if (old_v & COLOUR_MASK) == WHITE_COLOUR {
                new_v &= !COLOUR_MASK;
                new_v |= GREY_COLOUR;
                r = Some(lock_fn());
            }

            match self
                .v
                .compare_exchange_weak(old_v, new_v, Ordering::Relaxed, Ordering::Relaxed)
            {
                Ok(_) => break r,
                Err(x) => old_v = x,
            }
        }
    }

    /// Increment the reference counter.
    ///
    /// Uses a faster method, that's only usable if we are absolutely certain we have a non-zero reference counter.
    #[inline]
    pub(super) fn inc_strong(&self) {
        let old_v = self.v.fetch_add(1 << COUNT_SHIFT, Ordering::Relaxed);
        assert_ne!(old_v & COLOUR_MASK, EXPIRED_COLOUR);
        assert_ne!(old_v & !COLOUR_MASK, 0);
    }

    /// Decrease the reference counter.
    ///
    /// Will return a [RefCountResult], which indicates if a GC is required.
    pub(super) fn dec(&self) -> RefCountResult {
        let old_v = self.v.fetch_sub(1 << COUNT_SHIFT, Ordering::Relaxed);
        assert_ne!(old_v & !COLOUR_MASK, 0);
        if (old_v & !COLOUR_MASK) == 1 << COUNT_SHIFT {
            RefCountResult::NeedGC
        } else {
            RefCountResult::NoNeed
        }
    }

    /// Decrease the reference counter.
    pub(super) fn dec_reachable(&self) {
        let old_v = self.v.fetch_sub(1 << COUNT_SHIFT, Ordering::Relaxed);
        assert_ne!(old_v & !COLOUR_MASK, 0);
        assert_ne!(old_v & COLOUR_MASK, WHITE_COLOUR);
        assert_ne!(old_v & COLOUR_MASK, EXPIRED_COLOUR);
    }

    /// Mark an object as [Colour::Black].
    ///
    /// Requires the object is [Colour::Grey].
    pub(super) fn mark_black(&self) {
        let mut old_v = self.v.load(Ordering::Relaxed);
        loop {
            assert_eq!(old_v & COLOUR_MASK, GREY_COLOUR);
            let mut new_v = old_v;
            new_v &= !COLOUR_MASK;
            new_v |= BLACK_COLOUR;
            match self
                .v
                .compare_exchange_weak(old_v, new_v, Ordering::Relaxed, Ordering::Relaxed)
            {
                Ok(_) => break,
                Err(x) => old_v = x,
            }
        }
    }

    /// Attempt to mark an object as [Colour::Grey].
    ///
    /// Requires the object is valid.
    /// Only succeeds if the object is [Colour::White].
    pub(super) fn try_mark_grey<R>(&self, lock_fn: impl Fn() -> R) -> Result<R, ()> {
        let mut old_v = WHITE_COLOUR;
        loop {
            assert_ne!(old_v & COLOUR_MASK, EXPIRED_COLOUR);
            match old_v & COLOUR_MASK {
                WHITE_COLOUR => {
                    let r = lock_fn();
                    let mut new_v = old_v;
                    new_v &= !COLOUR_MASK;
                    new_v |= GREY_COLOUR;
                    match self.v.compare_exchange_weak(
                        old_v,
                        new_v,
                        Ordering::Relaxed,
                        Ordering::Relaxed,
                    ) {
                        Ok(_) => return Ok(r),
                        Err(x) => old_v = x,
                    }
                }
                _ => return Err(()),
            }
        }
    }
}