cranpose_core/snapshot_v2/
mod.rs

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

cranpose_core/snapshot_v2/
mod.rs
