cranpose_core/snapshot_v2/
mod.rs

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

cranpose_core/snapshot_v2/
mod.rs
