Skip to main content

cranpose_core/
snapshot_state_observer.rs

1use std::{
2    any::{Any, TypeId},
3    cell::{Cell, RefCell},
4    hash::{Hash, Hasher},
5    rc::{Rc, Weak},
6    sync::Arc,
7};
8
9use smallvec::SmallVec;
10
11use crate::{
12    RecomposeScope, RecomposeScopeInner, ScopeId,
13    collections::map::{HashMap, HashSet},
14    hash::default as default_hash,
15    snapshot_v2::{
16        ReadObserver, StateObjectId, TransparentObserverMutableSnapshot, register_apply_observer,
17    },
18    state::StateObject,
19};
20
21type Executor = dyn Fn(Box<dyn FnOnce() + 'static>) + 'static;
22
23trait ScopeChangedCallback: Fn(&dyn Any) + Any {}
24
25impl<F: Fn(&dyn Any) + Any> ScopeChangedCallback for F {}
26
27/// Observer that records state object reads performed inside a given scope and
28/// notifies the caller when any of the observed objects change.
29///
30/// This is a pragmatic Rust translation of Jetpack Compose's
31/// `SnapshotStateObserver`. The implementation focuses on the core behaviour
32/// needed by the Cranpose runtime:
33/// - Tracking state object reads per logical scope.
34/// - Reacting to snapshot apply notifications.
35/// - Scheduling invalidation callbacks via the supplied executor.
36///
37/// Advanced features from the Kotlin version (derived state tracking, change
38/// coalescing, queue minimisation) are deferred
39#[derive(Clone)]
40pub struct SnapshotStateObserver {
41    inner: Rc<SnapshotStateObserverInner>,
42}
43
44#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
45pub struct SnapshotStateObserverDebugStats {
46    pub scopes_len: usize,
47    pub scopes_cap: usize,
48    pub fast_scopes_len: usize,
49    pub fast_scopes_cap: usize,
50    pub stateless_scope_count: usize,
51    pub observed_state_count: usize,
52    pub observed_state_capacity: usize,
53}
54
55impl SnapshotStateObserver {
56    /// Create a new observer that schedules callbacks using `on_changed_executor`.
57    pub fn new(on_changed_executor: impl Fn(Box<dyn FnOnce() + 'static>) + 'static) -> Self {
58        let inner = Rc::new(SnapshotStateObserverInner::new(on_changed_executor));
59        inner.set_self(Rc::downgrade(&inner));
60        Self { inner }
61    }
62
63    /// Observe state object reads performed while executing `block`.
64    ///
65    /// Subsequent calls to `observe_reads` replace any previously recorded
66    /// observations for the provided `scope`. When one of the observed objects
67    /// mutates, `on_value_changed_for_scope` will be invoked on the executor.
68    pub fn observe_reads<T, R>(
69        &self,
70        scope: T,
71        on_value_changed_for_scope: impl Fn(&T) + 'static,
72        block: impl FnOnce() -> R,
73    ) -> R
74    where
75        T: Any + Clone + Eq + Hash + 'static,
76    {
77        self.inner
78            .observe_reads(scope, on_value_changed_for_scope, block)
79    }
80
81    /// Notify the observer that a new composition frame is starting.
82    pub fn begin_frame(&self) {
83        self.inner.begin_frame();
84    }
85
86    /// Drop bookkeeping for scopes that were released during the current frame.
87    pub fn prune_dead_scopes(&self) {
88        self.inner.prune_dead_scopes();
89    }
90
91    /// Temporarily pause read observation while executing `block`.
92    pub fn with_no_observations<R>(&self, block: impl FnOnce() -> R) -> R {
93        self.inner.with_no_observations(block)
94    }
95
96    /// Remove any recorded reads for `scope`.
97    pub fn clear<T>(&self, scope: &T)
98    where
99        T: Any + Eq + Hash + 'static,
100    {
101        self.inner.clear(scope);
102    }
103
104    /// Remove recorded reads for scopes that satisfy `predicate`.
105    pub fn clear_if(&self, predicate: impl Fn(&dyn Any) -> bool) {
106        self.inner.clear_if(predicate);
107    }
108
109    /// Remove all recorded observations.
110    pub fn clear_all(&self) {
111        self.inner.clear_all();
112    }
113
114    /// Begin listening for snapshot apply notifications.
115    pub fn start(&self) {
116        let weak = Rc::downgrade(&self.inner);
117        self.inner.start(weak);
118    }
119
120    /// Stop listening for snapshot apply notifications.
121    pub fn stop(&self) {
122        self.inner.stop();
123    }
124
125    pub fn debug_stats(&self) -> SnapshotStateObserverDebugStats {
126        self.inner.debug_stats()
127    }
128
129    #[cfg(test)]
130    pub fn notify_changes(&self, modified: &[Arc<dyn StateObject>]) {
131        self.inner.handle_apply(modified);
132    }
133}
134
135struct SnapshotStateObserverInner {
136    executor: Rc<Executor>,
137    owned_scopes: RefCell<HashMap<OwnedScopeIndexKey, OwnedScopeBucket>>,
138    fast_scopes: RefCell<HashMap<ScopeId, Rc<RefCell<ScopeEntry>>>>,
139    indexed_scopes: RefCell<HashMap<usize, Rc<RefCell<ScopeEntry>>>>,
140    observed_to_scopes: RefCell<HashMap<StateObjectId, HashSet<usize>>>,
141    pause_count: Rc<Cell<usize>>,
142    active_read_targets: Rc<RefCell<ReadObservationStack>>,
143    read_dispatcher: ReadObserver,
144    read_snapshot: RefCell<Option<Arc<TransparentObserverMutableSnapshot>>>,
145    apply_handle: RefCell<Option<crate::snapshot_v2::ObserverHandle>>,
146    weak_self: RefCell<Weak<SnapshotStateObserverInner>>,
147    frame_version: Cell<u64>,
148    next_entry_id: Cell<usize>,
149}
150
151#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
152struct OwnedScopeIndexKey {
153    type_id: TypeId,
154    value_hash: u64,
155}
156
157type OwnedScopeBucket = SmallVec<[Rc<RefCell<ScopeEntry>>; 1]>;
158
159fn owned_scope_index_key<T>(scope: &T) -> OwnedScopeIndexKey
160where
161    T: Any + Hash + 'static,
162{
163    let mut hasher = default_hash::new();
164    scope.hash(&mut hasher);
165    OwnedScopeIndexKey {
166        type_id: TypeId::of::<T>(),
167        value_hash: hasher.finish(),
168    }
169}
170
171impl SnapshotStateObserverInner {
172    const MIN_RETAINED_SCOPE_CAPACITY: usize = 256;
173
174    fn new(on_changed_executor: impl Fn(Box<dyn FnOnce() + 'static>) + 'static) -> Self {
175        let pause_count = Rc::new(Cell::new(0));
176        let active_read_targets = Rc::new(RefCell::new(ReadObservationStack::default()));
177        let dispatcher_pause_count = Rc::clone(&pause_count);
178        let dispatcher_targets = Rc::clone(&active_read_targets);
179        let read_dispatcher: ReadObserver = Arc::new(move |state| {
180            if dispatcher_pause_count.get() > 0 {
181                return;
182            }
183            let observed = dispatcher_targets.borrow().last().cloned();
184            if let Some(observed) = observed {
185                observed.borrow_mut().insert(state);
186            }
187        });
188
189        Self {
190            executor: Rc::new(on_changed_executor),
191            owned_scopes: RefCell::new(HashMap::default()),
192            fast_scopes: RefCell::new(HashMap::default()),
193            indexed_scopes: RefCell::new(HashMap::default()),
194            observed_to_scopes: RefCell::new(HashMap::default()),
195            pause_count,
196            active_read_targets,
197            read_dispatcher,
198            read_snapshot: RefCell::new(None),
199            apply_handle: RefCell::new(None),
200            weak_self: RefCell::new(Weak::new()),
201            frame_version: Cell::new(0),
202            next_entry_id: Cell::new(0),
203        }
204    }
205
206    fn set_self(&self, weak: Weak<SnapshotStateObserverInner>) {
207        self.weak_self.replace(weak);
208    }
209
210    fn begin_frame(&self) {
211        let next = self.frame_version.get().wrapping_add(1);
212        self.frame_version.set(next);
213        self.prune_dead_scopes();
214    }
215
216    fn observe_reads<T, R>(
217        &self,
218        scope: T,
219        on_value_changed_for_scope: impl Fn(&T) + 'static,
220        block: impl FnOnce() -> R,
221    ) -> R
222    where
223        T: Any + Clone + Eq + Hash + 'static,
224    {
225        let frame_version = self.frame_version.get();
226        let has_frame_version = frame_version != 0;
227
228        let existing_entry = self.find_scope_entry(&scope);
229        let on_changed = std::cell::LazyCell::new(|| {
230            let callback = move |scope_any: &dyn Any| {
231                if let Some(typed) = scope_any.downcast_ref::<T>() {
232                    on_value_changed_for_scope(typed);
233                }
234            };
235            match existing_entry.as_ref() {
236                Some(entry) => entry.borrow_mut().callback_reusing(callback),
237                None => Rc::new(callback),
238            }
239        });
240
241        if let Some(entry) = existing_entry.as_ref() {
242            let already_observed = {
243                let mut entry_mut = entry.borrow_mut();
244                entry_mut.update_scope(scope.clone());
245                has_frame_version && entry_mut.last_seen_version == frame_version
246            };
247            let callback = on_changed.clone();
248            entry.borrow_mut().on_changed = callback;
249            if already_observed {
250                return block();
251            }
252        }
253
254        let observed = self.active_read_targets.borrow_mut().push();
255        struct ActiveObservationGuard {
256            stack: Rc<RefCell<ReadObservationStack>>,
257        }
258        impl Drop for ActiveObservationGuard {
259            fn drop(&mut self) {
260                let target = self.stack.borrow_mut().pop();
261                let discarded = target.replace(ObservedIds::new());
262                drop(discarded);
263            }
264        }
265        let _guard = ActiveObservationGuard {
266            stack: Rc::clone(&self.active_read_targets),
267        };
268
269        let result = self.run_with_read_observer(block);
270
271        if observed.borrow().is_empty() {
272            if existing_entry.is_some() {
273                self.clear(&scope);
274            }
275            return result;
276        }
277
278        let observed = {
279            let mut observed = observed.borrow_mut();
280            std::mem::replace(&mut *observed, ObservedIds::new())
281        };
282        let entry = existing_entry
283            .clone()
284            .unwrap_or_else(|| self.insert_scope_entry(scope.clone(), on_changed.clone()));
285        {
286            let mut entry_mut = entry.borrow_mut();
287            entry_mut.update(scope, Rc::clone(&on_changed));
288            entry_mut.last_seen_version = if has_frame_version {
289                frame_version
290            } else {
291                u64::MAX
292            };
293        }
294        self.replace_observed_ids(&entry, observed);
295
296        result
297    }
298
299    fn with_no_observations<R>(&self, block: impl FnOnce() -> R) -> R {
300        self.pause_count.set(self.pause_count.get() + 1);
301        let result = block();
302        self.pause_count
303            .set(self.pause_count.get().saturating_sub(1));
304        result
305    }
306
307    fn clear<T>(&self, scope: &T)
308    where
309        T: Any + Eq + Hash + 'static,
310    {
311        if let Some(rc_scope) = (scope as &dyn Any).downcast_ref::<RecomposeScope>() {
312            if let Some(entry) = self.fast_scopes.borrow_mut().remove(&rc_scope.id()) {
313                self.unregister_entry(&entry);
314            }
315            return;
316        }
317
318        let removed = self.remove_owned_scope_entry(scope);
319        if let Some(entry) = removed {
320            self.unregister_entry(&entry);
321        }
322    }
323
324    fn clear_if(&self, predicate: impl Fn(&dyn Any) -> bool) {
325        let removed_fast = {
326            let mut fast_scopes = self.fast_scopes.borrow_mut();
327            let removed_ids: Vec<_> = fast_scopes
328                .iter()
329                .filter(|(_, entry)| entry.borrow().matches_predicate(&predicate))
330                .map(|(scope_id, _)| *scope_id)
331                .collect();
332            removed_ids
333                .into_iter()
334                .filter_map(|scope_id| fast_scopes.remove(&scope_id))
335                .collect::<Vec<_>>()
336        };
337        let removed_owned =
338            { self.partition_owned_scopes(|entry| entry.matches_predicate(&predicate)) };
339
340        for entry in removed_fast.into_iter().chain(removed_owned) {
341            self.unregister_entry(&entry);
342        }
343    }
344
345    fn clear_all(&self) {
346        self.fast_scopes.borrow_mut().clear();
347        self.owned_scopes.borrow_mut().clear();
348        self.indexed_scopes.borrow_mut().clear();
349        self.observed_to_scopes.borrow_mut().clear();
350    }
351
352    fn start(&self, weak_self: Weak<SnapshotStateObserverInner>) {
353        if self.apply_handle.borrow().is_some() {
354            return;
355        }
356
357        let handle = register_apply_observer(Rc::new(move |modified, _snapshot_id| {
358            if let Some(inner) = weak_self.upgrade() {
359                inner.handle_apply(modified);
360            }
361        }));
362        self.apply_handle.replace(Some(handle));
363    }
364
365    fn stop(&self) {
366        if let Some(handle) = self.apply_handle.borrow_mut().take() {
367            drop(handle);
368        }
369    }
370
371    fn find_scope_entry<T>(&self, scope: &T) -> Option<Rc<RefCell<ScopeEntry>>>
372    where
373        T: Any + Eq + Hash + 'static,
374    {
375        if let Some(scope) = (scope as &dyn Any).downcast_ref::<RecomposeScope>() {
376            return self.fast_scopes.borrow().get(&scope.id()).cloned();
377        }
378
379        self.find_owned_scope_entry(scope)
380    }
381
382    fn insert_scope_entry(
383        &self,
384        scope: impl Any + Clone + Eq + Hash + 'static,
385        on_changed: Rc<dyn ScopeChangedCallback>,
386    ) -> Rc<RefCell<ScopeEntry>> {
387        let entry_id = self.next_entry_id.get();
388        self.next_entry_id.set(entry_id.wrapping_add(1));
389        let recompose_scope_id = (&scope as &dyn Any)
390            .downcast_ref::<RecomposeScope>()
391            .map(RecomposeScope::id);
392        let owned_scope_key = recompose_scope_id
393            .is_none()
394            .then(|| owned_scope_index_key(&scope));
395        let entry = Rc::new(RefCell::new(ScopeEntry::new(entry_id, scope, on_changed)));
396        self.indexed_scopes
397            .borrow_mut()
398            .insert(entry_id, Rc::clone(&entry));
399        if let Some(scope_id) = recompose_scope_id {
400            self.fast_scopes
401                .borrow_mut()
402                .insert(scope_id, Rc::clone(&entry));
403        } else if let Some(scope_key) = owned_scope_key {
404            self.owned_scopes
405                .borrow_mut()
406                .entry(scope_key)
407                .or_default()
408                .push(Rc::clone(&entry));
409        }
410        entry
411    }
412
413    fn prune_dead_scopes(&self) {
414        let removed_fast = {
415            let mut fast_scopes = self.fast_scopes.borrow_mut();
416            let removed_ids: Vec<_> = fast_scopes
417                .iter()
418                .filter(|(_, entry)| !entry.borrow().should_retain())
419                .map(|(scope_id, _)| *scope_id)
420                .collect();
421            let removed = removed_ids
422                .into_iter()
423                .filter_map(|scope_id| fast_scopes.remove(&scope_id))
424                .collect::<Vec<_>>();
425            shrink_map_if_sparse(&mut fast_scopes, Self::MIN_RETAINED_SCOPE_CAPACITY);
426            removed
427        };
428
429        let removed_owned = { self.partition_owned_scopes(|entry| !entry.should_retain()) };
430
431        for entry in removed_fast.into_iter().chain(removed_owned) {
432            self.unregister_entry(&entry);
433        }
434    }
435
436    fn find_owned_scope_entry<T>(&self, scope: &T) -> Option<Rc<RefCell<ScopeEntry>>>
437    where
438        T: Any + Eq + Hash + 'static,
439    {
440        let key = owned_scope_index_key(scope);
441        self.owned_scopes.borrow().get(&key).and_then(|bucket| {
442            bucket
443                .iter()
444                .find(|entry| entry.borrow().matches_scope(scope))
445                .cloned()
446        })
447    }
448
449    fn remove_owned_scope_entry<T>(&self, scope: &T) -> Option<Rc<RefCell<ScopeEntry>>>
450    where
451        T: Any + Eq + Hash + 'static,
452    {
453        let key = owned_scope_index_key(scope);
454        let mut owned_scopes = self.owned_scopes.borrow_mut();
455        let mut removed = None;
456        let mut remove_bucket = false;
457        if let Some(bucket) = owned_scopes.get_mut(&key)
458            && let Some(index) = bucket
459                .iter()
460                .position(|entry| entry.borrow().matches_scope(scope))
461        {
462            removed = Some(bucket.remove(index));
463            remove_bucket = bucket.is_empty();
464        }
465        if remove_bucket {
466            owned_scopes.remove(&key);
467        }
468        shrink_map_if_sparse(&mut owned_scopes, Self::MIN_RETAINED_SCOPE_CAPACITY);
469        removed
470    }
471
472    fn partition_owned_scopes(
473        &self,
474        should_remove: impl Fn(&ScopeEntry) -> bool,
475    ) -> Vec<Rc<RefCell<ScopeEntry>>> {
476        let mut owned_scopes = self.owned_scopes.borrow_mut();
477        let mut retained = HashMap::default();
478        let mut removed = Vec::new();
479        for (key, mut bucket) in owned_scopes.drain() {
480            let mut retained_bucket = OwnedScopeBucket::new();
481            for entry in bucket.drain(..) {
482                if should_remove(&entry.borrow()) {
483                    removed.push(entry);
484                } else {
485                    retained_bucket.push(entry);
486                }
487            }
488            if !retained_bucket.is_empty() {
489                retained.insert(key, retained_bucket);
490            }
491        }
492        *owned_scopes = retained;
493        shrink_map_if_sparse(&mut owned_scopes, Self::MIN_RETAINED_SCOPE_CAPACITY);
494        removed
495    }
496
497    fn debug_stats(&self) -> SnapshotStateObserverDebugStats {
498        let owned_scopes = self.owned_scopes.borrow();
499        let fast_scopes = self.fast_scopes.borrow();
500        let owned_scope_len = owned_scopes.values().map(SmallVec::len).sum::<usize>();
501        let owned_scope_cap =
502            owned_scopes.capacity() + owned_scopes.values().map(SmallVec::capacity).sum::<usize>();
503        let scopes_len = owned_scope_len + fast_scopes.len();
504        let scopes_cap = owned_scope_cap + fast_scopes.capacity();
505        let mut observed_state_count = 0;
506        let mut observed_state_capacity = 0;
507        let mut stateless_scope_count = 0;
508
509        for entry in owned_scopes
510            .values()
511            .flat_map(|bucket| bucket.iter())
512            .chain(fast_scopes.values())
513        {
514            let entry = entry.borrow();
515            observed_state_count += entry.observed.len();
516            observed_state_capacity += entry.observed.capacity();
517            stateless_scope_count += usize::from(entry.observed.is_empty());
518        }
519
520        SnapshotStateObserverDebugStats {
521            scopes_len,
522            scopes_cap,
523            fast_scopes_len: fast_scopes.len(),
524            fast_scopes_cap: fast_scopes.capacity(),
525            stateless_scope_count,
526            observed_state_count,
527            observed_state_capacity,
528        }
529    }
530
531    fn run_with_read_observer<R>(&self, block: impl FnOnce() -> R) -> R {
532        use crate::snapshot_v2::take_transparent_observer_mutable_snapshot_reusing;
533
534        let mut snapshot = take_transparent_observer_mutable_snapshot_reusing(
535            Some(self.read_dispatcher.clone()),
536            None,
537            self.read_snapshot.take(),
538        );
539        let result = snapshot.enter(block);
540        snapshot.dispose();
541        if Arc::get_mut(&mut snapshot).is_some() && !snapshot.has_pending_changes() {
542            self.read_snapshot.replace(Some(snapshot));
543        }
544        result
545    }
546
547    fn handle_apply(&self, modified: &[Arc<dyn StateObject>]) {
548        if modified.is_empty() {
549            return;
550        }
551
552        let mut seen_scope_ids: HashSet<usize> = HashSet::default();
553        let mut to_notify: Vec<Rc<RefCell<ScopeEntry>>> = Vec::new();
554        {
555            let observed_to_scopes = self.observed_to_scopes.borrow();
556            let indexed_scopes = self.indexed_scopes.borrow();
557            for state in modified {
558                if let Some(scope_ids) = observed_to_scopes.get(&state.object_id().as_usize()) {
559                    let mut ordered_scope_ids: SmallVec<[usize; 8]> =
560                        scope_ids.iter().copied().collect();
561                    ordered_scope_ids.sort_unstable();
562                    for scope_id in ordered_scope_ids {
563                        if seen_scope_ids.insert(scope_id)
564                            && let Some(entry) = indexed_scopes.get(&scope_id)
565                        {
566                            to_notify.push(entry.clone());
567                        }
568                    }
569                }
570            }
571        }
572
573        if to_notify.is_empty() {
574            return;
575        }
576
577        for entry in to_notify {
578            let executor = self.executor.clone();
579            executor(Box::new(move || {
580                if let Ok(entry) = entry.try_borrow() {
581                    entry.notify();
582                }
583            }));
584        }
585    }
586
587    fn replace_observed_ids(&self, entry: &Rc<RefCell<ScopeEntry>>, observed: ObservedIds) {
588        let (entry_id, previous) = {
589            let mut entry_mut = entry.borrow_mut();
590            let entry_id = entry_mut.id;
591            let previous = std::mem::replace(&mut entry_mut.observed, observed);
592            (entry_id, previous)
593        };
594        let entry_ref = entry.borrow();
595        if previous.iter().eq(entry_ref.observed.iter()) {
596            return;
597        }
598        self.unregister_observed_ids(entry_id, &previous);
599        self.register_observed_ids(entry_id, &entry_ref.observed);
600    }
601
602    fn register_observed_ids(&self, entry_id: usize, observed: &ObservedIds) {
603        let mut observed_to_scopes = self.observed_to_scopes.borrow_mut();
604        for state_id in observed.iter() {
605            let scope_ids = observed_to_scopes.entry(state_id).or_default();
606            scope_ids.insert(entry_id);
607        }
608    }
609
610    fn unregister_observed_ids(&self, entry_id: usize, observed: &ObservedIds) {
611        let mut observed_to_scopes = self.observed_to_scopes.borrow_mut();
612        let mut emptied = SmallVec::<[StateObjectId; MAX_OBSERVED_STATES]>::new();
613        for state_id in observed.iter() {
614            if let Some(scope_ids) = observed_to_scopes.get_mut(&state_id) {
615                scope_ids.remove(&entry_id);
616                if scope_ids.is_empty() {
617                    emptied.push(state_id);
618                }
619            }
620        }
621        for state_id in emptied {
622            observed_to_scopes.remove(&state_id);
623        }
624        shrink_map_if_sparse(&mut observed_to_scopes, Self::MIN_RETAINED_SCOPE_CAPACITY);
625    }
626
627    fn unregister_entry(&self, entry: &Rc<RefCell<ScopeEntry>>) {
628        let (entry_id, observed) = {
629            let mut entry_mut = entry.borrow_mut();
630            let observed = std::mem::replace(&mut entry_mut.observed, ObservedIds::new());
631            (entry_mut.id, observed)
632        };
633        self.unregister_observed_ids(entry_id, &observed);
634        self.indexed_scopes.borrow_mut().remove(&entry_id);
635    }
636}
637
638fn shrink_map_if_sparse<K, V>(map: &mut HashMap<K, V>, min_retained_capacity: usize)
639where
640    K: Eq + std::hash::Hash,
641{
642    if map.capacity() <= map.len().max(min_retained_capacity).saturating_mul(4) {
643        return;
644    }
645
646    let retained = map.len().max(min_retained_capacity);
647    let mut rebuilt = HashMap::default();
648    rebuilt.reserve(retained);
649    rebuilt.extend(map.drain());
650    *map = rebuilt;
651}
652
653#[derive(Default)]
654struct ReadObservationStack {
655    targets: Vec<Rc<RefCell<ObservedIds>>>,
656    depth: usize,
657}
658
659impl ReadObservationStack {
660    fn push(&mut self) -> Rc<RefCell<ObservedIds>> {
661        if self.depth == self.targets.len() {
662            self.targets.push(Rc::new(RefCell::new(ObservedIds::new())));
663        }
664        let target = Rc::clone(&self.targets[self.depth]);
665        self.depth += 1;
666        target
667    }
668
669    fn last(&self) -> Option<&Rc<RefCell<ObservedIds>>> {
670        self.depth.checked_sub(1).map(|index| &self.targets[index])
671    }
672
673    fn pop(&mut self) -> Rc<RefCell<ObservedIds>> {
674        self.depth -= 1;
675        Rc::clone(&self.targets[self.depth])
676    }
677}
678
679enum ObservedIds {
680    Small(SmallVec<[ObservedState; MAX_OBSERVED_STATES]>),
681    Large(HashMap<StateObjectId, Option<Rc<dyn Any>>>),
682}
683
684struct ObservedState {
685    id: StateObjectId,
686    _lease: Option<Rc<dyn Any>>,
687}
688
689impl ObservedIds {
690    fn new() -> Self {
691        ObservedIds::Small(SmallVec::new())
692    }
693
694    fn insert(&mut self, state: &dyn StateObject) {
695        let id = state.object_id().as_usize();
696        match self {
697            ObservedIds::Small(small) => {
698                if small.iter().any(|observed| observed.id == id) {
699                    return;
700                }
701                if small.len() < MAX_OBSERVED_STATES {
702                    small.push(ObservedState {
703                        id,
704                        _lease: state.observation_lease(),
705                    });
706                } else {
707                    let mut large =
708                        HashMap::with_capacity_and_hasher(small.len() + 1, Default::default());
709                    for observed in small.drain(..) {
710                        large.insert(observed.id, observed._lease);
711                    }
712                    large.insert(id, state.observation_lease());
713                    *self = ObservedIds::Large(large);
714                }
715            }
716            ObservedIds::Large(large) => {
717                large.entry(id).or_insert_with(|| state.observation_lease());
718            }
719        }
720    }
721
722    fn is_empty(&self) -> bool {
723        match self {
724            ObservedIds::Small(small) => small.is_empty(),
725            ObservedIds::Large(large) => large.is_empty(),
726        }
727    }
728
729    fn len(&self) -> usize {
730        match self {
731            ObservedIds::Small(small) => small.len(),
732            ObservedIds::Large(large) => large.len(),
733        }
734    }
735
736    fn capacity(&self) -> usize {
737        match self {
738            ObservedIds::Small(small) => small.capacity(),
739            ObservedIds::Large(large) => large.capacity(),
740        }
741    }
742
743    fn iter(&self) -> impl Iterator<Item = StateObjectId> + '_ {
744        let (small, large) = match self {
745            ObservedIds::Small(small) => (Some(small.as_slice()), None),
746            ObservedIds::Large(large) => (None, Some(large)),
747        };
748        small
749            .into_iter()
750            .flatten()
751            .map(|observed| observed.id)
752            .chain(large.into_iter().flat_map(|states| states.keys().copied()))
753    }
754}
755
756const MAX_OBSERVED_STATES: usize = 8;
757
758enum ScopeStorage {
759    Owned(Box<dyn Any>),
760    RecomposeScope {
761        id: ScopeId,
762        weak: Weak<RecomposeScopeInner>,
763    },
764}
765
766struct ScopeEntry {
767    id: usize,
768    scope: ScopeStorage,
769    on_changed: Rc<dyn ScopeChangedCallback>,
770    observed: ObservedIds,
771    last_seen_version: u64,
772}
773
774impl ScopeEntry {
775    fn new<T>(id: usize, scope: T, on_changed: Rc<dyn ScopeChangedCallback>) -> Self
776    where
777        T: Any + 'static,
778    {
779        Self {
780            id,
781            scope: ScopeStorage::from_value(scope),
782            on_changed,
783            observed: ObservedIds::new(),
784            last_seen_version: u64::MAX,
785        }
786    }
787
788    fn callback_reusing<F: Fn(&dyn Any) + 'static>(
789        &mut self,
790        callback: F,
791    ) -> Rc<dyn ScopeChangedCallback> {
792        if let Some(stored) = Rc::get_mut(&mut self.on_changed)
793            .and_then(|stored| (stored as &mut dyn Any).downcast_mut::<F>())
794        {
795            *stored = callback;
796            Rc::clone(&self.on_changed)
797        } else {
798            Rc::new(callback)
799        }
800    }
801
802    fn update<T>(&mut self, new_scope: T, on_changed: Rc<dyn ScopeChangedCallback>)
803    where
804        T: Any + 'static,
805    {
806        self.update_scope(new_scope);
807        self.on_changed = on_changed;
808    }
809
810    fn update_scope<T>(&mut self, new_scope: T)
811    where
812        T: Any + 'static,
813    {
814        if let ScopeStorage::Owned(stored) = &mut self.scope
815            && let Some(stored) = stored.downcast_mut::<T>()
816        {
817            *stored = new_scope;
818        } else {
819            self.scope = ScopeStorage::from_value(new_scope);
820        }
821    }
822
823    fn matches_scope<T>(&self, scope: &T) -> bool
824    where
825        T: Any + Eq + 'static,
826    {
827        if let Some(scope) = (scope as &dyn Any).downcast_ref::<RecomposeScope>() {
828            return matches!(
829                &self.scope,
830                ScopeStorage::RecomposeScope { id, .. } if *id == scope.id()
831            );
832        }
833
834        match &self.scope {
835            ScopeStorage::Owned(stored) => stored
836                .downcast_ref::<T>()
837                .is_some_and(|stored| stored == scope),
838            ScopeStorage::RecomposeScope { .. } => false,
839        }
840    }
841
842    fn matches_predicate(&self, predicate: &impl Fn(&dyn Any) -> bool) -> bool {
843        match &self.scope {
844            ScopeStorage::Owned(scope) => predicate(scope.as_ref()),
845            ScopeStorage::RecomposeScope { weak, .. } => weak
846                .upgrade()
847                .is_none_or(|inner| predicate(&RecomposeScope { inner })),
848        }
849    }
850
851    fn should_retain(&self) -> bool {
852        match &self.scope {
853            ScopeStorage::Owned(_) => true,
854            ScopeStorage::RecomposeScope { weak, .. } => weak.upgrade().is_some(),
855        }
856    }
857
858    fn notify(&self) {
859        match &self.scope {
860            ScopeStorage::Owned(scope) => (self.on_changed)(scope.as_ref()),
861            ScopeStorage::RecomposeScope { weak, .. } => {
862                if let Some(inner) = weak.upgrade() {
863                    (self.on_changed)(&RecomposeScope { inner });
864                }
865            }
866        }
867    }
868}
869
870impl ScopeStorage {
871    fn from_value<T>(value: T) -> Self
872    where
873        T: Any + 'static,
874    {
875        let any = &value as &dyn Any;
876        if let Some(scope) = any.downcast_ref::<RecomposeScope>() {
877            Self::RecomposeScope {
878                id: scope.id(),
879                weak: scope.downgrade(),
880            }
881        } else {
882            Self::Owned(Box::new(value))
883        }
884    }
885}
886
887#[cfg(test)]
888mod tests {
889    use std::cell::Cell;
890
891    use super::*;
892    use crate::{
893        snapshot_v2::{TestRuntimeGuard, reset_runtime_for_tests, take_mutable_snapshot},
894        state::{NeverEqual, SnapshotMutableState},
895    };
896
897    fn reset_runtime() -> TestRuntimeGuard {
898        reset_runtime_for_tests()
899    }
900
901    #[derive(Clone, Eq, Hash, PartialEq)]
902    struct TestScope(&'static str);
903
904    #[test]
905    fn scope_update_reuses_storage_and_replaces_payload_and_callback() {
906        let first = Rc::new(String::from("first"));
907        let second = Rc::new(String::from("second"));
908        let delivered = Rc::new(RefCell::new(Vec::new()));
909        let mut entry = ScopeEntry::new(0, first.clone(), Rc::new(|_| panic!("stale callback")));
910        let ScopeStorage::Owned(stored) = &entry.scope else {
911            panic!("expected owned scope");
912        };
913        let address = stored.downcast_ref::<Rc<String>>().unwrap() as *const Rc<String>;
914        let received = delivered.clone();
915        entry.update(
916            second.clone(),
917            Rc::new(move |scope| {
918                received.borrow_mut().push(
919                    scope
920                        .downcast_ref::<Rc<String>>()
921                        .unwrap()
922                        .as_str()
923                        .to_owned(),
924                );
925            }),
926        );
927        assert_eq!(Rc::strong_count(&first), 1);
928        assert_eq!(Rc::strong_count(&second), 2);
929        entry.notify();
930        assert_eq!(*delivered.borrow(), vec!["second"]);
931        let ScopeStorage::Owned(stored) = &entry.scope else {
932            panic!("expected owned scope");
933        };
934        assert_eq!(
935            stored.downcast_ref::<Rc<String>>().unwrap() as *const Rc<String>,
936            address
937        );
938        drop(entry);
939        assert_eq!(Rc::strong_count(&second), 1);
940    }
941
942    #[test]
943    fn reobservation_refreshes_captures_and_preserves_shared_callbacks() {
944        let _guard = reset_runtime();
945        let state = SnapshotMutableState::new_in_arc(0, Arc::new(NeverEqual));
946        let delivered = Rc::new(RefCell::new(Vec::new()));
947        let callback = |generation| {
948            let delivered = delivered.clone();
949            move |scope: &TestScope| delivered.borrow_mut().push((generation, scope.0))
950        };
951        let scope = TestScope("callback");
952        let observer = SnapshotStateObserver::new(|callback| callback());
953        let read = || {
954            let _ = state.get();
955        };
956        observer.observe_reads(scope.clone(), callback(1), read);
957        let entry = observer.inner.find_scope_entry(&scope).unwrap();
958        let held = entry.borrow().on_changed.clone();
959        observer.observe_reads(scope.clone(), callback(2), read);
960        held(&scope);
961        entry.borrow().notify();
962        drop(held);
963
964        let allocation = Rc::as_ptr(&entry.borrow().on_changed);
965        observer.observe_reads(scope.clone(), callback(3), read);
966        assert!(std::ptr::addr_eq(
967            allocation,
968            Rc::as_ptr(&entry.borrow().on_changed)
969        ));
970        entry.borrow().notify();
971
972        let received = delivered.clone();
973        observer.observe_reads(
974            scope,
975            move |scope| received.borrow_mut().push((4, scope.0)),
976            read,
977        );
978        entry.borrow().notify();
979        assert_eq!(
980            *delivered.borrow(),
981            [
982                (1, "callback"),
983                (2, "callback"),
984                (3, "callback"),
985                (4, "callback")
986            ]
987        );
988    }
989
990    #[test]
991    fn reobservation_across_storage_thresholds_replaces_dependencies_and_callbacks() {
992        let _guard = reset_runtime();
993        let states: Vec<_> = (0..MAX_OBSERVED_STATES + 2)
994            .map(|_| SnapshotMutableState::new_in_arc(0, Arc::new(NeverEqual)))
995            .collect();
996        let notifications = Rc::new(RefCell::new(Vec::new()));
997        let observer = SnapshotStateObserver::new(|callback| callback());
998        observer.start();
999        for (generation, count) in [MAX_OBSERVED_STATES, MAX_OBSERVED_STATES + 1, 2, 0]
1000            .into_iter()
1001            .enumerate()
1002        {
1003            observer.begin_frame();
1004            let recorded = notifications.clone();
1005            observer.observe_reads(
1006                TestScope("changing"),
1007                move |scope| {
1008                    assert_eq!(scope.0, "changing");
1009                    recorded.borrow_mut().push(generation);
1010                },
1011                || {
1012                    for state in states.iter().take(count) {
1013                        let _ = state.get();
1014                        let _ = state.get();
1015                    }
1016                },
1017            );
1018            notifications.borrow_mut().clear();
1019            for (index, state) in states.iter().enumerate() {
1020                let snapshot = take_mutable_snapshot(None, None);
1021                snapshot.enter(|| state.set(generation as i32));
1022                snapshot.apply().check();
1023                let expected = (index + 1).min(count);
1024                assert_eq!(
1025                    *notifications.borrow(),
1026                    vec![generation; expected],
1027                    "count={count}, changed state={index}"
1028                );
1029            }
1030        }
1031    }
1032
1033    #[test]
1034    fn reobservation_preserves_notifications_when_dependencies_repeat_or_change() {
1035        let _guard = reset_runtime();
1036        let states: Vec<_> = (0..3)
1037            .map(|_| SnapshotMutableState::new_in_arc(0, Arc::new(NeverEqual)))
1038            .collect();
1039        let notifications = Rc::new(RefCell::new(Vec::new()));
1040        let observer = SnapshotStateObserver::new(|callback| callback());
1041        observer.start();
1042        for (generation, indices) in [[0, 1], [0, 1], [1, 0], [1, 2], [1, 2]]
1043            .into_iter()
1044            .enumerate()
1045        {
1046            observer.begin_frame();
1047            let received = notifications.clone();
1048            observer.observe_reads(
1049                TestScope("repeated"),
1050                move |_| received.borrow_mut().push(generation),
1051                || {
1052                    for index in indices {
1053                        let _ = states[index].get();
1054                    }
1055                },
1056            );
1057            for (index, state) in states.iter().enumerate() {
1058                notifications.borrow_mut().clear();
1059                let snapshot = take_mutable_snapshot(None, None);
1060                snapshot.enter(|| state.set(generation as i32));
1061                snapshot.apply().check();
1062                assert_eq!(
1063                    *notifications.borrow(),
1064                    if indices.contains(&index) {
1065                        vec![generation]
1066                    } else {
1067                        vec![]
1068                    },
1069                    "generation={generation}, state={index}"
1070                );
1071            }
1072        }
1073        observer.clear(&TestScope("repeated"));
1074        notifications.borrow_mut().clear();
1075        for state in states {
1076            observer.notify_changes(&[state]);
1077        }
1078        assert!(notifications.borrow().is_empty());
1079    }
1080
1081    #[test]
1082    fn stateless_scope_can_start_observing_and_replace_its_callback_before_the_block() {
1083        let _guard = reset_runtime();
1084        let observer = SnapshotStateObserver::new(|callback| callback());
1085        let state = SnapshotMutableState::new_in_arc(0, Arc::new(NeverEqual));
1086        let observed: Arc<dyn StateObject> = state.clone();
1087        let notifications = Rc::new(RefCell::new(Vec::new()));
1088        let discarded = notifications.clone();
1089        observer.observe_reads(
1090            TestScope("changing"),
1091            move |_| discarded.borrow_mut().push(0),
1092            || {},
1093        );
1094        assert_eq!(Rc::strong_count(&notifications), 1);
1095        assert_eq!(observer.debug_stats().scopes_len, 0);
1096        for generation in 1..=2 {
1097            observer.begin_frame();
1098            let delivered = notifications.clone();
1099            observer.observe_reads(
1100                TestScope("changing"),
1101                move |_| delivered.borrow_mut().push(generation),
1102                || {
1103                    observer.notify_changes(std::slice::from_ref(&observed));
1104                    let _ = state.get();
1105                },
1106            );
1107            observer.notify_changes(std::slice::from_ref(&observed));
1108        }
1109        assert_eq!(*notifications.borrow(), vec![1, 2, 2]);
1110        observer.clear(&TestScope("changing"));
1111        assert_eq!(Rc::strong_count(&notifications), 1);
1112    }
1113
1114    #[test]
1115    fn callback_captures_are_released_on_replacement_and_clear() {
1116        let _guard = reset_runtime();
1117        let state = SnapshotMutableState::new_in_arc(0, Arc::new(NeverEqual));
1118        let observer = SnapshotStateObserver::new(|callback| callback());
1119        let owners = [Rc::new(Cell::new(0)), Rc::new(Cell::new(0))];
1120        for owner in &owners {
1121            let captured = owner.clone();
1122            observer.observe_reads(
1123                TestScope("owner"),
1124                move |_| captured.set(captured.get() + 1),
1125                || {
1126                    let _ = state.get();
1127                },
1128            );
1129            assert_eq!(Rc::strong_count(owner), 2);
1130        }
1131        assert_eq!(Rc::strong_count(&owners[0]), 1);
1132        observer.clear(&TestScope("owner"));
1133        assert_eq!(Rc::strong_count(&owners[1]), 1);
1134    }
1135
1136    #[test]
1137    fn notifies_scope_when_state_changes() {
1138        let _guard = reset_runtime();
1139
1140        let state = SnapshotMutableState::new_in_arc(0, Arc::new(NeverEqual));
1141        let triggered = Rc::new(Cell::new(0));
1142        let observer_trigger = triggered.clone();
1143
1144        let observer = SnapshotStateObserver::new(|callback| callback());
1145        observer.start();
1146
1147        let scope = TestScope("scope");
1148        observer.observe_reads(
1149            scope,
1150            move |_| {
1151                observer_trigger.set(observer_trigger.get() + 1);
1152            },
1153            || {
1154                let _ = state.get();
1155            },
1156        );
1157
1158        let snapshot = take_mutable_snapshot(None, None);
1159        snapshot.enter(|| {
1160            state.set(1);
1161        });
1162        snapshot.apply().check();
1163
1164        assert_eq!(triggered.get(), 1);
1165        observer.stop();
1166    }
1167
1168    #[test]
1169    fn clear_removes_scope_observation() {
1170        let _guard = reset_runtime();
1171
1172        let state = SnapshotMutableState::new_in_arc(0, Arc::new(NeverEqual));
1173        let triggered = Rc::new(Cell::new(0));
1174        let observer_trigger = triggered.clone();
1175
1176        let observer = SnapshotStateObserver::new(|callback| callback());
1177        observer.start();
1178
1179        let scope = TestScope("scope");
1180        observer.observe_reads(
1181            scope.clone(),
1182            move |_| {
1183                observer_trigger.set(observer_trigger.get() + 1);
1184            },
1185            || {
1186                let _ = state.get();
1187            },
1188        );
1189
1190        observer.clear(&scope);
1191
1192        let snapshot = take_mutable_snapshot(None, None);
1193        snapshot.enter(|| {
1194            state.set(1);
1195        });
1196        snapshot.apply().check();
1197
1198        assert_eq!(triggered.get(), 0);
1199        observer.stop();
1200    }
1201
1202    #[test]
1203    fn repeated_owned_scope_observations_reuse_the_same_entry() {
1204        let _guard = reset_runtime();
1205
1206        let state = SnapshotMutableState::new_in_arc(0, Arc::new(NeverEqual));
1207        let observer = SnapshotStateObserver::new(|callback| callback());
1208        let scope = TestScope("scope");
1209
1210        observer.observe_reads(
1211            scope.clone(),
1212            |_| {},
1213            || {
1214                let _ = state.get();
1215            },
1216        );
1217        observer.observe_reads(
1218            scope,
1219            |_| {},
1220            || {
1221                let _ = state.get();
1222            },
1223        );
1224
1225        let stats = observer.debug_stats();
1226        assert_eq!(stats.scopes_len, 1);
1227        assert_eq!(stats.fast_scopes_len, 0);
1228    }
1229
1230    #[test]
1231    fn with_no_observations_skips_reads() {
1232        let _guard = reset_runtime();
1233
1234        let state = SnapshotMutableState::new_in_arc(0, Arc::new(NeverEqual));
1235        let triggered = Rc::new(Cell::new(0));
1236        let observer_trigger = triggered.clone();
1237
1238        let observer = SnapshotStateObserver::new(|callback| callback());
1239        observer.start();
1240
1241        let scope = TestScope("scope");
1242        observer.observe_reads(
1243            scope,
1244            move |_| {
1245                observer_trigger.set(observer_trigger.get() + 1);
1246            },
1247            || {
1248                observer.with_no_observations(|| {
1249                    let _ = state.get();
1250                });
1251            },
1252        );
1253
1254        let snapshot = take_mutable_snapshot(None, None);
1255        snapshot.enter(|| {
1256            state.set(1);
1257        });
1258        snapshot.apply().check();
1259
1260        assert_eq!(triggered.get(), 0);
1261        observer.stop();
1262    }
1263
1264    #[test]
1265    fn recycled_observation_refreshes_snapshot_state() {
1266        let _guard = reset_runtime();
1267        let state = SnapshotMutableState::new_in_arc(0, Arc::new(NeverEqual));
1268        let observer = SnapshotStateObserver::new(|callback| callback());
1269        let mut allocation = None;
1270        for value in 1..=3 {
1271            let parent = take_mutable_snapshot(None, None);
1272            parent.enter(|| {
1273                state.set(value);
1274                let expected = crate::snapshot_v2::current_snapshot().unwrap();
1275                observer.inner.run_with_read_observer(|| {
1276                    let crate::snapshot_v2::AnySnapshot::TransparentMutable(current) =
1277                        crate::snapshot_v2::current_snapshot().unwrap()
1278                    else {
1279                        panic!("expected an observation snapshot");
1280                    };
1281                    assert_eq!(current.snapshot_id(), expected.snapshot_id());
1282                    assert_eq!(current.invalid(), expected.invalid());
1283                    assert!(!current.is_disposed());
1284                    assert!(!current.has_pending_changes());
1285                    assert_eq!(state.get(), value);
1286                    let address = Arc::as_ptr(&current) as usize;
1287                    assert_eq!(*allocation.get_or_insert(address), address);
1288                });
1289            });
1290            parent.apply().check();
1291        }
1292    }
1293
1294    #[test]
1295    fn recycled_observation_preserves_escaped_snapshots() {
1296        let _guard = reset_runtime();
1297        let observer = SnapshotStateObserver::new(|callback| callback());
1298        let escaped = observer
1299            .inner
1300            .run_with_read_observer(|| crate::snapshot_v2::current_snapshot().unwrap());
1301        let id = escaped.snapshot_id();
1302        observer.inner.run_with_read_observer(|| {
1303            let current = crate::snapshot_v2::current_snapshot().unwrap();
1304            let crate::snapshot_v2::AnySnapshot::TransparentMutable(escaped) = &escaped else {
1305                panic!("expected an observation snapshot");
1306            };
1307            assert!(!current.is_same_transparent(escaped));
1308            assert_eq!(escaped.snapshot_id(), id);
1309            assert!(escaped.is_disposed());
1310        });
1311        let weak = observer.inner.run_with_read_observer(|| {
1312            let crate::snapshot_v2::AnySnapshot::TransparentMutable(current) =
1313                crate::snapshot_v2::current_snapshot().unwrap()
1314            else {
1315                panic!("expected an observation snapshot");
1316            };
1317            Arc::downgrade(&current)
1318        });
1319        assert!(weak.upgrade().is_none());
1320        observer.inner.run_with_read_observer(|| {
1321            assert!(weak.upgrade().is_none());
1322        });
1323    }
1324
1325    #[test]
1326    fn recycled_observation_does_not_retain_written_state() {
1327        let _guard = reset_runtime();
1328        let observer = SnapshotStateObserver::new(|callback| callback());
1329        let state = SnapshotMutableState::new_in_arc(0, Arc::new(NeverEqual));
1330        let owners = Arc::strong_count(&state);
1331        observer.inner.run_with_read_observer(|| {
1332            let current = crate::snapshot_v2::current_snapshot().unwrap();
1333            current.record_write(state.clone());
1334        });
1335        assert_eq!(Arc::strong_count(&state), owners);
1336    }
1337
1338    #[test]
1339    fn nested_observe_reads_attributes_state_to_innermost_scope_only() {
1340        let _guard = reset_runtime();
1341
1342        let state = SnapshotMutableState::new_in_arc(0, Arc::new(NeverEqual));
1343        let outer_state = SnapshotMutableState::new_in_arc(0, Arc::new(NeverEqual));
1344        let outer_triggered = Rc::new(Cell::new(0));
1345        let inner_triggered = Rc::new(Cell::new(0));
1346
1347        let observer = SnapshotStateObserver::new(|callback| callback());
1348        observer.start();
1349
1350        let outer_scope = TestScope("outer");
1351        let inner_scope = TestScope("inner");
1352        observer.observe_reads(
1353            outer_scope,
1354            {
1355                let outer_triggered = Rc::clone(&outer_triggered);
1356                move |_| outer_triggered.set(outer_triggered.get() + 1)
1357            },
1358            || {
1359                let _ = outer_state.get();
1360                observer.observe_reads(
1361                    inner_scope.clone(),
1362                    {
1363                        let inner_triggered = Rc::clone(&inner_triggered);
1364                        move |_| inner_triggered.set(inner_triggered.get() + 1)
1365                    },
1366                    || {
1367                        let _ = state.get();
1368                    },
1369                );
1370            },
1371        );
1372
1373        let snapshot = take_mutable_snapshot(None, None);
1374        snapshot.enter(|| {
1375            state.set(1);
1376        });
1377        snapshot.apply().check();
1378
1379        assert_eq!(outer_triggered.get(), 0);
1380        assert_eq!(inner_triggered.get(), 1);
1381        let snapshot = take_mutable_snapshot(None, None);
1382        snapshot.enter(|| outer_state.set(1));
1383        snapshot.apply().check();
1384        assert_eq!(outer_triggered.get(), 1);
1385        assert_eq!(inner_triggered.get(), 1);
1386        observer.stop();
1387    }
1388
1389    #[test]
1390    fn unwound_observation_does_not_leak_reads_into_reused_storage() {
1391        let _guard = reset_runtime();
1392        let abandoned = SnapshotMutableState::new_in_arc(0, Arc::new(NeverEqual));
1393        let live = SnapshotMutableState::new_in_arc(0, Arc::new(NeverEqual));
1394        let triggered = Rc::new(Cell::new(0));
1395        let observer = SnapshotStateObserver::new(|callback| callback());
1396        observer.start();
1397
1398        let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
1399            observer.observe_reads(
1400                TestScope("abandoned"),
1401                |_| {},
1402                || {
1403                    let _ = abandoned.get();
1404                    panic!("abandon observation");
1405                },
1406            );
1407        }));
1408        assert!(result.is_err());
1409        observer.observe_reads(
1410            TestScope("live"),
1411            {
1412                let triggered = Rc::clone(&triggered);
1413                move |_| triggered.set(triggered.get() + 1)
1414            },
1415            || {
1416                let _ = live.get();
1417            },
1418        );
1419
1420        let snapshot = take_mutable_snapshot(None, None);
1421        snapshot.enter(|| abandoned.set(1));
1422        snapshot.apply().check();
1423        assert_eq!(triggered.get(), 0);
1424        let snapshot = take_mutable_snapshot(None, None);
1425        snapshot.enter(|| live.set(1));
1426        snapshot.apply().check();
1427        assert_eq!(triggered.get(), 1);
1428        observer.stop();
1429    }
1430
1431    #[test]
1432    fn clearing_one_scope_keeps_shared_state_registered_for_other_scope() {
1433        let _guard = reset_runtime();
1434
1435        let state = SnapshotMutableState::new_in_arc(0, Arc::new(NeverEqual));
1436        let first_triggered = Rc::new(Cell::new(0));
1437        let second_triggered = Rc::new(Cell::new(0));
1438
1439        let observer = SnapshotStateObserver::new(|callback| callback());
1440        observer.start();
1441
1442        let first_scope = TestScope("first");
1443        let second_scope = TestScope("second");
1444        observer.observe_reads(
1445            first_scope.clone(),
1446            {
1447                let first_triggered = Rc::clone(&first_triggered);
1448                move |_| first_triggered.set(first_triggered.get() + 1)
1449            },
1450            || {
1451                let _ = state.get();
1452            },
1453        );
1454        observer.observe_reads(
1455            second_scope,
1456            {
1457                let second_triggered = Rc::clone(&second_triggered);
1458                move |_| second_triggered.set(second_triggered.get() + 1)
1459            },
1460            || {
1461                let _ = state.get();
1462            },
1463        );
1464
1465        observer.clear(&first_scope);
1466
1467        let snapshot = take_mutable_snapshot(None, None);
1468        snapshot.enter(|| {
1469            state.set(1);
1470        });
1471        snapshot.apply().check();
1472
1473        assert_eq!(first_triggered.get(), 0);
1474        assert_eq!(second_triggered.get(), 1);
1475        observer.stop();
1476    }
1477
1478    #[test]
1479    fn shared_state_notifies_scopes_in_registration_order() {
1480        let _guard = reset_runtime();
1481
1482        let state = SnapshotMutableState::new_in_arc(0, Arc::new(NeverEqual));
1483        let notifications = Rc::new(RefCell::new(Vec::new()));
1484
1485        let observer = SnapshotStateObserver::new(|callback| callback());
1486        observer.start();
1487
1488        observer.observe_reads(
1489            TestScope("first"),
1490            {
1491                let notifications = Rc::clone(&notifications);
1492                move |_| notifications.borrow_mut().push("first")
1493            },
1494            || {
1495                let _ = state.get();
1496            },
1497        );
1498        observer.observe_reads(
1499            TestScope("second"),
1500            {
1501                let notifications = Rc::clone(&notifications);
1502                move |_| notifications.borrow_mut().push("second")
1503            },
1504            || {
1505                let _ = state.get();
1506            },
1507        );
1508
1509        let snapshot = take_mutable_snapshot(None, None);
1510        snapshot.enter(|| {
1511            state.set(1);
1512        });
1513        snapshot.apply().check();
1514
1515        assert_eq!(notifications.borrow().as_slice(), &["first", "second"]);
1516        observer.stop();
1517    }
1518
1519    #[test]
1520    fn stateless_recompose_scope_does_not_retain_observer_entry() {
1521        let _guard = reset_runtime();
1522
1523        let observer = SnapshotStateObserver::new(|callback| callback());
1524        let runtime = crate::TestRuntime::new();
1525        let scope = RecomposeScope::new_for_test(runtime.handle());
1526
1527        observer.observe_reads(scope, |_| {}, || {});
1528
1529        let stats = observer.debug_stats();
1530        assert_eq!(stats.scopes_len, 0);
1531        assert_eq!(stats.fast_scopes_len, 0);
1532        assert_eq!(stats.stateless_scope_count, 0);
1533    }
1534
1535    #[test]
1536    fn scope_that_stops_reading_state_is_removed_immediately() {
1537        let _guard = reset_runtime();
1538
1539        let state = SnapshotMutableState::new_in_arc(0, Arc::new(NeverEqual));
1540        let observer = SnapshotStateObserver::new(|callback| callback());
1541        let runtime = crate::TestRuntime::new();
1542        let scope = RecomposeScope::new_for_test(runtime.handle());
1543        let triggered = Rc::new(Cell::new(0));
1544        let observer_trigger = Rc::clone(&triggered);
1545
1546        observer.observe_reads(
1547            scope.clone(),
1548            move |_| observer_trigger.set(observer_trigger.get() + 1),
1549            || {
1550                let _ = state.get();
1551            },
1552        );
1553
1554        let after_stateful = observer.debug_stats();
1555        assert_eq!(after_stateful.scopes_len, 1);
1556        assert_eq!(after_stateful.fast_scopes_len, 1);
1557
1558        observer.observe_reads(scope, |_| {}, || {});
1559
1560        let after_stateless = observer.debug_stats();
1561        assert_eq!(after_stateless.scopes_len, 0);
1562        assert_eq!(after_stateless.fast_scopes_len, 0);
1563
1564        let snapshot = take_mutable_snapshot(None, None);
1565        snapshot.enter(|| {
1566            state.set(1);
1567        });
1568        snapshot.apply().check();
1569
1570        assert_eq!(triggered.get(), 0);
1571    }
1572
1573    #[test]
1574    fn begin_frame_prunes_dropped_recompose_scope_entries() {
1575        let _guard = reset_runtime();
1576
1577        let state = SnapshotMutableState::new_in_arc(0, Arc::new(NeverEqual));
1578        let observer = SnapshotStateObserver::new(|callback| callback());
1579        let runtime = crate::TestRuntime::new();
1580        let scope = RecomposeScope::new_for_test(runtime.handle());
1581
1582        observer.observe_reads(
1583            scope.clone(),
1584            |_| {},
1585            || {
1586                let _ = state.get();
1587            },
1588        );
1589
1590        let before_prune = observer.debug_stats();
1591        assert_eq!(before_prune.scopes_len, 1);
1592        assert_eq!(before_prune.fast_scopes_len, 1);
1593
1594        drop(scope);
1595        observer.begin_frame();
1596
1597        let after_prune = observer.debug_stats();
1598        assert_eq!(after_prune.scopes_len, 0);
1599        assert_eq!(after_prune.fast_scopes_len, 0);
1600    }
1601}