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

cranpose_core/snapshot_v2/
mod.rs
