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

cranpose_core/snapshot_v2/
mod.rs
