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    hash::{Hash, Hasher},
26    rc::{Rc, Weak},
27};
28
29use crate::{
30    collections::map::{HashMap, HashSet},
31    snapshot_id_set::{SnapshotId, SnapshotIdSet},
32    snapshot_pinning::{self, PinHandle},
33    snapshot_weak_set::SnapshotWeakSetDebugStats,
34    state::{StateObject, StateRecord},
35};
36
37mod global;
38mod mutable;
39mod nested;
40mod readonly;
41mod runtime;
42mod transparent;
43
44#[cfg(test)]
45#[path = "tests/integration_tests.rs"]
46mod integration_tests;
47
48pub use global::{GlobalSnapshot, advance_global_snapshot};
49pub use mutable::MutableSnapshot;
50pub use nested::{NestedMutableSnapshot, NestedReadonlySnapshot};
51pub use readonly::ReadonlySnapshot;
52#[cfg(test)]
53pub(crate) use runtime::{TestRuntimeGuard, reset_runtime_for_tests};
54pub(crate) use runtime::{allocate_snapshot, close_snapshot, with_runtime};
55pub use transparent::{TransparentObserverMutableSnapshot, TransparentObserverSnapshot};
56
57/// Observer that is called when a state object is read.
58pub type ReadObserver = Rc<dyn Fn(&dyn StateObject) + 'static>;
59
60/// Observer that is called when a state object is written.
61pub type WriteObserver = Rc<dyn Fn(&dyn StateObject) + 'static>;
62
63/// Apply observer that is called when a snapshot is applied.
64pub type ApplyObserver = Rc<dyn Fn(&[Rc<dyn StateObject>], SnapshotId) + 'static>;
65
66/// Result of applying a mutable snapshot.
67#[derive(Debug, Clone, PartialEq, Eq)]
68pub enum SnapshotApplyResult {
69    /// The snapshot was applied successfully.
70    Success,
71    /// The snapshot could not be applied due to conflicts.
72    Failure,
73}
74
75impl SnapshotApplyResult {
76    /// Check if the result is successful.
77    pub fn is_success(&self) -> bool {
78        matches!(self, SnapshotApplyResult::Success)
79    }
80
81    /// Check if the result is a failure.
82    pub fn is_failure(&self) -> bool {
83        matches!(self, SnapshotApplyResult::Failure)
84    }
85
86    /// Panic if the result is a failure (for use in tests).
87    #[track_caller]
88    pub fn check(&self) {
89        assert!(!self.is_failure(), "Snapshot apply failed");
90    }
91}
92
93/// Unique identifier for a state object in the modified set.
94pub type StateObjectId = usize;
95
96/// Enum wrapper for all snapshot types.
97///
98/// This provides a type-safe way to work with different snapshot types
99/// without requiring trait objects, which avoids object-safety issues.
100#[derive(Clone)]
101pub enum AnySnapshot {
102    Readonly(Rc<ReadonlySnapshot>),
103    Mutable(Rc<MutableSnapshot>),
104    NestedReadonly(Rc<NestedReadonlySnapshot>),
105    NestedMutable(Rc<NestedMutableSnapshot>),
106    Global(Rc<GlobalSnapshot>),
107    TransparentMutable(Rc<TransparentObserverMutableSnapshot>),
108    TransparentReadonly(Rc<TransparentObserverSnapshot>),
109}
110
111/// Enum wrapper for mutable snapshot types.
112///
113/// This allows `take_mutable_snapshot` to return either a root MutableSnapshot
114/// or a NestedMutableSnapshot depending on the current context, matching Kotlin's
115/// behavior where `takeMutableSnapshot` creates nested snapshots when inside a
116/// mutable snapshot.
117#[derive(Clone)]
118pub enum AnyMutableSnapshot {
119    Root(Rc<MutableSnapshot>),
120    Nested(Rc<NestedMutableSnapshot>),
121}
122
123impl AnyMutableSnapshot {
124    /// Get the snapshot ID.
125    pub fn snapshot_id(&self) -> SnapshotId {
126        match self {
127            AnyMutableSnapshot::Root(s) => s.snapshot_id(),
128            AnyMutableSnapshot::Nested(s) => s.snapshot_id(),
129        }
130    }
131
132    /// Get the set of invalid snapshot IDs.
133    pub fn invalid(&self) -> SnapshotIdSet {
134        match self {
135            AnyMutableSnapshot::Root(s) => s.invalid(),
136            AnyMutableSnapshot::Nested(s) => s.invalid(),
137        }
138    }
139
140    /// Enter this snapshot, making it current for the duration of the closure.
141    pub fn enter<T>(&self, f: impl FnOnce() -> T) -> T {
142        match self {
143            AnyMutableSnapshot::Root(s) => s.enter(f),
144            AnyMutableSnapshot::Nested(s) => s.enter(f),
145        }
146    }
147
148    /// Apply the snapshot.
149    pub fn apply(&self) -> SnapshotApplyResult {
150        match self {
151            AnyMutableSnapshot::Root(s) => s.apply(),
152            AnyMutableSnapshot::Nested(s) => s.apply(),
153        }
154    }
155
156    /// Dispose the snapshot.
157    pub fn dispose(&self) {
158        match self {
159            AnyMutableSnapshot::Root(s) => s.dispose(),
160            AnyMutableSnapshot::Nested(s) => s.dispose(),
161        }
162    }
163}
164
165impl AnySnapshot {
166    /// Get the snapshot ID.
167    pub fn snapshot_id(&self) -> SnapshotId {
168        match self {
169            AnySnapshot::Readonly(s) => s.snapshot_id(),
170            AnySnapshot::Mutable(s) => s.snapshot_id(),
171            AnySnapshot::NestedReadonly(s) => s.snapshot_id(),
172            AnySnapshot::NestedMutable(s) => s.snapshot_id(),
173            AnySnapshot::Global(s) => s.snapshot_id(),
174            AnySnapshot::TransparentMutable(s) => s.snapshot_id(),
175            AnySnapshot::TransparentReadonly(s) => s.snapshot_id(),
176        }
177    }
178
179    /// Get the set of invalid snapshot IDs.
180    fn state(&self) -> &SnapshotState {
181        match self {
182            AnySnapshot::Readonly(s) => &s.state,
183            AnySnapshot::Mutable(s) => &s.state,
184            AnySnapshot::NestedReadonly(s) => &s.state,
185            AnySnapshot::NestedMutable(s) => &s.state,
186            AnySnapshot::Global(s) => &s.state,
187            AnySnapshot::TransparentMutable(s) => &s.state,
188            AnySnapshot::TransparentReadonly(s) => &s.state,
189        }
190    }
191
192    pub(crate) fn with_invalid<R>(&self, f: impl FnOnce(&SnapshotIdSet) -> R) -> R {
193        f(&self.state().invalid.borrow())
194    }
195
196    pub fn invalid(&self) -> SnapshotIdSet {
197        match self {
198            AnySnapshot::Readonly(s) => s.invalid(),
199            AnySnapshot::Mutable(s) => s.invalid(),
200            AnySnapshot::NestedReadonly(s) => s.invalid(),
201            AnySnapshot::NestedMutable(s) => s.invalid(),
202            AnySnapshot::Global(s) => s.invalid(),
203            AnySnapshot::TransparentMutable(s) => s.invalid(),
204            AnySnapshot::TransparentReadonly(s) => s.invalid(),
205        }
206    }
207
208    /// Check if a snapshot ID is valid in this snapshot.
209    pub fn is_valid(&self, id: SnapshotId) -> bool {
210        let snapshot_id = self.snapshot_id();
211        id <= snapshot_id && !self.invalid().get(id)
212    }
213
214    /// Check if this is a read-only snapshot.
215    pub fn read_only(&self) -> bool {
216        match self {
217            AnySnapshot::Readonly(_) => true,
218            AnySnapshot::Mutable(_) => false,
219            AnySnapshot::NestedReadonly(_) => true,
220            AnySnapshot::NestedMutable(_) => false,
221            AnySnapshot::Global(_) => false,
222            AnySnapshot::TransparentMutable(_) => false,
223            AnySnapshot::TransparentReadonly(_) => true,
224        }
225    }
226
227    /// Get the root snapshot.
228    pub fn root(&self) -> AnySnapshot {
229        match self {
230            AnySnapshot::Readonly(s) => AnySnapshot::Readonly(s.root_readonly()),
231            AnySnapshot::Mutable(s) => AnySnapshot::Mutable(s.root_mutable()),
232            AnySnapshot::NestedReadonly(s) => AnySnapshot::NestedReadonly(s.root_nested_readonly()),
233            AnySnapshot::NestedMutable(s) => AnySnapshot::Mutable(s.root_mutable()),
234            AnySnapshot::Global(_) => AnySnapshot::Global(GlobalSnapshot::get_or_create()),
235            AnySnapshot::TransparentMutable(s) => {
236                AnySnapshot::TransparentMutable(s.root_transparent_mutable())
237            }
238            AnySnapshot::TransparentReadonly(s) => {
239                AnySnapshot::TransparentReadonly(s.root_transparent_readonly())
240            }
241        }
242    }
243
244    /// Check if this snapshot refers to the same transparent snapshot.
245    pub fn is_same_transparent(&self, other: &Rc<TransparentObserverMutableSnapshot>) -> bool {
246        matches!(self, AnySnapshot::TransparentMutable(snapshot) if Rc::ptr_eq(snapshot, other))
247    }
248
249    /// Check if this snapshot refers to the same transparent mutable snapshot.
250    pub fn is_same_transparent_mutable(
251        &self,
252        other: &Rc<TransparentObserverMutableSnapshot>,
253    ) -> bool {
254        self.is_same_transparent(other)
255    }
256
257    /// Check if this snapshot refers to the same transparent readonly snapshot.
258    pub fn is_same_transparent_readonly(&self, other: &Rc<TransparentObserverSnapshot>) -> bool {
259        matches!(self, AnySnapshot::TransparentReadonly(snapshot) if Rc::ptr_eq(snapshot, other))
260    }
261
262    /// Enter this snapshot, making it current for the duration of the closure.
263    pub fn enter<T>(&self, f: impl FnOnce() -> T) -> T {
264        match self {
265            AnySnapshot::Readonly(s) => s.enter(f),
266            AnySnapshot::Mutable(s) => s.enter(f),
267            AnySnapshot::NestedReadonly(s) => s.enter(f),
268            AnySnapshot::NestedMutable(s) => s.enter(f),
269            AnySnapshot::Global(s) => s.enter(f),
270            AnySnapshot::TransparentMutable(s) => s.enter(f),
271            AnySnapshot::TransparentReadonly(s) => s.enter(f),
272        }
273    }
274
275    /// Take a nested read-only snapshot.
276    pub fn take_nested_snapshot(&self, read_observer: Option<ReadObserver>) -> AnySnapshot {
277        match self {
278            AnySnapshot::Readonly(s) => {
279                AnySnapshot::Readonly(s.take_nested_snapshot(read_observer))
280            }
281            AnySnapshot::Mutable(s) => AnySnapshot::Readonly(s.take_nested_snapshot(read_observer)),
282            AnySnapshot::NestedReadonly(s) => {
283                AnySnapshot::NestedReadonly(s.take_nested_snapshot(read_observer))
284            }
285            AnySnapshot::NestedMutable(s) => {
286                AnySnapshot::Readonly(s.take_nested_snapshot(read_observer))
287            }
288            AnySnapshot::Global(s) => AnySnapshot::Readonly(s.take_nested_snapshot(read_observer)),
289            AnySnapshot::TransparentMutable(s) => {
290                AnySnapshot::Readonly(s.take_nested_snapshot(read_observer))
291            }
292            AnySnapshot::TransparentReadonly(s) => {
293                AnySnapshot::TransparentReadonly(s.take_nested_snapshot(read_observer))
294            }
295        }
296    }
297
298    /// Check if there are pending changes.
299    pub fn has_pending_changes(&self) -> bool {
300        match self {
301            AnySnapshot::Readonly(s) => s.has_pending_changes(),
302            AnySnapshot::Mutable(s) => s.has_pending_changes(),
303            AnySnapshot::NestedReadonly(s) => s.has_pending_changes(),
304            AnySnapshot::NestedMutable(s) => s.has_pending_changes(),
305            AnySnapshot::Global(s) => s.has_pending_changes(),
306            AnySnapshot::TransparentMutable(s) => s.has_pending_changes(),
307            AnySnapshot::TransparentReadonly(s) => s.has_pending_changes(),
308        }
309    }
310
311    /// Dispose of this snapshot.
312    pub fn dispose(&self) {
313        match self {
314            AnySnapshot::Readonly(s) => s.dispose(),
315            AnySnapshot::Mutable(s) => s.dispose(),
316            AnySnapshot::NestedReadonly(s) => s.dispose(),
317            AnySnapshot::NestedMutable(s) => s.dispose(),
318            AnySnapshot::Global(s) => s.dispose(),
319            AnySnapshot::TransparentMutable(s) => s.dispose(),
320            AnySnapshot::TransparentReadonly(s) => s.dispose(),
321        }
322    }
323
324    /// Check if disposed.
325    pub fn is_disposed(&self) -> bool {
326        match self {
327            AnySnapshot::Readonly(s) => s.is_disposed(),
328            AnySnapshot::Mutable(s) => s.is_disposed(),
329            AnySnapshot::NestedReadonly(s) => s.is_disposed(),
330            AnySnapshot::NestedMutable(s) => s.is_disposed(),
331            AnySnapshot::Global(s) => s.is_disposed(),
332            AnySnapshot::TransparentMutable(s) => s.is_disposed(),
333            AnySnapshot::TransparentReadonly(s) => s.is_disposed(),
334        }
335    }
336
337    /// Record a read.
338    pub fn record_read(&self, state: &dyn StateObject) {
339        match self {
340            AnySnapshot::Readonly(s) => s.record_read(state),
341            AnySnapshot::Mutable(s) => s.record_read(state),
342            AnySnapshot::NestedReadonly(s) => s.record_read(state),
343            AnySnapshot::NestedMutable(s) => s.record_read(state),
344            AnySnapshot::Global(s) => s.record_read(state),
345            AnySnapshot::TransparentMutable(s) => s.record_read(state),
346            AnySnapshot::TransparentReadonly(s) => s.record_read(state),
347        }
348    }
349
350    /// Record a write.
351    pub fn record_write(&self, state: Rc<dyn StateObject>) {
352        match self {
353            AnySnapshot::Readonly(s) => s.record_write(state),
354            AnySnapshot::Mutable(s) => s.record_write(state),
355            AnySnapshot::NestedReadonly(s) => s.record_write(state),
356            AnySnapshot::NestedMutable(s) => s.record_write(state),
357            AnySnapshot::Global(s) => s.record_write(state),
358            AnySnapshot::TransparentMutable(s) => s.record_write(state),
359            AnySnapshot::TransparentReadonly(s) => s.record_write(state),
360        }
361    }
362
363    /// Apply changes (only valid for mutable snapshots).
364    pub fn apply(&self) -> SnapshotApplyResult {
365        match self {
366            AnySnapshot::Mutable(s) => s.apply(),
367            AnySnapshot::NestedMutable(s) => s.apply(),
368            AnySnapshot::Global(s) => s.apply(),
369            AnySnapshot::TransparentMutable(s) => s.apply(),
370            _ => panic!("Cannot apply a read-only snapshot"),
371        }
372    }
373
374    /// Take a nested mutable snapshot (only valid for mutable snapshots).
375    pub fn take_nested_mutable_snapshot(
376        &self,
377        read_observer: Option<ReadObserver>,
378        write_observer: Option<WriteObserver>,
379    ) -> AnySnapshot {
380        match self {
381            AnySnapshot::Mutable(s) => AnySnapshot::NestedMutable(
382                s.take_nested_mutable_snapshot(read_observer, write_observer),
383            ),
384            AnySnapshot::NestedMutable(s) => AnySnapshot::NestedMutable(
385                s.take_nested_mutable_snapshot(read_observer, write_observer),
386            ),
387            AnySnapshot::Global(s) => {
388                AnySnapshot::Mutable(s.take_nested_mutable_snapshot(read_observer, write_observer))
389            }
390            AnySnapshot::TransparentMutable(s) => AnySnapshot::TransparentMutable(
391                s.take_nested_mutable_snapshot(read_observer, write_observer),
392            ),
393            _ => panic!("Cannot take nested mutable snapshot from read-only snapshot"),
394        }
395    }
396}
397
398thread_local! {
399    static CURRENT_SNAPSHOT: RefCell<Option<AnySnapshot>> = const { RefCell::new(None) };
400}
401
402/// Get the current snapshot, or None if not in a snapshot context.
403pub fn current_snapshot() -> Option<AnySnapshot> {
404    CURRENT_SNAPSHOT
405        .try_with(|cell| cell.borrow().clone())
406        .unwrap_or(None)
407}
408
409pub(crate) fn current_snapshot_reads_into(observer: &ReadObserver) -> bool {
410    CURRENT_SNAPSHOT
411        .try_with(|cell| match &*cell.borrow() {
412            Some(AnySnapshot::TransparentMutable(snapshot)) => {
413                snapshot.can_reuse() && snapshot.reads_into(observer)
414            }
415            _ => false,
416        })
417        .unwrap_or(false)
418}
419
420pub(crate) fn set_current_snapshot(snapshot: Option<AnySnapshot>) {
421    let _ = CURRENT_SNAPSHOT.try_with(|cell| {
422        *cell.borrow_mut() = snapshot;
423    });
424}
425
426struct CurrentSnapshotGuard {
427    previous: Option<AnySnapshot>,
428}
429
430impl CurrentSnapshotGuard {
431    fn enter(snapshot: AnySnapshot) -> Self {
432        let previous = CURRENT_SNAPSHOT
433            .try_with(|cell| cell.replace(Some(snapshot)))
434            .unwrap_or(None);
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) -> Rc<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<Rc<TransparentObserverMutableSnapshot>>,
495) -> Rc<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(Rc::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)) => Rc::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: &[Rc<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, Rc<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
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(Rc::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 Rc-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 Rc-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, (Rc<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: Rc<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: &Rc<P>,
938    child_id: SnapshotId,
939) -> impl FnOnce() + 'static
940where
941    P: NestedMutableHost + 'static,
942{
943    let weak = Rc::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: &Rc<P>,
960    root: Weak<MutableSnapshot>,
961    read_observer: Option<ReadObserver>,
962    write_observer: Option<WriteObserver>,
963) -> Rc<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
