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

cranpose_core/snapshot_v2/
mod.rs
