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

cranpose_core/snapshot_v2/
mod.rs
