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)]
888#[path = "tests/snapshot_state_observer_tests.rs"]
889mod tests;
cranpose_core/snapshot_state_observer.rs

cranpose_core/
snapshot_state_observer.rs
