lazy-cow 0.1.0

Copy-on-write pointers with lazy modification support to minimise clones with a cost counter to limit work duplication.
Documentation
use std::{hint::unreachable_unchecked, ptr::NonNull, sync::atomic::{AtomicBool, AtomicUsize, Ordering}};

pub trait LazyModifiable: Clone {
    /// A difference type containing lazy edits.
    type Diff: Clone + Default;

    fn budget(&self) -> usize;

    fn cost(work: &Self::Diff) -> usize;

    fn fastforward(&mut self, work: Self::Diff);

    fn fastforward_from(&self, work: Self::Diff) -> Self {
        let mut other = self.clone();
        other.fastforward(work);
        other
    }

    fn fastforward_with(&self, work: &Self::Diff) -> Self {
        self.fastforward_from(work.clone())
    }
}

impl<T: Clone> LazyModifiable for Vec<T> {
    type Diff = Vec<T>;

    fn budget(&self) -> usize {
        self.len()
    }

    fn cost(work: &Self::Diff) -> usize {
        work.len()
    }

    fn fastforward(&mut self, mut work: Self::Diff) {
        self.append(&mut work);
    }

    fn fastforward_from(&self, work: Self::Diff) -> Self {
        self.iter().map(Clone::clone).chain(work.into_iter()).collect()
    }

    fn fastforward_with(&self, work: &Self::Diff) -> Self {
        self.iter().chain(work.iter()).map(Clone::clone).collect()
    }
}

struct Inner<T: LazyModifiable> {
    strong: AtomicUsize,
    spent: AtomicUsize,
    data: T
}

impl<T: LazyModifiable> Inner<T> {
    pub fn new(value: T) -> Inner<T> {
        Inner {
            strong: 1.into(),
            spent: 0.into(),
            data: value,
        }
    }
}

enum State<T: LazyModifiable> {
    Strict {
        is_unique: AtomicBool
    },
    Lazy(T::Diff)
}

pub enum LazyView<'a, T: LazyModifiable> {
    Strict(&'a T),
    Lazy(&'a T, &'a T::Diff)
}

pub enum LazyViewMut<'a, T: LazyModifiable> {
    Strict(&'a mut T),
    Lazy(&'a T, &'a mut T::Diff)
}

pub struct LazyCow<T: LazyModifiable> {
    state: State<T>,
    share: NonNull<Inner<T>>,
}

impl<T: LazyModifiable> LazyCow<T> {
    pub fn new(value: T) -> LazyCow<T> {
        LazyCow {
            state: State::Strict { is_unique: true.into() },
            share: unsafe { NonNull::new_unchecked(Box::into_raw(Box::new(Inner::new(value)))) },
        }
    }

    fn get_inner(&self) -> &T {
        &unsafe { self.share.as_ref() }.data
    }

    fn increment_strong(&self) {
        unsafe { self.share.as_ref() }.strong.fetch_add(1, Ordering::Relaxed);
    }

    fn decrement_strong(&mut self) {
        if unsafe { self.share.as_ref() }.strong.fetch_sub(1, Ordering::Release) == 1 {
            drop(unsafe { Box::from_raw(self.share.as_ptr()) })
        }
    }

    fn try_unique(&mut self) -> Result<&mut T, &mut Self> {
        match &mut self.state {
            State::Strict { is_unique } => if *is_unique.get_mut() {
                Ok(&mut unsafe { self.share.as_mut() }.data)
            } else {
                if unsafe { self.share.as_ref() }.strong.load(Ordering::Acquire) == 1 {
                    *is_unique.get_mut() = true;
                    *unsafe { self.share.as_mut() }.spent.get_mut() = 0;
                    Ok(&mut unsafe { self.share.as_mut() }.data)
                } else {
                    Err(self)
                }
            },
            State::Lazy(_) => {
                if unsafe { self.share.as_ref() }.strong.load(Ordering::Acquire) == 1 {
                    let State::Lazy(work) = std::mem::replace(&mut self.state, State::Strict { is_unique: true.into() }) else {
                        unsafe { unreachable_unchecked() }
                    };
                    unsafe { self.share.as_mut() }.data.fastforward(work);
                    *unsafe { self.share.as_mut() }.spent.get_mut() = 0;
                    Ok(&mut unsafe { self.share.as_mut() }.data)
                } else {
                    Err(self)
                }
            },
        }
    }

    pub fn get(&self) -> LazyView<'_, T> {
        match &self.state {
            State::Strict {..} => LazyView::Strict(self.get_inner()),
            State::Lazy(work) => LazyView::Lazy(self.get_inner(), work),
        }
    }

    pub fn get_mut(&mut self) -> LazyViewMut<'_, T> {
        match self.try_unique() {
            Ok(ref_mut) => LazyViewMut::Strict(ref_mut),
            Err(this) => match &mut this.state {
                state@State::Strict { .. } => {
                    let State::Strict { is_unique } = state else {
                        unsafe { unreachable_unchecked() }
                    };
                    if *is_unique.get_mut() {
                        LazyViewMut::Strict(&mut unsafe { this.share.as_mut() }.data)
                    } else {
                        *state = State::Lazy(T::Diff::default());
                        let State::Lazy(work) = state else {
                            unsafe { unreachable_unchecked() }
                        };
                        LazyViewMut::Lazy(&unsafe { this.share.as_ref() }.data, work)
                    }
                },
                State::Lazy(work) => {
                    LazyViewMut::Lazy(&unsafe { this.share.as_ref() }.data, work)
                },
            }
        }
    }

    pub fn get_strict(&mut self) -> &mut T {
        match self.try_unique() {
            Ok(ref_mut) => ref_mut,
            Err(this) => match &mut this.state {
                State::Strict { is_unique } => if *is_unique.get_mut() {
                    &mut unsafe { this.share.as_mut() }.data
                } else {
                    let value = this.get_inner().clone();
                    this.decrement_strong();
                    let mut unique = unsafe { NonNull::new_unchecked(Box::into_raw(Box::new(Inner::new(value)))) };
                    this.share = unique;
                    this.state = State::Strict { is_unique: true.into() };
                    &mut unsafe { unique.as_mut() }.data
                },
                State::Lazy(_) => {
                    let State::Lazy(work) = std::mem::replace(
                        &mut this.state,
                        State::Strict { is_unique: false.into() }
                    ) else {
                        unsafe { unreachable_unchecked() }
                    };
                    let value = this.get_inner().fastforward_from(work);
                    this.decrement_strong();
                    let mut unique = unsafe { NonNull::new_unchecked(Box::into_raw(Box::new(Inner::new(value)))) };
                    this.share = unique;
                    this.state = State::Strict { is_unique: true.into() };
                    &mut unsafe { unique.as_mut() }.data
                },
            }
        }
    }

    pub fn into_inner(mut self) -> T {
        match std::mem::replace(&mut self.state, State::Strict { is_unique: false.into() }) {
            State::Strict { is_unique } => if is_unique.into_inner() {
                let ptr = self.share.as_ptr();
                std::mem::forget(self);
                unsafe { Box::from_raw(ptr) }.data
            } else {
                if unsafe { self.share.as_ref() }.strong.load(Ordering::Acquire) == 1 {
                    let ptr = self.share.as_ptr();
                    std::mem::forget(self);
                    unsafe { Box::from_raw(ptr) }.data
                } else {
                    unsafe { self.share.as_ref() }.data.clone()
                }
            },
            State::Lazy(work) => {
                if unsafe { self.share.as_ref() }.strong.load(Ordering::Acquire) == 1 {
                    let ptr = self.share.as_ptr();
                    std::mem::forget(self);
                    let mut value = unsafe { Box::from_raw(ptr) }.data;
                    value.fastforward(work);
                    value
                } else {
                    unsafe { self.share.as_ref() }.data.fastforward_from(work)
                }
            },
        }
    }

    pub fn eval(&self) -> T {
        match self.get() {
            LazyView::Strict(t) => t.clone(),
            LazyView::Lazy(t, work) => t.fastforward_with(work),
        }
    }
}

impl<T: LazyModifiable> Clone for LazyCow<T> {
    fn clone(&self) -> Self {
        match &self.state {
            State::Strict { is_unique } => {
                is_unique.store(false, Ordering::Relaxed);
                self.increment_strong();
                LazyCow {
                    state: State::Strict { is_unique: false.into() },
                    share: self.share,
                }
            },
            State::Lazy(work) => {
                let inner = unsafe { self.share.as_ref() };
                let cost = T::cost(work);
                if loop {
                    let spent = inner.spent.load(Ordering::Relaxed);
                    let floating_spent = spent.saturating_add(cost);
                    if floating_spent > inner.data.budget() {
                        break false;
                    }
                    if let Ok(..) = inner.spent.compare_exchange_weak(spent, floating_spent, Ordering::Relaxed, Ordering::Relaxed) {
                        break true;
                    }
                } {
                    let work = work.clone();
                    self.increment_strong();
                    LazyCow {
                        state: State::Lazy(work),
                        share: self.share,
                    }
                } else {
                    let value = inner.data.fastforward_with(work);
                    LazyCow {
                        state: State::Strict { is_unique: true.into() },
                        share: unsafe { NonNull::new_unchecked(Box::into_raw(Box::new(Inner::new(value)))) },
                    }
                }
            },
        }
    }
}

impl<T: LazyModifiable> Drop for LazyCow<T> {
    fn drop(&mut self) {
        self.decrement_strong();
    }
}