use std::{hint::unreachable_unchecked, ptr::NonNull, sync::atomic::{AtomicBool, AtomicUsize, Ordering}};
pub trait LazyModifiable: Clone {
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();
}
}