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 set_current_snapshot(snapshot: Option<AnySnapshot>) {
395    let _ = CURRENT_SNAPSHOT.try_with(|cell| {
396        *cell.borrow_mut() = snapshot;
397    });
398}
399
400struct CurrentSnapshotGuard {
401    previous: Option<AnySnapshot>,
402}
403
404impl CurrentSnapshotGuard {
405    fn enter(snapshot: AnySnapshot) -> Self {
406        let previous = current_snapshot();
407        set_current_snapshot(Some(snapshot));
408        Self { previous }
409    }
410}
411
412impl Drop for CurrentSnapshotGuard {
413    fn drop(&mut self) {
414        set_current_snapshot(self.previous.take());
415    }
416}
417
418pub(crate) fn enter_snapshot_scope<T>(snapshot: AnySnapshot, f: impl FnOnce() -> T) -> T {
419    let _guard = CurrentSnapshotGuard::enter(snapshot);
420    f()
421}
422
423/// Creates a mutable snapshot, matching Kotlin's `Snapshot.takeMutableSnapshot` semantics.
424///
425/// If called while inside a MutableSnapshot, creates a nested snapshot that will
426/// apply to the parent when `apply()` is called. This ensures proper isolation
427/// between nested operations (like event handlers during animations).
428///
429/// If called while inside a GlobalSnapshot or no snapshot, creates a root
430/// mutable snapshot that applies to the global state.
431pub fn take_mutable_snapshot(
432    read_observer: Option<ReadObserver>,
433    write_observer: Option<WriteObserver>,
434) -> AnyMutableSnapshot {
435    match current_snapshot() {
436        Some(AnySnapshot::Mutable(parent)) => AnyMutableSnapshot::Nested(
437            parent.take_nested_mutable_snapshot(read_observer, write_observer),
438        ),
439        Some(AnySnapshot::NestedMutable(parent)) => AnyMutableSnapshot::Nested(
440            parent.take_nested_mutable_snapshot(read_observer, write_observer),
441        ),
442        _ => AnyMutableSnapshot::Root(
443            GlobalSnapshot::get_or_create()
444                .take_nested_mutable_snapshot(read_observer, write_observer),
445        ),
446    }
447}
448
449/// Take a transparent observer mutable snapshot with optional observers.
450///
451/// This type of snapshot is used for read observation during composition,
452/// matching Kotlin's Snapshot.observeInternal behavior. It allows writes
453/// to happen during observation.
454///
455/// Transparent snapshots DO NOT allocate new IDs - they delegate to the
456/// current/global snapshot, making them "transparent" to the snapshot system.
457pub fn take_transparent_observer_mutable_snapshot(
458    read_observer: Option<ReadObserver>,
459    write_observer: Option<WriteObserver>,
460) -> Arc<TransparentObserverMutableSnapshot> {
461    take_transparent_observer_mutable_snapshot_reusing(read_observer, write_observer, None)
462}
463
464pub(crate) fn take_transparent_observer_mutable_snapshot_reusing(
465    read_observer: Option<ReadObserver>,
466    write_observer: Option<WriteObserver>,
467    recycled: Option<Arc<TransparentObserverMutableSnapshot>>,
468) -> Arc<TransparentObserverMutableSnapshot> {
469    let parent = current_snapshot();
470    match parent {
471        Some(AnySnapshot::TransparentMutable(transparent)) if transparent.can_reuse() => {
472            let (parent_read, parent_write) = transparent.observers();
473            if already_observes(&read_observer, &parent_read)
474                && already_observes(&write_observer, &parent_write)
475            {
476                return transparent;
477            }
478            TransparentObserverMutableSnapshot::new_reusing(
479                recycled,
480                transparent.snapshot_id(),
481                transparent.invalid(),
482                merge_read_observers(read_observer, parent_read),
483                merge_write_observers(write_observer, parent_write),
484                Some(Arc::downgrade(&transparent)),
485            )
486        }
487        _ => {
488            let current = current_snapshot()
489                .unwrap_or_else(|| AnySnapshot::Global(GlobalSnapshot::get_or_create()));
490            let id = current.snapshot_id();
491            let invalid = current.invalid();
492            TransparentObserverMutableSnapshot::new_reusing(
493                recycled,
494                id,
495                invalid,
496                read_observer,
497                write_observer,
498                None,
499            )
500        }
501    }
502}
503
504fn already_observes(requested: &Option<ReadObserver>, installed: &Option<ReadObserver>) -> bool {
505    match (requested, installed) {
506        (None, _) => true,
507        (Some(requested), Some(installed)) => Arc::ptr_eq(requested, installed),
508        (Some(_), None) => false,
509    }
510}
511
512/// Allocate a new record identifier that is distinct from any active snapshot id.
513pub fn allocate_record_id() -> SnapshotId {
514    runtime::allocate_record_id()
515}
516
517pub(crate) fn peek_next_snapshot_id() -> SnapshotId {
518    runtime::peek_next_snapshot_id()
519}
520
521#[derive(Clone)]
522struct ObserverId(Rc<()>);
523
524impl ObserverId {
525    fn new() -> Self {
526        Self(Rc::new(()))
527    }
528}
529
530impl PartialEq for ObserverId {
531    fn eq(&self, other: &Self) -> bool {
532        Rc::ptr_eq(&self.0, &other.0)
533    }
534}
535
536impl Eq for ObserverId {}
537
538impl Hash for ObserverId {
539    fn hash<H: Hasher>(&self, state: &mut H) {
540        Rc::as_ptr(&self.0).hash(state);
541    }
542}
543
544thread_local! {
545    static APPLY_OBSERVERS: RefCell<HashMap<ObserverId, ApplyObserver>> = RefCell::new(HashMap::default());
546}
547
548thread_local! {
549    static LAST_WRITES: RefCell<HashMap<StateObjectId, SnapshotId>> = RefCell::new(HashMap::default());
550}
551
552thread_local! {
553    static EXTRA_STATE_OBJECTS: RefCell<crate::snapshot_weak_set::SnapshotWeakSet> = RefCell::new(crate::snapshot_weak_set::SnapshotWeakSet::new());
554}
555
556const UNUSED_RECORD_CLEANUP_INTERVAL: SnapshotId = 2;
557const UNUSED_RECORD_CLEANUP_BUSY_INTERVAL: SnapshotId = 1;
558const UNUSED_RECORD_CLEANUP_MIN_SIZE: usize = 64;
559
560thread_local! {
561    static LAST_UNUSED_RECORD_CLEANUP: Cell<SnapshotId> = const { Cell::new(0) };
562}
563
564#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
565pub struct SnapshotV2DebugStats {
566    pub apply_observers_len: usize,
567    pub apply_observers_cap: usize,
568    pub last_writes_len: usize,
569    pub last_writes_cap: usize,
570    pub extra_state_objects_len: usize,
571    pub extra_state_objects_cap: usize,
572    pub last_unused_record_cleanup: SnapshotId,
573}
574
575pub fn debug_snapshot_v2_stats() -> SnapshotV2DebugStats {
576    let (apply_observers_len, apply_observers_cap) = APPLY_OBSERVERS.with(|cell| {
577        let observers = cell.borrow();
578        (observers.len(), observers.capacity())
579    });
580    let (last_writes_len, last_writes_cap) = LAST_WRITES.with(|cell| {
581        let writes = cell.borrow();
582        (writes.len(), writes.capacity())
583    });
584    let SnapshotWeakSetDebugStats {
585        len: extra_state_objects_len,
586        capacity: extra_state_objects_cap,
587    } = EXTRA_STATE_OBJECTS.with(|cell| cell.borrow().debug_stats());
588    let last_unused_record_cleanup = LAST_UNUSED_RECORD_CLEANUP.with(Cell::get);
589
590    SnapshotV2DebugStats {
591        apply_observers_len,
592        apply_observers_cap,
593        last_writes_len,
594        last_writes_cap,
595        extra_state_objects_len,
596        extra_state_objects_cap,
597        last_unused_record_cleanup,
598    }
599}
600
601/// Register an apply observer.
602///
603/// Returns a handle that will automatically unregister the observer when dropped.
604pub fn register_apply_observer(observer: ApplyObserver) -> ObserverHandle {
605    let id = ObserverId::new();
606    APPLY_OBSERVERS.with(|cell| {
607        cell.borrow_mut().insert(id.clone(), observer);
608    });
609    ObserverHandle {
610        kind: ObserverKind::Apply,
611        id,
612    }
613}
614
615/// Handle for unregistering observers.
616///
617/// When dropped, automatically removes the associated observer.
618pub struct ObserverHandle {
619    kind: ObserverKind,
620    id: ObserverId,
621}
622
623enum ObserverKind {
624    Apply,
625}
626
627impl Drop for ObserverHandle {
628    fn drop(&mut self) {
629        match self.kind {
630            ObserverKind::Apply => {
631                APPLY_OBSERVERS.with(|cell| {
632                    cell.borrow_mut().remove(&self.id);
633                });
634            }
635        }
636    }
637}
638
639pub(crate) fn notify_apply_observers(modified: &[Arc<dyn StateObject>], snapshot_id: SnapshotId) {
640    APPLY_OBSERVERS.with(|cell| {
641        let observers: Vec<ApplyObserver> = cell.borrow().values().cloned().collect();
642        for observer in observers.into_iter() {
643            observer(modified, snapshot_id);
644        }
645    });
646}
647
648pub(crate) fn set_last_write(id: StateObjectId, snapshot_id: SnapshotId) {
649    LAST_WRITES.with(|cell| {
650        cell.borrow_mut().insert(id, snapshot_id);
651    });
652}
653
654#[cfg(test)]
655pub(crate) fn clear_last_writes() {
656    LAST_WRITES.with(|cell| {
657        cell.borrow_mut().clear();
658    });
659}
660
661pub(crate) fn check_and_overwrite_unused_records_locked() {
662    EXTRA_STATE_OBJECTS.with(|cell| {
663        cell.borrow_mut()
664            .remove_if(super::state::StateObject::overwrite_unused_records);
665    });
666}
667
668pub(crate) fn maybe_check_and_overwrite_unused_records_locked(current_snapshot_id: SnapshotId) {
669    let should_run = EXTRA_STATE_OBJECTS.with(|cell| {
670        let set = cell.borrow();
671        if set.is_empty() {
672            return false;
673        }
674        let last_cleanup = LAST_UNUSED_RECORD_CLEANUP.with(Cell::get);
675        let interval = if set.len() >= UNUSED_RECORD_CLEANUP_MIN_SIZE {
676            UNUSED_RECORD_CLEANUP_BUSY_INTERVAL
677        } else {
678            UNUSED_RECORD_CLEANUP_INTERVAL
679        };
680        current_snapshot_id.saturating_sub(last_cleanup) >= interval
681    });
682
683    if should_run {
684        LAST_UNUSED_RECORD_CLEANUP.with(|cell| cell.set(current_snapshot_id));
685        check_and_overwrite_unused_records_locked();
686    }
687}
688
689#[cfg(test)]
690pub(crate) fn clear_unused_record_cleanup_for_tests() {
691    LAST_UNUSED_RECORD_CLEANUP.with(|cell| cell.set(0));
692}
693
694pub(crate) fn optimistic_merges(
695    current_snapshot_id: SnapshotId,
696    base_parent_id: SnapshotId,
697    modified_objects: &[(StateObjectId, Arc<dyn StateObject>, SnapshotId)],
698    invalid_snapshots: &SnapshotIdSet,
699    applying_invalid: &SnapshotIdSet,
700) -> Option<HashMap<usize, Rc<StateRecord>>> {
701    if modified_objects.is_empty() {
702        return None;
703    }
704
705    let mut result: Option<HashMap<usize, Rc<StateRecord>>> = None;
706
707    for (_, state, writer_id) in modified_objects {
708        let head = state.first_record();
709
710        let Some(current) =
711            crate::state::readable_record_for(&head, current_snapshot_id, invalid_snapshots)
712        else {
713            continue;
714        };
715
716        let (previous_opt, found_base) =
717            mutable::find_previous_record(&head, base_parent_id, applying_invalid);
718        let previous = previous_opt?;
719
720        if !found_base || previous.snapshot_id() == crate::state::PREEXISTING_SNAPSHOT_ID {
721            continue;
722        }
723
724        if Rc::ptr_eq(&current, &previous) {
725            continue;
726        }
727
728        let applied = mutable::find_record_by_id(&head, *writer_id)?;
729
730        let merged = state.merge_records(
731            Rc::clone(&previous),
732            Rc::clone(&current),
733            Rc::clone(&applied),
734        )?;
735
736        result
737            .get_or_insert_with(HashMap::default)
738            .insert(Rc::as_ptr(&current) as usize, merged);
739    }
740
741    result
742}
743
744#[expect(clippy::arc_with_non_send_sync)]
745fn merge_observers(a: Option<ReadObserver>, b: Option<ReadObserver>) -> Option<ReadObserver> {
746    match (a, b) {
747        (None, None) => None,
748        (Some(a), None) => Some(a),
749        (None, Some(b)) => Some(b),
750        (Some(a), Some(b)) => Some(Arc::new(move |state: &dyn StateObject| {
751            a(state);
752            b(state);
753        })),
754    }
755}
756
757/// Merge two read observers into one.
758///
759/// # Thread Safety
760/// The resulting Arc-wrapped closure may capture non-Send closures. This is safe
761/// because observers are only invoked on the UI thread where they were created.
762pub fn merge_read_observers(
763    a: Option<ReadObserver>,
764    b: Option<ReadObserver>,
765) -> Option<ReadObserver> {
766    merge_observers(a, b)
767}
768
769/// Merge two write observers into one.
770///
771/// # Thread Safety
772/// The resulting Arc-wrapped closure may capture non-Send closures. This is safe
773/// because observers are only invoked on the UI thread where they were created.
774pub fn merge_write_observers(
775    a: Option<WriteObserver>,
776    b: Option<WriteObserver>,
777) -> Option<WriteObserver> {
778    merge_observers(a, b)
779}
780
781pub(crate) struct SnapshotState {
782    pub(crate) id: Cell<SnapshotId>,
783    pub(crate) invalid: RefCell<SnapshotIdSet>,
784    pub(crate) pin_handle: Cell<PinHandle>,
785    pub(crate) disposed: Cell<bool>,
786    pub(crate) read_observer: RefCell<Option<ReadObserver>>,
787    pub(crate) write_observer: RefCell<Option<WriteObserver>>,
788    #[expect(clippy::type_complexity)]
789    pub(crate) modified: RefCell<HashMap<StateObjectId, (Arc<dyn StateObject>, SnapshotId)>>,
790    on_dispose: RefCell<Option<Box<dyn FnOnce()>>>,
791    runtime_tracked: bool,
792    pending_children: RefCell<HashSet<SnapshotId>>,
793}
794
795impl SnapshotState {
796    pub(crate) fn new(
797        id: SnapshotId,
798        invalid: SnapshotIdSet,
799        read_observer: Option<ReadObserver>,
800        write_observer: Option<WriteObserver>,
801        runtime_tracked: bool,
802    ) -> Self {
803        Self::new_with_pinning(
804            id,
805            invalid,
806            read_observer,
807            write_observer,
808            runtime_tracked,
809            true,
810        )
811    }
812
813    pub(crate) fn new_with_pinning(
814        id: SnapshotId,
815        invalid: SnapshotIdSet,
816        read_observer: Option<ReadObserver>,
817        write_observer: Option<WriteObserver>,
818        runtime_tracked: bool,
819        should_pin: bool,
820    ) -> Self {
821        let pin_handle = if should_pin {
822            snapshot_pinning::track_pinning(id, &invalid)
823        } else {
824            snapshot_pinning::PinHandle::INVALID
825        };
826        Self {
827            id: Cell::new(id),
828            invalid: RefCell::new(invalid),
829            pin_handle: Cell::new(pin_handle),
830            disposed: Cell::new(false),
831            read_observer: RefCell::new(read_observer),
832            write_observer: RefCell::new(write_observer),
833            modified: RefCell::new(HashMap::default()),
834            on_dispose: RefCell::new(None),
835            runtime_tracked,
836            pending_children: RefCell::new(HashSet::default()),
837        }
838    }
839
840    pub(crate) fn record_read(&self, state: &dyn StateObject) {
841        if let Some(observer) = self.read_observer.borrow().as_ref() {
842            observer(state);
843        }
844    }
845
846    pub(crate) fn record_write(&self, state: Arc<dyn StateObject>, writer_id: SnapshotId) {
847        let state_id = state.object_id().as_usize();
848
849        let mut modified = self.modified.borrow_mut();
850
851        match modified.entry(state_id) {
852            std::collections::hash_map::Entry::Vacant(e) => {
853                if let Some(observer) = self.write_observer.borrow().as_ref() {
854                    observer(&*state);
855                }
856                e.insert((state, writer_id));
857            }
858            std::collections::hash_map::Entry::Occupied(mut e) => {
859                e.insert((state, writer_id));
860            }
861        }
862    }
863
864    pub(crate) fn dispose(&self) {
865        if !self.disposed.replace(true) {
866            let pin_handle = self.pin_handle.get();
867            snapshot_pinning::release_pinning(pin_handle);
868            if let Some(cb) = self.on_dispose.borrow_mut().take() {
869                cb();
870            }
871            if self.runtime_tracked {
872                close_snapshot(self.id.get());
873            }
874        }
875    }
876
877    pub(crate) fn add_pending_child(&self, id: SnapshotId) {
878        self.pending_children.borrow_mut().insert(id);
879    }
880
881    pub(crate) fn remove_pending_child(&self, id: SnapshotId) {
882        self.pending_children.borrow_mut().remove(&id);
883    }
884
885    pub(crate) fn has_pending_children(&self) -> bool {
886        !self.pending_children.borrow().is_empty()
887    }
888
889    pub(crate) fn pending_children(&self) -> Vec<SnapshotId> {
890        self.pending_children.borrow().iter().copied().collect()
891    }
892
893    pub(crate) fn set_on_dispose<F>(&self, f: F)
894    where
895        F: FnOnce() + 'static,
896    {
897        *self.on_dispose.borrow_mut() = Some(Box::new(f));
898    }
899}
900
901pub(crate) trait NestedMutableHost {
902    fn snapshot_state(&self) -> &SnapshotState;
903    fn nested_count(&self) -> &Cell<usize>;
904}
905
906pub(crate) fn clear_nested_child_on_dispose<P>(
907    parent: &Arc<P>,
908    child_id: SnapshotId,
909) -> impl FnOnce() + 'static
910where
911    P: NestedMutableHost + 'static,
912{
913    let weak = Arc::downgrade(parent);
914    move || {
915        if let Some(parent) = weak.upgrade() {
916            let nested_count = parent.nested_count();
917            if nested_count.get() > 0 {
918                nested_count.set(nested_count.get().saturating_sub(1));
919            }
920            let state = parent.snapshot_state();
921            let new_invalid = state.invalid.borrow().clone().clear(child_id);
922            state.invalid.replace(new_invalid);
923            state.remove_pending_child(child_id);
924        }
925    }
926}
927
928pub(crate) fn allocate_nested_mutable_snapshot<P>(
929    parent: &Arc<P>,
930    root: Weak<MutableSnapshot>,
931    read_observer: Option<ReadObserver>,
932    write_observer: Option<WriteObserver>,
933) -> Arc<NestedMutableSnapshot>
934where
935    P: NestedMutableHost + 'static,
936{
937    let state = parent.snapshot_state();
938    let merged_read = merge_read_observers(read_observer, state.read_observer.borrow().clone());
939    let merged_write = merge_write_observers(write_observer, state.write_observer.borrow().clone());
940
941    let parent_id = state.id.get();
942    let current_invalid = state.invalid.borrow().clone();
943
944    let (new_id, _runtime_invalid) = allocate_snapshot();
945
946    let parent_invalid_with_child = current_invalid.set(new_id);
947    state.invalid.replace(parent_invalid_with_child);
948
949    let invalid = current_invalid.add_range(parent_id + 1, new_id);
950
951    let nested = NestedMutableSnapshot::new(
952        new_id,
953        invalid,
954        merged_read,
955        merged_write,
956        root,
957        state.id.get(),
958    );
959
960    let nested_count = parent.nested_count();
961    nested_count.set(nested_count.get() + 1);
962    state.add_pending_child(new_id);
963
964    nested.set_on_dispose(clear_nested_child_on_dispose(parent, new_id));
965
966    nested
967}
968
969#[cfg(test)]
970#[path = "tests/snapshot_v2_tests.rs"]
971mod tests;
cranpose_core/snapshot_v2/mod.rs

cranpose_core/snapshot_v2/
mod.rs
