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    let _retiring = super::state::RetiredValuesScope::enter();
693    EXTRA_STATE_OBJECTS.with(|cell| {
694        cell.borrow_mut()
695            .remove_if(super::state::StateObject::overwrite_unused_records);
696    });
697}
698
699pub(crate) fn maybe_check_and_overwrite_unused_records_locked(current_snapshot_id: SnapshotId) {
700    let should_run = EXTRA_STATE_OBJECTS.with(|cell| {
701        let set = cell.borrow();
702        if set.is_empty() {
703            return false;
704        }
705        let last_cleanup = LAST_UNUSED_RECORD_CLEANUP.with(Cell::get);
706        let interval = if set.len() >= UNUSED_RECORD_CLEANUP_MIN_SIZE {
707            UNUSED_RECORD_CLEANUP_BUSY_INTERVAL
708        } else {
709            UNUSED_RECORD_CLEANUP_INTERVAL
710        };
711        current_snapshot_id.saturating_sub(last_cleanup) >= interval
712    });
713
714    if should_run {
715        LAST_UNUSED_RECORD_CLEANUP.with(|cell| cell.set(current_snapshot_id));
716        check_and_overwrite_unused_records_locked();
717    }
718}
719
720#[cfg(test)]
721pub(crate) fn clear_unused_record_cleanup_for_tests() {
722    LAST_UNUSED_RECORD_CLEANUP.with(|cell| cell.set(0));
723}
724
725pub(crate) fn optimistic_merges(
726    current_snapshot_id: SnapshotId,
727    base_parent_id: SnapshotId,
728    modified_objects: &[(StateObjectId, Arc<dyn StateObject>, SnapshotId)],
729    invalid_snapshots: &SnapshotIdSet,
730    applying_invalid: &SnapshotIdSet,
731) -> Option<HashMap<usize, Rc<StateRecord>>> {
732    if modified_objects.is_empty() {
733        return None;
734    }
735
736    let mut result: Option<HashMap<usize, Rc<StateRecord>>> = None;
737
738    for (_, state, writer_id) in modified_objects {
739        let head = state.first_record();
740
741        let Some(current) =
742            crate::state::readable_record_for(&head, current_snapshot_id, invalid_snapshots)
743        else {
744            continue;
745        };
746
747        let (previous_opt, found_base) =
748            mutable::find_previous_record(&head, base_parent_id, applying_invalid);
749        let previous = previous_opt?;
750
751        if !found_base || previous.snapshot_id() == crate::state::PREEXISTING_SNAPSHOT_ID {
752            continue;
753        }
754
755        if Rc::ptr_eq(&current, &previous) {
756            continue;
757        }
758
759        let applied = mutable::find_record_by_id(&head, *writer_id)?;
760
761        let merged = state.merge_records(
762            Rc::clone(&previous),
763            Rc::clone(&current),
764            Rc::clone(&applied),
765        )?;
766
767        result
768            .get_or_insert_with(HashMap::default)
769            .insert(Rc::as_ptr(&current) as usize, merged);
770    }
771
772    result
773}
774
775#[expect(clippy::arc_with_non_send_sync)]
776fn merge_observers(a: Option<ReadObserver>, b: Option<ReadObserver>) -> Option<ReadObserver> {
777    match (a, b) {
778        (None, None) => None,
779        (Some(a), None) => Some(a),
780        (None, Some(b)) => Some(b),
781        (Some(a), Some(b)) => Some(Arc::new(move |state: &dyn StateObject| {
782            a(state);
783            b(state);
784        })),
785    }
786}
787
788/// Merge two read observers into one.
789///
790/// # Thread Safety
791/// The resulting Arc-wrapped closure may capture non-Send closures. This is safe
792/// because observers are only invoked on the UI thread where they were created.
793pub fn merge_read_observers(
794    a: Option<ReadObserver>,
795    b: Option<ReadObserver>,
796) -> Option<ReadObserver> {
797    merge_observers(a, b)
798}
799
800/// Merge two write observers into one.
801///
802/// # Thread Safety
803/// The resulting Arc-wrapped closure may capture non-Send closures. This is safe
804/// because observers are only invoked on the UI thread where they were created.
805pub fn merge_write_observers(
806    a: Option<WriteObserver>,
807    b: Option<WriteObserver>,
808) -> Option<WriteObserver> {
809    merge_observers(a, b)
810}
811
812pub(crate) struct SnapshotState {
813    pub(crate) id: Cell<SnapshotId>,
814    pub(crate) invalid: RefCell<SnapshotIdSet>,
815    pub(crate) pin_handle: Cell<PinHandle>,
816    pub(crate) disposed: Cell<bool>,
817    pub(crate) read_observer: RefCell<Option<ReadObserver>>,
818    pub(crate) write_observer: RefCell<Option<WriteObserver>>,
819    #[expect(clippy::type_complexity)]
820    pub(crate) modified: RefCell<HashMap<StateObjectId, (Arc<dyn StateObject>, SnapshotId)>>,
821    on_dispose: RefCell<Option<Box<dyn FnOnce()>>>,
822    runtime_tracked: bool,
823    pending_children: RefCell<HashSet<SnapshotId>>,
824}
825
826impl SnapshotState {
827    pub(crate) fn new(
828        id: SnapshotId,
829        invalid: SnapshotIdSet,
830        read_observer: Option<ReadObserver>,
831        write_observer: Option<WriteObserver>,
832        runtime_tracked: bool,
833    ) -> Self {
834        Self::new_with_pinning(
835            id,
836            invalid,
837            read_observer,
838            write_observer,
839            runtime_tracked,
840            true,
841        )
842    }
843
844    pub(crate) fn new_with_pinning(
845        id: SnapshotId,
846        invalid: SnapshotIdSet,
847        read_observer: Option<ReadObserver>,
848        write_observer: Option<WriteObserver>,
849        runtime_tracked: bool,
850        should_pin: bool,
851    ) -> Self {
852        let pin_handle = if should_pin {
853            snapshot_pinning::track_pinning(id, &invalid)
854        } else {
855            snapshot_pinning::PinHandle::INVALID
856        };
857        Self {
858            id: Cell::new(id),
859            invalid: RefCell::new(invalid),
860            pin_handle: Cell::new(pin_handle),
861            disposed: Cell::new(false),
862            read_observer: RefCell::new(read_observer),
863            write_observer: RefCell::new(write_observer),
864            modified: RefCell::new(HashMap::default()),
865            on_dispose: RefCell::new(None),
866            runtime_tracked,
867            pending_children: RefCell::new(HashSet::default()),
868        }
869    }
870
871    pub(crate) fn record_read(&self, state: &dyn StateObject) {
872        if let Some(observer) = self.read_observer.borrow().as_ref() {
873            observer(state);
874        }
875    }
876
877    pub(crate) fn record_write(&self, state: Arc<dyn StateObject>, writer_id: SnapshotId) {
878        let state_id = state.object_id().as_usize();
879
880        let mut modified = self.modified.borrow_mut();
881
882        match modified.entry(state_id) {
883            std::collections::hash_map::Entry::Vacant(e) => {
884                if let Some(observer) = self.write_observer.borrow().as_ref() {
885                    observer(&*state);
886                }
887                e.insert((state, writer_id));
888            }
889            std::collections::hash_map::Entry::Occupied(mut e) => {
890                e.insert((state, writer_id));
891            }
892        }
893    }
894
895    pub(crate) fn dispose(&self) {
896        if !self.disposed.replace(true) {
897            let pin_handle = self.pin_handle.get();
898            snapshot_pinning::release_pinning(pin_handle);
899            if let Some(cb) = self.on_dispose.borrow_mut().take() {
900                cb();
901            }
902            if self.runtime_tracked {
903                close_snapshot(self.id.get());
904            }
905        }
906    }
907
908    pub(crate) fn add_pending_child(&self, id: SnapshotId) {
909        self.pending_children.borrow_mut().insert(id);
910    }
911
912    pub(crate) fn remove_pending_child(&self, id: SnapshotId) {
913        self.pending_children.borrow_mut().remove(&id);
914    }
915
916    pub(crate) fn has_pending_children(&self) -> bool {
917        !self.pending_children.borrow().is_empty()
918    }
919
920    pub(crate) fn pending_children(&self) -> Vec<SnapshotId> {
921        self.pending_children.borrow().iter().copied().collect()
922    }
923
924    pub(crate) fn set_on_dispose<F>(&self, f: F)
925    where
926        F: FnOnce() + 'static,
927    {
928        *self.on_dispose.borrow_mut() = Some(Box::new(f));
929    }
930}
931
932pub(crate) trait NestedMutableHost {
933    fn snapshot_state(&self) -> &SnapshotState;
934    fn nested_count(&self) -> &Cell<usize>;
935}
936
937pub(crate) fn clear_nested_child_on_dispose<P>(
938    parent: &Arc<P>,
939    child_id: SnapshotId,
940) -> impl FnOnce() + 'static
941where
942    P: NestedMutableHost + 'static,
943{
944    let weak = Arc::downgrade(parent);
945    move || {
946        if let Some(parent) = weak.upgrade() {
947            let nested_count = parent.nested_count();
948            if nested_count.get() > 0 {
949                nested_count.set(nested_count.get().saturating_sub(1));
950            }
951            let state = parent.snapshot_state();
952            let new_invalid = state.invalid.borrow().clone().clear(child_id);
953            state.invalid.replace(new_invalid);
954            state.remove_pending_child(child_id);
955        }
956    }
957}
958
959pub(crate) fn allocate_nested_mutable_snapshot<P>(
960    parent: &Arc<P>,
961    root: Weak<MutableSnapshot>,
962    read_observer: Option<ReadObserver>,
963    write_observer: Option<WriteObserver>,
964) -> Arc<NestedMutableSnapshot>
965where
966    P: NestedMutableHost + 'static,
967{
968    let state = parent.snapshot_state();
969    let merged_read = merge_read_observers(read_observer, state.read_observer.borrow().clone());
970    let merged_write = merge_write_observers(write_observer, state.write_observer.borrow().clone());
971
972    let parent_id = state.id.get();
973    let current_invalid = state.invalid.borrow().clone();
974
975    let (new_id, _runtime_invalid) = allocate_snapshot();
976
977    let parent_invalid_with_child = current_invalid.set(new_id);
978    state.invalid.replace(parent_invalid_with_child);
979
980    let invalid = current_invalid.add_range(parent_id + 1, new_id);
981
982    let nested = NestedMutableSnapshot::new(
983        new_id,
984        invalid,
985        merged_read,
986        merged_write,
987        root,
988        state.id.get(),
989    );
990
991    let nested_count = parent.nested_count();
992    nested_count.set(nested_count.get() + 1);
993    state.add_pending_child(new_id);
994
995    nested.set_on_dispose(clear_nested_child_on_dispose(parent, new_id));
996
997    nested
998}
999
1000#[cfg(test)]
1001#[path = "tests/snapshot_v2_tests.rs"]
1002mod tests;
cranpose_core/snapshot_v2/mod.rs

cranpose_core/snapshot_v2/
mod.rs
