cranpose_core/snapshot_v2/
mod.rs

1//! Snapshot system for managing isolated state changes.
2//!
3//! This module implements Jetpack Compose's snapshot isolation system, allowing
4//! state changes to be isolated, composed, and atomically applied.
5//!
6//! # Snapshot Types
7//!
8//! - **ReadonlySnapshot**: Immutable view of state at a point in time
9//! - **MutableSnapshot**: Allows isolated state mutations
10//! - **NestedReadonlySnapshot**: Readonly snapshot nested in a parent
11//! - **NestedMutableSnapshot**: Mutable snapshot nested in a parent
12//! - **GlobalSnapshot**: Special global mutable snapshot
13//! - **TransparentObserverMutableSnapshot**: Optimized for observer chaining
14//! - **TransparentObserverSnapshot**: Readonly version of transparent observer
15//!
16//! # Thread Local Storage
17//!
18//! The current snapshot is stored in thread-local storage and automatically
19//! managed by the snapshot system. Snapshot handles and observers use [`Rc`]
20//! and stay on their creating thread. Background work sends results to the UI
21//! runtime instead of sharing snapshots across threads.
22
23use std::{
24    cell::{Cell, RefCell},
25    rc::{Rc, Weak},
26};
27
28use crate::{
29    collections::map::{HashMap, HashSet},
30    snapshot_id_set::{SnapshotId, SnapshotIdSet},
31    snapshot_pinning::{self, PinHandle},
32    snapshot_weak_set::SnapshotWeakSetDebugStats,
33    state::{StateObject, StateRecord},
34};
35
36mod global;
37mod mutable;
38mod nested;
39mod readonly;
40mod runtime;
41mod transparent;
42
43#[cfg(test)]
44#[path = "tests/integration_tests.rs"]
45mod integration_tests;
46
47pub use global::{GlobalSnapshot, advance_global_snapshot};
48pub use mutable::MutableSnapshot;
49pub use nested::{NestedMutableSnapshot, NestedReadonlySnapshot};
50pub use readonly::ReadonlySnapshot;
51pub(crate) use runtime::{SnapshotRuntime, allocate_snapshot, close_snapshot, with_runtime};
52#[cfg(test)]
53pub(crate) use runtime::{TestRuntimeGuard, reset_runtime_for_tests};
54pub use transparent::{TransparentObserverMutableSnapshot, TransparentObserverSnapshot};
55
56/// Observer that is called when a state object is read.
57pub type ReadObserver = Rc<dyn Fn(&dyn StateObject) + 'static>;
58
59/// Observer that is called when a state object is written.
60pub type WriteObserver = Rc<dyn Fn(&dyn StateObject) + 'static>;
61
62/// Apply observer that is called when a snapshot is applied.
63pub type ApplyObserver = Rc<dyn Fn(&[Rc<dyn StateObject>], SnapshotId) + 'static>;
64
65/// Result of applying a mutable snapshot.
66#[derive(Debug, Clone, PartialEq, Eq)]
67pub enum SnapshotApplyResult {
68    /// The snapshot was applied successfully.
69    Success,
70    /// The snapshot could not be applied due to conflicts.
71    Failure,
72}
73
74impl SnapshotApplyResult {
75    /// Check if the result is successful.
76    pub fn is_success(&self) -> bool {
77        matches!(self, SnapshotApplyResult::Success)
78    }
79
80    /// Check if the result is a failure.
81    pub fn is_failure(&self) -> bool {
82        matches!(self, SnapshotApplyResult::Failure)
83    }
84
85    /// Panic if the result is a failure (for use in tests).
86    #[track_caller]
87    pub fn check(&self) {
88        assert!(!self.is_failure(), "Snapshot apply failed");
89    }
90}
91
92/// Unique identifier for a state object in the modified set.
93pub type StateObjectId = usize;
94
95/// Enum wrapper for all snapshot types.
96///
97/// This provides a type-safe way to work with different snapshot types
98/// without requiring trait objects, which avoids object-safety issues.
99#[derive(Clone)]
100pub enum AnySnapshot {
101    /// A root read-only snapshot.
102    Readonly(Rc<ReadonlySnapshot>),
103    /// A root mutable snapshot, applied to the global state.
104    Mutable(Rc<MutableSnapshot>),
105    /// A read-only snapshot nested in a parent snapshot.
106    NestedReadonly(Rc<NestedReadonlySnapshot>),
107    /// A mutable snapshot nested in a mutable parent, applied to that parent.
108    NestedMutable(Rc<NestedMutableSnapshot>),
109    /// The global snapshot of the thread.
110    Global(Rc<GlobalSnapshot>),
111    /// A mutable snapshot whose observers can be replaced after creation.
112    TransparentMutable(Rc<TransparentObserverMutableSnapshot>),
113    /// A read-only snapshot whose observer can be replaced after creation.
114    TransparentReadonly(Rc<TransparentObserverSnapshot>),
115}
116
117/// Enum wrapper for mutable snapshot types.
118///
119/// This allows `take_mutable_snapshot` to return either a root MutableSnapshot
120/// or a NestedMutableSnapshot depending on the current context, matching Kotlin's
121/// behavior where `takeMutableSnapshot` creates nested snapshots when inside a
122/// mutable snapshot.
123#[derive(Clone)]
124pub enum AnyMutableSnapshot {
125    /// A root mutable snapshot, applied to the global state.
126    Root(Rc<MutableSnapshot>),
127    /// A mutable snapshot nested in a mutable parent, applied to that parent.
128    Nested(Rc<NestedMutableSnapshot>),
129}
130
131impl AnyMutableSnapshot {
132    /// Get the snapshot ID.
133    pub fn snapshot_id(&self) -> SnapshotId {
134        match self {
135            AnyMutableSnapshot::Root(s) => s.snapshot_id(),
136            AnyMutableSnapshot::Nested(s) => s.snapshot_id(),
137        }
138    }
139
140    /// Get the set of invalid snapshot IDs.
141    pub fn invalid(&self) -> SnapshotIdSet {
142        match self {
143            AnyMutableSnapshot::Root(s) => s.invalid(),
144            AnyMutableSnapshot::Nested(s) => s.invalid(),
145        }
146    }
147
148    /// Enter this snapshot, making it current for the duration of the closure.
149    pub fn enter<T>(&self, f: impl FnOnce() -> T) -> T {
150        match self {
151            AnyMutableSnapshot::Root(s) => s.enter(f),
152            AnyMutableSnapshot::Nested(s) => s.enter(f),
153        }
154    }
155
156    /// Apply the snapshot.
157    pub fn apply(&self) -> SnapshotApplyResult {
158        match self {
159            AnyMutableSnapshot::Root(s) => s.apply(),
160            AnyMutableSnapshot::Nested(s) => s.apply(),
161        }
162    }
163
164    /// Dispose the snapshot.
165    pub fn dispose(&self) {
166        match self {
167            AnyMutableSnapshot::Root(s) => s.dispose(),
168            AnyMutableSnapshot::Nested(s) => s.dispose(),
169        }
170    }
171}
172
173impl AnySnapshot {
174    /// Get the snapshot ID.
175    pub fn snapshot_id(&self) -> SnapshotId {
176        match self {
177            AnySnapshot::Readonly(s) => s.snapshot_id(),
178            AnySnapshot::Mutable(s) => s.snapshot_id(),
179            AnySnapshot::NestedReadonly(s) => s.snapshot_id(),
180            AnySnapshot::NestedMutable(s) => s.snapshot_id(),
181            AnySnapshot::Global(s) => s.snapshot_id(),
182            AnySnapshot::TransparentMutable(s) => s.snapshot_id(),
183            AnySnapshot::TransparentReadonly(s) => s.snapshot_id(),
184        }
185    }
186
187    /// The state every snapshot kind shares: id, invalid set, observers and flags.
188    fn state(&self) -> &SnapshotState {
189        match self {
190            AnySnapshot::Readonly(s) => &s.state,
191            AnySnapshot::Mutable(s) => &s.state,
192            AnySnapshot::NestedReadonly(s) => &s.state,
193            AnySnapshot::NestedMutable(s) => &s.state,
194            AnySnapshot::Global(s) => &s.state,
195            AnySnapshot::TransparentMutable(s) => &s.state,
196            AnySnapshot::TransparentReadonly(s) => &s.state,
197        }
198    }
199
200    pub(crate) fn with_invalid<R>(&self, f: impl FnOnce(&SnapshotIdSet) -> R) -> R {
201        f(&self.state().invalid.borrow())
202    }
203
204    /// Returns the ids of snapshots whose writes this snapshot does not see.
205    pub fn invalid(&self) -> SnapshotIdSet {
206        match self {
207            AnySnapshot::Readonly(s) => s.invalid(),
208            AnySnapshot::Mutable(s) => s.invalid(),
209            AnySnapshot::NestedReadonly(s) => s.invalid(),
210            AnySnapshot::NestedMutable(s) => s.invalid(),
211            AnySnapshot::Global(s) => s.invalid(),
212            AnySnapshot::TransparentMutable(s) => s.invalid(),
213            AnySnapshot::TransparentReadonly(s) => s.invalid(),
214        }
215    }
216
217    /// Check if a snapshot ID is valid in this snapshot.
218    pub fn is_valid(&self, id: SnapshotId) -> bool {
219        let snapshot_id = self.snapshot_id();
220        id <= snapshot_id && !self.invalid().get(id)
221    }
222
223    /// Check if this is a read-only snapshot.
224    pub fn read_only(&self) -> bool {
225        match self {
226            AnySnapshot::Readonly(_) => true,
227            AnySnapshot::Mutable(_) => false,
228            AnySnapshot::NestedReadonly(_) => true,
229            AnySnapshot::NestedMutable(_) => false,
230            AnySnapshot::Global(_) => false,
231            AnySnapshot::TransparentMutable(_) => false,
232            AnySnapshot::TransparentReadonly(_) => true,
233        }
234    }
235
236    /// Get the root snapshot.
237    pub fn root(&self) -> AnySnapshot {
238        match self {
239            AnySnapshot::Readonly(s) => AnySnapshot::Readonly(s.root_readonly()),
240            AnySnapshot::Mutable(s) => AnySnapshot::Mutable(s.root_mutable()),
241            AnySnapshot::NestedReadonly(s) => AnySnapshot::NestedReadonly(s.root_nested_readonly()),
242            AnySnapshot::NestedMutable(s) => AnySnapshot::Mutable(s.root_mutable()),
243            AnySnapshot::Global(_) => AnySnapshot::Global(GlobalSnapshot::get_or_create()),
244            AnySnapshot::TransparentMutable(s) => {
245                AnySnapshot::TransparentMutable(s.root_transparent_mutable())
246            }
247            AnySnapshot::TransparentReadonly(s) => {
248                AnySnapshot::TransparentReadonly(s.root_transparent_readonly())
249            }
250        }
251    }
252
253    /// Check if this snapshot refers to the same transparent snapshot.
254    pub fn is_same_transparent(&self, other: &Rc<TransparentObserverMutableSnapshot>) -> bool {
255        matches!(self, AnySnapshot::TransparentMutable(snapshot) if Rc::ptr_eq(snapshot, other))
256    }
257
258    /// Check if this snapshot refers to the same transparent mutable snapshot.
259    pub fn is_same_transparent_mutable(
260        &self,
261        other: &Rc<TransparentObserverMutableSnapshot>,
262    ) -> bool {
263        self.is_same_transparent(other)
264    }
265
266    /// Check if this snapshot refers to the same transparent readonly snapshot.
267    pub fn is_same_transparent_readonly(&self, other: &Rc<TransparentObserverSnapshot>) -> bool {
268        matches!(self, AnySnapshot::TransparentReadonly(snapshot) if Rc::ptr_eq(snapshot, other))
269    }
270
271    /// Enter this snapshot, making it current for the duration of the closure.
272    pub fn enter<T>(&self, f: impl FnOnce() -> T) -> T {
273        match self {
274            AnySnapshot::Readonly(s) => s.enter(f),
275            AnySnapshot::Mutable(s) => s.enter(f),
276            AnySnapshot::NestedReadonly(s) => s.enter(f),
277            AnySnapshot::NestedMutable(s) => s.enter(f),
278            AnySnapshot::Global(s) => s.enter(f),
279            AnySnapshot::TransparentMutable(s) => s.enter(f),
280            AnySnapshot::TransparentReadonly(s) => s.enter(f),
281        }
282    }
283
284    /// Take a nested read-only snapshot.
285    pub fn take_nested_snapshot(&self, read_observer: Option<ReadObserver>) -> AnySnapshot {
286        match self {
287            AnySnapshot::Readonly(s) => {
288                AnySnapshot::Readonly(s.take_nested_snapshot(read_observer))
289            }
290            AnySnapshot::Mutable(s) => AnySnapshot::Readonly(s.take_nested_snapshot(read_observer)),
291            AnySnapshot::NestedReadonly(s) => {
292                AnySnapshot::NestedReadonly(s.take_nested_snapshot(read_observer))
293            }
294            AnySnapshot::NestedMutable(s) => {
295                AnySnapshot::Readonly(s.take_nested_snapshot(read_observer))
296            }
297            AnySnapshot::Global(s) => AnySnapshot::Readonly(s.take_nested_snapshot(read_observer)),
298            AnySnapshot::TransparentMutable(s) => {
299                AnySnapshot::Readonly(s.take_nested_snapshot(read_observer))
300            }
301            AnySnapshot::TransparentReadonly(s) => {
302                AnySnapshot::TransparentReadonly(s.take_nested_snapshot(read_observer))
303            }
304        }
305    }
306
307    /// Check if there are pending changes.
308    pub fn has_pending_changes(&self) -> bool {
309        match self {
310            AnySnapshot::Readonly(s) => s.has_pending_changes(),
311            AnySnapshot::Mutable(s) => s.has_pending_changes(),
312            AnySnapshot::NestedReadonly(s) => s.has_pending_changes(),
313            AnySnapshot::NestedMutable(s) => s.has_pending_changes(),
314            AnySnapshot::Global(s) => s.has_pending_changes(),
315            AnySnapshot::TransparentMutable(s) => s.has_pending_changes(),
316            AnySnapshot::TransparentReadonly(s) => s.has_pending_changes(),
317        }
318    }
319
320    /// Dispose of this snapshot.
321    pub fn dispose(&self) {
322        match self {
323            AnySnapshot::Readonly(s) => s.dispose(),
324            AnySnapshot::Mutable(s) => s.dispose(),
325            AnySnapshot::NestedReadonly(s) => s.dispose(),
326            AnySnapshot::NestedMutable(s) => s.dispose(),
327            AnySnapshot::Global(s) => s.dispose(),
328            AnySnapshot::TransparentMutable(s) => s.dispose(),
329            AnySnapshot::TransparentReadonly(s) => s.dispose(),
330        }
331    }
332
333    /// Check if disposed.
334    pub fn is_disposed(&self) -> bool {
335        match self {
336            AnySnapshot::Readonly(s) => s.is_disposed(),
337            AnySnapshot::Mutable(s) => s.is_disposed(),
338            AnySnapshot::NestedReadonly(s) => s.is_disposed(),
339            AnySnapshot::NestedMutable(s) => s.is_disposed(),
340            AnySnapshot::Global(s) => s.is_disposed(),
341            AnySnapshot::TransparentMutable(s) => s.is_disposed(),
342            AnySnapshot::TransparentReadonly(s) => s.is_disposed(),
343        }
344    }
345
346    /// Record a read.
347    pub fn record_read(&self, state: &dyn StateObject) {
348        match self {
349            AnySnapshot::Readonly(s) => s.record_read(state),
350            AnySnapshot::Mutable(s) => s.record_read(state),
351            AnySnapshot::NestedReadonly(s) => s.record_read(state),
352            AnySnapshot::NestedMutable(s) => s.record_read(state),
353            AnySnapshot::Global(s) => s.record_read(state),
354            AnySnapshot::TransparentMutable(s) => s.record_read(state),
355            AnySnapshot::TransparentReadonly(s) => s.record_read(state),
356        }
357    }
358
359    /// Record a write.
360    pub fn record_write(&self, state: Rc<dyn StateObject>) {
361        match self {
362            AnySnapshot::Readonly(s) => s.record_write(state),
363            AnySnapshot::Mutable(s) => s.record_write(state),
364            AnySnapshot::NestedReadonly(s) => s.record_write(state),
365            AnySnapshot::NestedMutable(s) => s.record_write(state),
366            AnySnapshot::Global(s) => s.record_write(state),
367            AnySnapshot::TransparentMutable(s) => s.record_write(state),
368            AnySnapshot::TransparentReadonly(s) => s.record_write(state),
369        }
370    }
371
372    /// Apply changes (only valid for mutable snapshots).
373    pub fn apply(&self) -> SnapshotApplyResult {
374        match self {
375            AnySnapshot::Mutable(s) => s.apply(),
376            AnySnapshot::NestedMutable(s) => s.apply(),
377            AnySnapshot::Global(s) => s.apply(),
378            AnySnapshot::TransparentMutable(s) => s.apply(),
379            _ => panic!("Cannot apply a read-only snapshot"),
380        }
381    }
382
383    /// Take a nested mutable snapshot (only valid for mutable snapshots).
384    pub fn take_nested_mutable_snapshot(
385        &self,
386        read_observer: Option<ReadObserver>,
387        write_observer: Option<WriteObserver>,
388    ) -> AnySnapshot {
389        match self {
390            AnySnapshot::Mutable(s) => AnySnapshot::NestedMutable(
391                s.take_nested_mutable_snapshot(read_observer, write_observer),
392            ),
393            AnySnapshot::NestedMutable(s) => AnySnapshot::NestedMutable(
394                s.take_nested_mutable_snapshot(read_observer, write_observer),
395            ),
396            AnySnapshot::Global(s) => {
397                AnySnapshot::Mutable(s.take_nested_mutable_snapshot(read_observer, write_observer))
398            }
399            AnySnapshot::TransparentMutable(s) => AnySnapshot::TransparentMutable(
400                s.take_nested_mutable_snapshot(read_observer, write_observer),
401            ),
402            _ => panic!("Cannot take nested mutable snapshot from read-only snapshot"),
403        }
404    }
405}
406
407/// Get the current snapshot, or None if not in a snapshot context.
408pub fn current_snapshot() -> Option<AnySnapshot> {
409    crate::thread_state::try_snapshots(|state| state.current.borrow().clone()).flatten()
410}
411
412pub(crate) fn current_snapshot_reads_into(observer: &ReadObserver) -> bool {
413    crate::thread_state::try_snapshots(|state| match &*state.current.borrow() {
414        Some(AnySnapshot::TransparentMutable(snapshot)) => {
415            snapshot.can_reuse() && snapshot.reads_into(observer)
416        }
417        _ => false,
418    })
419    .unwrap_or(false)
420}
421
422pub(crate) fn set_current_snapshot(snapshot: Option<AnySnapshot>) {
423    crate::thread_state::try_snapshots(|state| {
424        *state.current.borrow_mut() = snapshot;
425    });
426}
427
428struct CurrentSnapshotGuard {
429    previous: Option<AnySnapshot>,
430}
431
432impl CurrentSnapshotGuard {
433    fn enter(snapshot: AnySnapshot) -> Self {
434        let previous =
435            crate::thread_state::try_snapshots(|state| state.current.replace(Some(snapshot)))
436                .flatten();
437        Self { previous }
438    }
439}
440
441impl Drop for CurrentSnapshotGuard {
442    fn drop(&mut self) {
443        set_current_snapshot(self.previous.take());
444    }
445}
446
447pub(crate) fn enter_snapshot_scope<T>(snapshot: AnySnapshot, f: impl FnOnce() -> T) -> T {
448    let _guard = CurrentSnapshotGuard::enter(snapshot);
449    f()
450}
451
452/// Creates a mutable snapshot, matching Kotlin's `Snapshot.takeMutableSnapshot` semantics.
453///
454/// If called while inside a MutableSnapshot, creates a nested snapshot that will
455/// apply to the parent when `apply()` is called. This ensures proper isolation
456/// between nested operations (like event handlers during animations).
457///
458/// If called while inside a GlobalSnapshot or no snapshot, creates a root
459/// mutable snapshot that applies to the global state.
460pub fn take_mutable_snapshot(
461    read_observer: Option<ReadObserver>,
462    write_observer: Option<WriteObserver>,
463) -> AnyMutableSnapshot {
464    match current_snapshot() {
465        Some(AnySnapshot::Mutable(parent)) => AnyMutableSnapshot::Nested(
466            parent.take_nested_mutable_snapshot(read_observer, write_observer),
467        ),
468        Some(AnySnapshot::NestedMutable(parent)) => AnyMutableSnapshot::Nested(
469            parent.take_nested_mutable_snapshot(read_observer, write_observer),
470        ),
471        _ => AnyMutableSnapshot::Root(
472            GlobalSnapshot::get_or_create()
473                .take_nested_mutable_snapshot(read_observer, write_observer),
474        ),
475    }
476}
477
478/// Take a transparent observer mutable snapshot with optional observers.
479///
480/// This type of snapshot is used for read observation during composition,
481/// matching Kotlin's Snapshot.observeInternal behavior. It allows writes
482/// to happen during observation.
483///
484/// Transparent snapshots DO NOT allocate new IDs - they delegate to the
485/// current/global snapshot, making them "transparent" to the snapshot system.
486pub fn take_transparent_observer_mutable_snapshot(
487    read_observer: Option<ReadObserver>,
488    write_observer: Option<WriteObserver>,
489) -> Rc<TransparentObserverMutableSnapshot> {
490    take_transparent_observer_mutable_snapshot_reusing(read_observer, write_observer, None)
491}
492
493pub(crate) fn take_transparent_observer_mutable_snapshot_reusing(
494    read_observer: Option<ReadObserver>,
495    write_observer: Option<WriteObserver>,
496    recycled: Option<Rc<TransparentObserverMutableSnapshot>>,
497) -> Rc<TransparentObserverMutableSnapshot> {
498    let parent = current_snapshot();
499    match parent {
500        Some(AnySnapshot::TransparentMutable(transparent)) if transparent.can_reuse() => {
501            let (parent_read, parent_write) = transparent.observers();
502            if already_observes(&read_observer, &parent_read)
503                && already_observes(&write_observer, &parent_write)
504            {
505                return transparent;
506            }
507            TransparentObserverMutableSnapshot::new_reusing(
508                recycled,
509                transparent.snapshot_id(),
510                transparent.invalid(),
511                merge_read_observers(read_observer, parent_read),
512                merge_write_observers(write_observer, parent_write),
513                Some(Rc::downgrade(&transparent)),
514            )
515        }
516        _ => {
517            let current = current_snapshot()
518                .unwrap_or_else(|| AnySnapshot::Global(GlobalSnapshot::get_or_create()));
519            let id = current.snapshot_id();
520            let invalid = current.invalid();
521            TransparentObserverMutableSnapshot::new_reusing(
522                recycled,
523                id,
524                invalid,
525                read_observer,
526                write_observer,
527                None,
528            )
529        }
530    }
531}
532
533fn already_observes(requested: &Option<ReadObserver>, installed: &Option<ReadObserver>) -> bool {
534    match (requested, installed) {
535        (None, _) => true,
536        (Some(requested), Some(installed)) => Rc::ptr_eq(requested, installed),
537        (Some(_), None) => false,
538    }
539}
540
541/// Allocate a new record identifier that is distinct from any active snapshot id.
542pub fn allocate_record_id() -> SnapshotId {
543    runtime::allocate_record_id()
544}
545
546pub(crate) fn peek_next_snapshot_id() -> SnapshotId {
547    runtime::peek_next_snapshot_id()
548}
549
550#[derive(Clone)]
551pub(crate) struct ObserverId(Rc<()>);
552
553impl ObserverId {
554    fn new() -> Self {
555        Self(Rc::new(()))
556    }
557}
558
559impl PartialEq for ObserverId {
560    fn eq(&self, other: &Self) -> bool {
561        Rc::ptr_eq(&self.0, &other.0)
562    }
563}
564
565impl Eq for ObserverId {}
566
567/// The apply observers of a thread in the order they registered, and the
568/// list a notification runs, kept until an observer registers or leaves.
569pub(crate) struct ApplyObservers {
570    entries: Vec<(ObserverId, ApplyObserver)>,
571    running: Option<Rc<[ApplyObserver]>>,
572}
573
574impl ApplyObservers {
575    pub(crate) const fn new() -> Self {
576        Self {
577            entries: Vec::new(),
578            running: None,
579        }
580    }
581
582    fn insert(&mut self, id: ObserverId, observer: ApplyObserver) {
583        self.entries.push((id, observer));
584        self.running = None;
585    }
586
587    fn remove(&mut self, id: &ObserverId) {
588        if let Some(at) = self.entries.iter().position(|(entry, _)| entry == id) {
589            self.entries.remove(at);
590            self.running = None;
591        }
592    }
593
594    /// The observers a notification runs, apart from the registry: an
595    /// observer may register or drop observers while it runs.
596    fn running(&mut self) -> Rc<[ApplyObserver]> {
597        let entries = &self.entries;
598        Rc::clone(self.running.get_or_insert_with(|| {
599            entries
600                .iter()
601                .map(|(_, observer)| Rc::clone(observer))
602                .collect()
603        }))
604    }
605}
606
607const UNUSED_RECORD_CLEANUP_INTERVAL: SnapshotId = 2;
608const UNUSED_RECORD_CLEANUP_BUSY_INTERVAL: SnapshotId = 1;
609const UNUSED_RECORD_CLEANUP_MIN_SIZE: usize = 64;
610
611/// Size figures of the snapshot system's bookkeeping on the current thread, from
612/// `debug_snapshot_v2_stats`. `*_len` is the number of entries and `*_cap` the entries that fit
613/// before the storage grows.
614#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
615pub struct SnapshotV2DebugStats {
616    /// Registered apply observers.
617    pub apply_observers_len: usize,
618    /// Apply observers the observer storage holds before it grows.
619    pub apply_observers_cap: usize,
620    /// State objects with a recorded last-write snapshot id.
621    pub last_writes_len: usize,
622    /// Entries the last-write record holds before it grows.
623    pub last_writes_cap: usize,
624    /// State objects tracked for later cleanup of unused records.
625    pub extra_state_objects_len: usize,
626    /// Entries the tracked state object set holds before it grows.
627    pub extra_state_objects_cap: usize,
628    /// Id of the snapshot at which unused records were last cleaned up.
629    pub last_unused_record_cleanup: SnapshotId,
630}
631
632/// Returns the size figures of the snapshot system's bookkeeping on the current thread.
633pub fn debug_snapshot_v2_stats() -> SnapshotV2DebugStats {
634    crate::thread_state::snapshots(|state| {
635        let observers = state.apply_observers.borrow();
636        let writes = state.last_writes.borrow();
637        let SnapshotWeakSetDebugStats {
638            len: extra_state_objects_len,
639            capacity: extra_state_objects_cap,
640        } = state.extra_state_objects.borrow().debug_stats();
641        SnapshotV2DebugStats {
642            apply_observers_len: observers.entries.len(),
643            apply_observers_cap: observers.entries.capacity(),
644            last_writes_len: writes.len(),
645            last_writes_cap: writes.capacity(),
646            extra_state_objects_len,
647            extra_state_objects_cap,
648            last_unused_record_cleanup: state.last_unused_record_cleanup.get(),
649        }
650    })
651}
652
653/// Register an apply observer.
654///
655/// Returns a handle that will automatically unregister the observer when dropped.
656pub fn register_apply_observer(observer: ApplyObserver) -> ObserverHandle {
657    crate::thread_state::watch_thread_exit();
658    let id = ObserverId::new();
659    crate::thread_state::snapshots(|state| {
660        state
661            .apply_observers
662            .borrow_mut()
663            .insert(id.clone(), observer);
664    });
665    ObserverHandle {
666        kind: ObserverKind::Apply,
667        id,
668    }
669}
670
671/// Handle for unregistering observers.
672///
673/// When dropped, automatically removes the associated observer.
674pub struct ObserverHandle {
675    kind: ObserverKind,
676    id: ObserverId,
677}
678
679enum ObserverKind {
680    Apply,
681}
682
683impl Drop for ObserverHandle {
684    fn drop(&mut self) {
685        match self.kind {
686            ObserverKind::Apply => {
687                // A handle dropped while the thread tears its locals down,
688                // as a coroutine's state can be, finds the registry gone
689                // and has nothing left to remove.
690                crate::thread_state::try_snapshots(|state| {
691                    state.apply_observers.borrow_mut().remove(&self.id);
692                });
693            }
694        }
695    }
696}
697
698pub(crate) fn notify_apply_observers(modified: &[Rc<dyn StateObject>], snapshot_id: SnapshotId) {
699    let observers =
700        crate::thread_state::snapshots(|state| state.apply_observers.borrow_mut().running());
701    for observer in observers.iter() {
702        observer(modified, snapshot_id);
703    }
704}
705
706pub(crate) fn set_last_write(id: StateObjectId, snapshot_id: SnapshotId) {
707    crate::thread_state::snapshots(|state| {
708        state.last_writes.borrow_mut().insert(id, snapshot_id);
709    });
710}
711
712#[cfg(test)]
713pub(crate) fn clear_last_writes() {
714    crate::thread_state::snapshots(|state| state.last_writes.borrow_mut().clear());
715}
716
717pub(crate) fn check_and_overwrite_unused_records_locked() {
718    let _retiring = super::state::RetiredValuesScope::enter();
719    crate::thread_state::snapshots(|state| {
720        state
721            .extra_state_objects
722            .borrow_mut()
723            .remove_if(super::state::StateObject::overwrite_unused_records);
724    });
725}
726
727pub(crate) fn maybe_check_and_overwrite_unused_records_locked(current_snapshot_id: SnapshotId) {
728    let should_run = crate::thread_state::snapshots(|state| {
729        let set = state.extra_state_objects.borrow();
730        if set.is_empty() {
731            return false;
732        }
733        let interval = if set.len() >= UNUSED_RECORD_CLEANUP_MIN_SIZE {
734            UNUSED_RECORD_CLEANUP_BUSY_INTERVAL
735        } else {
736            UNUSED_RECORD_CLEANUP_INTERVAL
737        };
738        let due =
739            current_snapshot_id.saturating_sub(state.last_unused_record_cleanup.get()) >= interval;
740        if due {
741            state.last_unused_record_cleanup.set(current_snapshot_id);
742        }
743        due
744    });
745
746    if should_run {
747        check_and_overwrite_unused_records_locked();
748    }
749}
750
751#[cfg(test)]
752pub(crate) fn clear_unused_record_cleanup_for_tests() {
753    crate::thread_state::snapshots(|state| state.last_unused_record_cleanup.set(0));
754}
755
756pub(crate) fn optimistic_merges(
757    current_snapshot_id: SnapshotId,
758    base_parent_id: SnapshotId,
759    modified_objects: &[(StateObjectId, Rc<dyn StateObject>, SnapshotId)],
760    invalid_snapshots: &SnapshotIdSet,
761    applying_invalid: &SnapshotIdSet,
762) -> Option<HashMap<usize, Rc<StateRecord>>> {
763    if modified_objects.is_empty() {
764        return None;
765    }
766
767    let mut result: Option<HashMap<usize, Rc<StateRecord>>> = None;
768
769    for (_, state, writer_id) in modified_objects {
770        let head = state.first_record();
771
772        let Some(current) =
773            crate::state::readable_record_for(&head, current_snapshot_id, invalid_snapshots)
774        else {
775            continue;
776        };
777
778        let (previous_opt, found_base) =
779            mutable::find_previous_record(&head, base_parent_id, applying_invalid);
780        let previous = previous_opt?;
781
782        if !found_base || previous.snapshot_id() == crate::state::PREEXISTING_SNAPSHOT_ID {
783            continue;
784        }
785
786        if Rc::ptr_eq(&current, &previous) {
787            continue;
788        }
789
790        let applied = mutable::find_record_by_id(&head, *writer_id)?;
791
792        let merged = state.merge_records(
793            Rc::clone(&previous),
794            Rc::clone(&current),
795            Rc::clone(&applied),
796        )?;
797
798        result
799            .get_or_insert_with(HashMap::default)
800            .insert(Rc::as_ptr(&current) as usize, merged);
801    }
802
803    result
804}
805
806fn merge_observers(a: Option<ReadObserver>, b: Option<ReadObserver>) -> Option<ReadObserver> {
807    match (a, b) {
808        (None, None) => None,
809        (Some(a), None) => Some(a),
810        (None, Some(b)) => Some(b),
811        (Some(a), Some(b)) => Some(Rc::new(move |state: &dyn StateObject| {
812            a(state);
813            b(state);
814        })),
815    }
816}
817
818/// Merge two read observers into one.
819///
820/// # Thread Safety
821/// The resulting Rc-wrapped closure may capture non-Send closures. This is safe
822/// because observers are only invoked on the UI thread where they were created.
823pub fn merge_read_observers(
824    a: Option<ReadObserver>,
825    b: Option<ReadObserver>,
826) -> Option<ReadObserver> {
827    merge_observers(a, b)
828}
829
830/// Merge two write observers into one.
831///
832/// # Thread Safety
833/// The resulting Rc-wrapped closure may capture non-Send closures. This is safe
834/// because observers are only invoked on the UI thread where they were created.
835pub fn merge_write_observers(
836    a: Option<WriteObserver>,
837    b: Option<WriteObserver>,
838) -> Option<WriteObserver> {
839    merge_observers(a, b)
840}
841
842pub(crate) struct SnapshotState {
843    pub(crate) id: Cell<SnapshotId>,
844    pub(crate) invalid: RefCell<SnapshotIdSet>,
845    pub(crate) pin_handle: Cell<PinHandle>,
846    pub(crate) disposed: Cell<bool>,
847    pub(crate) read_observer: RefCell<Option<ReadObserver>>,
848    pub(crate) write_observer: RefCell<Option<WriteObserver>>,
849    #[expect(clippy::type_complexity)]
850    pub(crate) modified: RefCell<HashMap<StateObjectId, (Rc<dyn StateObject>, SnapshotId)>>,
851    on_dispose: RefCell<Option<Box<dyn FnOnce()>>>,
852    runtime_tracked: bool,
853    pending_children: RefCell<HashSet<SnapshotId>>,
854}
855
856impl SnapshotState {
857    pub(crate) fn new(
858        id: SnapshotId,
859        invalid: SnapshotIdSet,
860        read_observer: Option<ReadObserver>,
861        write_observer: Option<WriteObserver>,
862        runtime_tracked: bool,
863    ) -> Self {
864        Self::new_with_pinning(
865            id,
866            invalid,
867            read_observer,
868            write_observer,
869            runtime_tracked,
870            true,
871        )
872    }
873
874    pub(crate) fn new_with_pinning(
875        id: SnapshotId,
876        invalid: SnapshotIdSet,
877        read_observer: Option<ReadObserver>,
878        write_observer: Option<WriteObserver>,
879        runtime_tracked: bool,
880        should_pin: bool,
881    ) -> Self {
882        let pin_handle = if should_pin {
883            snapshot_pinning::track_pinning(id, &invalid)
884        } else {
885            snapshot_pinning::PinHandle::INVALID
886        };
887        Self {
888            id: Cell::new(id),
889            invalid: RefCell::new(invalid),
890            pin_handle: Cell::new(pin_handle),
891            disposed: Cell::new(false),
892            read_observer: RefCell::new(read_observer),
893            write_observer: RefCell::new(write_observer),
894            modified: RefCell::new(HashMap::default()),
895            on_dispose: RefCell::new(None),
896            runtime_tracked,
897            pending_children: RefCell::new(HashSet::default()),
898        }
899    }
900
901    pub(crate) fn record_read(&self, state: &dyn StateObject) {
902        if let Some(observer) = self.read_observer.borrow().as_ref() {
903            observer(state);
904        }
905    }
906
907    pub(crate) fn record_write(&self, state: Rc<dyn StateObject>, writer_id: SnapshotId) {
908        let state_id = state.object_id().as_usize();
909
910        let mut modified = self.modified.borrow_mut();
911
912        match modified.entry(state_id) {
913            std::collections::hash_map::Entry::Vacant(e) => {
914                if let Some(observer) = self.write_observer.borrow().as_ref() {
915                    observer(&*state);
916                }
917                e.insert((state, writer_id));
918            }
919            std::collections::hash_map::Entry::Occupied(mut e) => {
920                e.insert((state, writer_id));
921            }
922        }
923    }
924
925    pub(crate) fn dispose(&self) {
926        if !self.disposed.replace(true) {
927            let pin_handle = self.pin_handle.get();
928            snapshot_pinning::release_pinning(pin_handle);
929            if let Some(cb) = self.on_dispose.borrow_mut().take() {
930                cb();
931            }
932            if self.runtime_tracked {
933                close_snapshot(self.id.get());
934            }
935        }
936    }
937
938    pub(crate) fn add_pending_child(&self, id: SnapshotId) {
939        self.pending_children.borrow_mut().insert(id);
940    }
941
942    pub(crate) fn remove_pending_child(&self, id: SnapshotId) {
943        self.pending_children.borrow_mut().remove(&id);
944    }
945
946    pub(crate) fn has_pending_children(&self) -> bool {
947        !self.pending_children.borrow().is_empty()
948    }
949
950    pub(crate) fn pending_children(&self) -> Vec<SnapshotId> {
951        self.pending_children.borrow().iter().copied().collect()
952    }
953
954    pub(crate) fn set_on_dispose<F>(&self, f: F)
955    where
956        F: FnOnce() + 'static,
957    {
958        *self.on_dispose.borrow_mut() = Some(Box::new(f));
959    }
960}
961
962pub(crate) trait NestedMutableHost {
963    fn snapshot_state(&self) -> &SnapshotState;
964    fn nested_count(&self) -> &Cell<usize>;
965}
966
967pub(crate) fn clear_nested_child_on_dispose<P>(
968    parent: &Rc<P>,
969    child_id: SnapshotId,
970) -> impl FnOnce() + 'static
971where
972    P: NestedMutableHost + 'static,
973{
974    let weak = Rc::downgrade(parent);
975    move || {
976        if let Some(parent) = weak.upgrade() {
977            let nested_count = parent.nested_count();
978            if nested_count.get() > 0 {
979                nested_count.set(nested_count.get().saturating_sub(1));
980            }
981            let state = parent.snapshot_state();
982            let new_invalid = state.invalid.borrow().clone().clear(child_id);
983            state.invalid.replace(new_invalid);
984            state.remove_pending_child(child_id);
985        }
986    }
987}
988
989pub(crate) fn allocate_nested_mutable_snapshot<P>(
990    parent: &Rc<P>,
991    root: Weak<MutableSnapshot>,
992    read_observer: Option<ReadObserver>,
993    write_observer: Option<WriteObserver>,
994) -> Rc<NestedMutableSnapshot>
995where
996    P: NestedMutableHost + 'static,
997{
998    let state = parent.snapshot_state();
999    let merged_read = merge_read_observers(read_observer, state.read_observer.borrow().clone());
1000    let merged_write = merge_write_observers(write_observer, state.write_observer.borrow().clone());
1001
1002    let parent_id = state.id.get();
1003    let current_invalid = state.invalid.borrow().clone();
1004
1005    let (new_id, _runtime_invalid) = allocate_snapshot();
1006
1007    let parent_invalid_with_child = current_invalid.set(new_id);
1008    state.invalid.replace(parent_invalid_with_child);
1009
1010    let invalid = current_invalid.add_range(parent_id + 1, new_id);
1011
1012    let nested = NestedMutableSnapshot::new(
1013        new_id,
1014        invalid,
1015        merged_read,
1016        merged_write,
1017        root,
1018        state.id.get(),
1019    );
1020
1021    let nested_count = parent.nested_count();
1022    nested_count.set(nested_count.get() + 1);
1023    state.add_pending_child(new_id);
1024
1025    nested.set_on_dispose(clear_nested_child_on_dispose(parent, new_id));
1026
1027    nested
1028}
1029
1030#[cfg(test)]
1031#[path = "tests/snapshot_v2_tests.rs"]
1032mod tests;
cranpose_core/snapshot_v2/mod.rs

cranpose_core/snapshot_v2/
mod.rs
