1use std::{hint::unreachable_unchecked, ptr::NonNull, sync::atomic::{AtomicBool, AtomicUsize, Ordering}};
2
3pub trait LazyModifiable: Clone {
4 type Diff: Clone + Default;
6
7 fn budget(&self) -> usize;
8
9 fn cost(work: &Self::Diff) -> usize;
10
11 fn fastforward(&mut self, work: Self::Diff);
12
13 fn fastforward_from(&self, work: Self::Diff) -> Self {
14 let mut other = self.clone();
15 other.fastforward(work);
16 other
17 }
18
19 fn fastforward_with(&self, work: &Self::Diff) -> Self {
20 self.fastforward_from(work.clone())
21 }
22}
23
24impl<T: Clone> LazyModifiable for Vec<T> {
25 type Diff = Vec<T>;
26
27 fn budget(&self) -> usize {
28 self.len()
29 }
30
31 fn cost(work: &Self::Diff) -> usize {
32 work.len()
33 }
34
35 fn fastforward(&mut self, mut work: Self::Diff) {
36 self.append(&mut work);
37 }
38
39 fn fastforward_from(&self, work: Self::Diff) -> Self {
40 self.iter().map(Clone::clone).chain(work.into_iter()).collect()
41 }
42
43 fn fastforward_with(&self, work: &Self::Diff) -> Self {
44 self.iter().chain(work.iter()).map(Clone::clone).collect()
45 }
46}
47
48struct Inner<T: LazyModifiable> {
49 strong: AtomicUsize,
50 spent: AtomicUsize,
51 data: T
52}
53
54impl<T: LazyModifiable> Inner<T> {
55 pub fn new(value: T) -> Inner<T> {
56 Inner {
57 strong: 1.into(),
58 spent: 0.into(),
59 data: value,
60 }
61 }
62}
63
64enum State<T: LazyModifiable> {
65 Strict {
66 is_unique: AtomicBool
67 },
68 Lazy(T::Diff)
69}
70
71pub enum LazyView<'a, T: LazyModifiable> {
72 Strict(&'a T),
73 Lazy(&'a T, &'a T::Diff)
74}
75
76pub enum LazyViewMut<'a, T: LazyModifiable> {
77 Strict(&'a mut T),
78 Lazy(&'a T, &'a mut T::Diff)
79}
80
81pub struct LazyCow<T: LazyModifiable> {
82 state: State<T>,
83 share: NonNull<Inner<T>>,
84}
85
86impl<T: LazyModifiable> LazyCow<T> {
87 pub fn new(value: T) -> LazyCow<T> {
88 LazyCow {
89 state: State::Strict { is_unique: true.into() },
90 share: unsafe { NonNull::new_unchecked(Box::into_raw(Box::new(Inner::new(value)))) },
91 }
92 }
93
94 fn get_inner(&self) -> &T {
95 &unsafe { self.share.as_ref() }.data
96 }
97
98 fn increment_strong(&self) {
99 unsafe { self.share.as_ref() }.strong.fetch_add(1, Ordering::Relaxed);
100 }
101
102 fn decrement_strong(&mut self) {
103 if unsafe { self.share.as_ref() }.strong.fetch_sub(1, Ordering::Release) == 1 {
104 drop(unsafe { Box::from_raw(self.share.as_ptr()) })
105 }
106 }
107
108 fn try_unique(&mut self) -> Result<&mut T, &mut Self> {
109 match &mut self.state {
110 State::Strict { is_unique } => if *is_unique.get_mut() {
111 Ok(&mut unsafe { self.share.as_mut() }.data)
112 } else {
113 if unsafe { self.share.as_ref() }.strong.load(Ordering::Acquire) == 1 {
114 *is_unique.get_mut() = true;
115 *unsafe { self.share.as_mut() }.spent.get_mut() = 0;
116 Ok(&mut unsafe { self.share.as_mut() }.data)
117 } else {
118 Err(self)
119 }
120 },
121 State::Lazy(_) => {
122 if unsafe { self.share.as_ref() }.strong.load(Ordering::Acquire) == 1 {
123 let State::Lazy(work) = std::mem::replace(&mut self.state, State::Strict { is_unique: true.into() }) else {
124 unsafe { unreachable_unchecked() }
125 };
126 unsafe { self.share.as_mut() }.data.fastforward(work);
127 *unsafe { self.share.as_mut() }.spent.get_mut() = 0;
128 Ok(&mut unsafe { self.share.as_mut() }.data)
129 } else {
130 Err(self)
131 }
132 },
133 }
134 }
135
136 pub fn get(&self) -> LazyView<'_, T> {
137 match &self.state {
138 State::Strict {..} => LazyView::Strict(self.get_inner()),
139 State::Lazy(work) => LazyView::Lazy(self.get_inner(), work),
140 }
141 }
142
143 pub fn get_mut(&mut self) -> LazyViewMut<'_, T> {
144 match self.try_unique() {
145 Ok(ref_mut) => LazyViewMut::Strict(ref_mut),
146 Err(this) => match &mut this.state {
147 state@State::Strict { .. } => {
148 let State::Strict { is_unique } = state else {
149 unsafe { unreachable_unchecked() }
150 };
151 if *is_unique.get_mut() {
152 LazyViewMut::Strict(&mut unsafe { this.share.as_mut() }.data)
153 } else {
154 *state = State::Lazy(T::Diff::default());
155 let State::Lazy(work) = state else {
156 unsafe { unreachable_unchecked() }
157 };
158 LazyViewMut::Lazy(&unsafe { this.share.as_ref() }.data, work)
159 }
160 },
161 State::Lazy(work) => {
162 LazyViewMut::Lazy(&unsafe { this.share.as_ref() }.data, work)
163 },
164 }
165 }
166 }
167
168 pub fn get_strict(&mut self) -> &mut T {
169 match self.try_unique() {
170 Ok(ref_mut) => ref_mut,
171 Err(this) => match &mut this.state {
172 State::Strict { is_unique } => if *is_unique.get_mut() {
173 &mut unsafe { this.share.as_mut() }.data
174 } else {
175 let value = this.get_inner().clone();
176 this.decrement_strong();
177 let mut unique = unsafe { NonNull::new_unchecked(Box::into_raw(Box::new(Inner::new(value)))) };
178 this.share = unique;
179 this.state = State::Strict { is_unique: true.into() };
180 &mut unsafe { unique.as_mut() }.data
181 },
182 State::Lazy(_) => {
183 let State::Lazy(work) = std::mem::replace(
184 &mut this.state,
185 State::Strict { is_unique: false.into() }
186 ) else {
187 unsafe { unreachable_unchecked() }
188 };
189 let value = this.get_inner().fastforward_from(work);
190 this.decrement_strong();
191 let mut unique = unsafe { NonNull::new_unchecked(Box::into_raw(Box::new(Inner::new(value)))) };
192 this.share = unique;
193 this.state = State::Strict { is_unique: true.into() };
194 &mut unsafe { unique.as_mut() }.data
195 },
196 }
197 }
198 }
199
200 pub fn into_inner(mut self) -> T {
201 match std::mem::replace(&mut self.state, State::Strict { is_unique: false.into() }) {
202 State::Strict { is_unique } => if is_unique.into_inner() {
203 let ptr = self.share.as_ptr();
204 std::mem::forget(self);
205 unsafe { Box::from_raw(ptr) }.data
206 } else {
207 if unsafe { self.share.as_ref() }.strong.load(Ordering::Acquire) == 1 {
208 let ptr = self.share.as_ptr();
209 std::mem::forget(self);
210 unsafe { Box::from_raw(ptr) }.data
211 } else {
212 unsafe { self.share.as_ref() }.data.clone()
213 }
214 },
215 State::Lazy(work) => {
216 if unsafe { self.share.as_ref() }.strong.load(Ordering::Acquire) == 1 {
217 let ptr = self.share.as_ptr();
218 std::mem::forget(self);
219 let mut value = unsafe { Box::from_raw(ptr) }.data;
220 value.fastforward(work);
221 value
222 } else {
223 unsafe { self.share.as_ref() }.data.fastforward_from(work)
224 }
225 },
226 }
227 }
228
229 pub fn eval(&self) -> T {
230 match self.get() {
231 LazyView::Strict(t) => t.clone(),
232 LazyView::Lazy(t, work) => t.fastforward_with(work),
233 }
234 }
235}
236
237impl<T: LazyModifiable> Clone for LazyCow<T> {
238 fn clone(&self) -> Self {
239 match &self.state {
240 State::Strict { is_unique } => {
241 is_unique.store(false, Ordering::Relaxed);
242 self.increment_strong();
243 LazyCow {
244 state: State::Strict { is_unique: false.into() },
245 share: self.share,
246 }
247 },
248 State::Lazy(work) => {
249 let inner = unsafe { self.share.as_ref() };
250 let cost = T::cost(work);
251 if loop {
252 let spent = inner.spent.load(Ordering::Relaxed);
253 let floating_spent = spent.saturating_add(cost);
254 if floating_spent > inner.data.budget() {
255 break false;
256 }
257 if let Ok(..) = inner.spent.compare_exchange_weak(spent, floating_spent, Ordering::Relaxed, Ordering::Relaxed) {
258 break true;
259 }
260 } {
261 let work = work.clone();
262 self.increment_strong();
263 LazyCow {
264 state: State::Lazy(work),
265 share: self.share,
266 }
267 } else {
268 let value = inner.data.fastforward_with(work);
269 LazyCow {
270 state: State::Strict { is_unique: true.into() },
271 share: unsafe { NonNull::new_unchecked(Box::into_raw(Box::new(Inner::new(value)))) },
272 }
273 }
274 },
275 }
276 }
277}
278
279impl<T: LazyModifiable> Drop for LazyCow<T> {
280 fn drop(&mut self) {
281 self.decrement_strong();
282 }
283}