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    /// One `Rc` per type of callback that captures nothing: see
150    /// [`SnapshotStateObserverInner::capture_free_callback`].
151    capture_free_callbacks: RefCell<CaptureFreeCallbacks>,
152}
153
154#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
155struct OwnedScopeIndexKey {
156    type_id: TypeId,
157    value_hash: u64,
158}
159
160type OwnedScopeBucket = SmallVec<[Rc<RefCell<ScopeEntry>>; 1]>;
161
162/// The shared `Rc` of each type of callback that captures nothing.
163type CaptureFreeCallbacks = SmallVec<[(TypeId, Rc<dyn ScopeChangedCallback>); 2]>;
164
165fn owned_scope_index_key<T>(scope: &T) -> OwnedScopeIndexKey
166where
167    T: Any + Hash + 'static,
168{
169    let mut hasher = default_hash::new();
170    scope.hash(&mut hasher);
171    OwnedScopeIndexKey {
172        type_id: TypeId::of::<T>(),
173        value_hash: hasher.finish(),
174    }
175}
176
177impl SnapshotStateObserverInner {
178    const MIN_RETAINED_SCOPE_CAPACITY: usize = 256;
179
180    fn new(on_changed_executor: impl Fn(Box<dyn FnOnce() + 'static>) + 'static) -> Self {
181        let pause_count = Rc::new(Cell::new(0));
182        let active_read_targets = Rc::new(RefCell::new(ReadObservationStack::default()));
183        let dispatcher_pause_count = Rc::clone(&pause_count);
184        let dispatcher_targets = Rc::clone(&active_read_targets);
185        let read_dispatcher: ReadObserver = Arc::new(move |state| {
186            if dispatcher_pause_count.get() > 0 {
187                return;
188            }
189            let observed = dispatcher_targets.borrow().last().cloned();
190            if let Some(observed) = observed {
191                observed.borrow_mut().insert(state);
192            }
193        });
194
195        Self {
196            executor: Rc::new(on_changed_executor),
197            owned_scopes: RefCell::new(HashMap::default()),
198            fast_scopes: RefCell::new(HashMap::default()),
199            indexed_scopes: RefCell::new(HashMap::default()),
200            observed_to_scopes: RefCell::new(HashMap::default()),
201            pause_count,
202            active_read_targets,
203            read_dispatcher,
204            read_snapshot: RefCell::new(None),
205            apply_handle: RefCell::new(None),
206            weak_self: RefCell::new(Weak::new()),
207            frame_version: Cell::new(0),
208            next_entry_id: Cell::new(0),
209            capture_free_callbacks: RefCell::new(SmallVec::new()),
210        }
211    }
212
213    fn set_self(&self, weak: Weak<SnapshotStateObserverInner>) {
214        self.weak_self.replace(weak);
215    }
216
217    fn begin_frame(&self) {
218        let next = self.frame_version.get().wrapping_add(1);
219        self.frame_version.set(next);
220        self.prune_dead_scopes();
221    }
222
223    fn observe_reads<T, R>(
224        &self,
225        scope: T,
226        on_value_changed_for_scope: impl Fn(&T) + 'static,
227        block: impl FnOnce() -> R,
228    ) -> R
229    where
230        T: Any + Clone + Eq + Hash + 'static,
231    {
232        let frame_version = self.frame_version.get();
233        let has_frame_version = frame_version != 0;
234
235        let existing_entry = self.find_scope_entry(&scope);
236        let on_changed = std::cell::LazyCell::new(|| {
237            let callback = move |scope_any: &dyn Any| {
238                if let Some(typed) = scope_any.downcast_ref::<T>() {
239                    on_value_changed_for_scope(typed);
240                }
241            };
242            if std::mem::size_of_val(&callback) == 0 {
243                return self.capture_free_callback(callback);
244            }
245            match existing_entry.as_ref() {
246                Some(entry) => entry.borrow_mut().callback_reusing(callback),
247                None => Rc::new(callback),
248            }
249        });
250
251        if let Some(entry) = existing_entry.as_ref() {
252            let already_observed = {
253                let mut entry_mut = entry.borrow_mut();
254                entry_mut.update_scope(scope.clone());
255                has_frame_version && entry_mut.last_seen_version == frame_version
256            };
257            let callback = on_changed.clone();
258            entry.borrow_mut().on_changed = callback;
259            if already_observed {
260                return block();
261            }
262        }
263
264        let observed = self.active_read_targets.borrow_mut().push();
265        struct ActiveObservationGuard {
266            stack: Rc<RefCell<ReadObservationStack>>,
267        }
268        impl Drop for ActiveObservationGuard {
269            fn drop(&mut self) {
270                let target = self.stack.borrow_mut().pop();
271                let discarded = target.replace(ObservedIds::new());
272                drop(discarded);
273            }
274        }
275        let _guard = ActiveObservationGuard {
276            stack: Rc::clone(&self.active_read_targets),
277        };
278
279        let result = self.run_with_read_observer(block);
280
281        if observed.borrow().is_empty() {
282            if existing_entry.is_some() {
283                self.clear(&scope);
284            }
285            return result;
286        }
287
288        let observed = {
289            let mut observed = observed.borrow_mut();
290            std::mem::replace(&mut *observed, ObservedIds::new())
291        };
292        let entry = existing_entry
293            .clone()
294            .unwrap_or_else(|| self.insert_scope_entry(scope.clone(), on_changed.clone()));
295        {
296            let mut entry_mut = entry.borrow_mut();
297            entry_mut.update(scope, Rc::clone(&on_changed));
298            entry_mut.last_seen_version = if has_frame_version {
299                frame_version
300            } else {
301                u64::MAX
302            };
303        }
304        self.replace_observed_ids(&entry, observed);
305
306        result
307    }
308
309    fn with_no_observations<R>(&self, block: impl FnOnce() -> R) -> R {
310        self.pause_count.set(self.pause_count.get() + 1);
311        let result = block();
312        self.pause_count
313            .set(self.pause_count.get().saturating_sub(1));
314        result
315    }
316
317    fn clear<T>(&self, scope: &T)
318    where
319        T: Any + Eq + Hash + 'static,
320    {
321        if let Some(rc_scope) = (scope as &dyn Any).downcast_ref::<RecomposeScope>() {
322            if let Some(entry) = self.fast_scopes.borrow_mut().remove(&rc_scope.id()) {
323                self.unregister_entry(&entry);
324            }
325            return;
326        }
327
328        let removed = self.remove_owned_scope_entry(scope);
329        if let Some(entry) = removed {
330            self.unregister_entry(&entry);
331        }
332    }
333
334    fn clear_if(&self, predicate: impl Fn(&dyn Any) -> bool) {
335        let removed_fast = {
336            let mut fast_scopes = self.fast_scopes.borrow_mut();
337            let removed_ids: Vec<_> = fast_scopes
338                .iter()
339                .filter(|(_, entry)| entry.borrow().matches_predicate(&predicate))
340                .map(|(scope_id, _)| *scope_id)
341                .collect();
342            removed_ids
343                .into_iter()
344                .filter_map(|scope_id| fast_scopes.remove(&scope_id))
345                .collect::<Vec<_>>()
346        };
347        let removed_owned =
348            { self.partition_owned_scopes(|entry| entry.matches_predicate(&predicate)) };
349
350        for entry in removed_fast.into_iter().chain(removed_owned) {
351            self.unregister_entry(&entry);
352        }
353    }
354
355    fn clear_all(&self) {
356        self.fast_scopes.borrow_mut().clear();
357        self.owned_scopes.borrow_mut().clear();
358        self.indexed_scopes.borrow_mut().clear();
359        self.observed_to_scopes.borrow_mut().clear();
360    }
361
362    fn start(&self, weak_self: Weak<SnapshotStateObserverInner>) {
363        if self.apply_handle.borrow().is_some() {
364            return;
365        }
366
367        let handle = register_apply_observer(Rc::new(move |modified, _snapshot_id| {
368            if let Some(inner) = weak_self.upgrade() {
369                inner.handle_apply(modified);
370            }
371        }));
372        self.apply_handle.replace(Some(handle));
373    }
374
375    fn stop(&self) {
376        if let Some(handle) = self.apply_handle.borrow_mut().take() {
377            drop(handle);
378        }
379    }
380
381    fn find_scope_entry<T>(&self, scope: &T) -> Option<Rc<RefCell<ScopeEntry>>>
382    where
383        T: Any + Eq + Hash + 'static,
384    {
385        if let Some(scope) = (scope as &dyn Any).downcast_ref::<RecomposeScope>() {
386            return self.fast_scopes.borrow().get(&scope.id()).cloned();
387        }
388
389        self.find_owned_scope_entry(scope)
390    }
391
392    /// The `Rc` of a callback that captures nothing. Every closure of such a
393    /// type does the same thing, as `RecomposeScope::invalidate` does for
394    /// every group, so one `Rc` serves all its scopes instead of one each.
395    fn capture_free_callback<F: Fn(&dyn Any) + 'static>(
396        &self,
397        callback: F,
398    ) -> Rc<dyn ScopeChangedCallback> {
399        let type_id = TypeId::of::<F>();
400        let mut shared = self.capture_free_callbacks.borrow_mut();
401        if let Some((_, callback)) = shared.iter().find(|(id, _)| *id == type_id) {
402            return Rc::clone(callback);
403        }
404        let callback: Rc<dyn ScopeChangedCallback> = Rc::new(callback);
405        shared.push((type_id, Rc::clone(&callback)));
406        callback
407    }
408
409    fn insert_scope_entry(
410        &self,
411        scope: impl Any + Clone + Eq + Hash + 'static,
412        on_changed: Rc<dyn ScopeChangedCallback>,
413    ) -> Rc<RefCell<ScopeEntry>> {
414        let entry_id = self.next_entry_id.get();
415        self.next_entry_id.set(entry_id.wrapping_add(1));
416        let recompose_scope_id = (&scope as &dyn Any)
417            .downcast_ref::<RecomposeScope>()
418            .map(RecomposeScope::id);
419        let owned_scope_key = recompose_scope_id
420            .is_none()
421            .then(|| owned_scope_index_key(&scope));
422        let entry = Rc::new(RefCell::new(ScopeEntry::new(entry_id, scope, on_changed)));
423        self.indexed_scopes
424            .borrow_mut()
425            .insert(entry_id, Rc::clone(&entry));
426        if let Some(scope_id) = recompose_scope_id {
427            self.fast_scopes
428                .borrow_mut()
429                .insert(scope_id, Rc::clone(&entry));
430        } else if let Some(scope_key) = owned_scope_key {
431            self.owned_scopes
432                .borrow_mut()
433                .entry(scope_key)
434                .or_default()
435                .push(Rc::clone(&entry));
436        }
437        entry
438    }
439
440    fn prune_dead_scopes(&self) {
441        let removed_fast = {
442            let mut fast_scopes = self.fast_scopes.borrow_mut();
443            let removed_ids: Vec<_> = fast_scopes
444                .iter()
445                .filter(|(_, entry)| !entry.borrow().should_retain())
446                .map(|(scope_id, _)| *scope_id)
447                .collect();
448            let removed = removed_ids
449                .into_iter()
450                .filter_map(|scope_id| fast_scopes.remove(&scope_id))
451                .collect::<Vec<_>>();
452            shrink_map_if_sparse(&mut fast_scopes, Self::MIN_RETAINED_SCOPE_CAPACITY);
453            removed
454        };
455
456        let removed_owned = { self.partition_owned_scopes(|entry| !entry.should_retain()) };
457
458        for entry in removed_fast.into_iter().chain(removed_owned) {
459            self.unregister_entry(&entry);
460        }
461    }
462
463    fn find_owned_scope_entry<T>(&self, scope: &T) -> Option<Rc<RefCell<ScopeEntry>>>
464    where
465        T: Any + Eq + Hash + 'static,
466    {
467        let key = owned_scope_index_key(scope);
468        self.owned_scopes.borrow().get(&key).and_then(|bucket| {
469            bucket
470                .iter()
471                .find(|entry| entry.borrow().matches_scope(scope))
472                .cloned()
473        })
474    }
475
476    fn remove_owned_scope_entry<T>(&self, scope: &T) -> Option<Rc<RefCell<ScopeEntry>>>
477    where
478        T: Any + Eq + Hash + 'static,
479    {
480        let key = owned_scope_index_key(scope);
481        let mut owned_scopes = self.owned_scopes.borrow_mut();
482        let mut removed = None;
483        let mut remove_bucket = false;
484        if let Some(bucket) = owned_scopes.get_mut(&key)
485            && let Some(index) = bucket
486                .iter()
487                .position(|entry| entry.borrow().matches_scope(scope))
488        {
489            removed = Some(bucket.remove(index));
490            remove_bucket = bucket.is_empty();
491        }
492        if remove_bucket {
493            owned_scopes.remove(&key);
494        }
495        shrink_map_if_sparse(&mut owned_scopes, Self::MIN_RETAINED_SCOPE_CAPACITY);
496        removed
497    }
498
499    fn partition_owned_scopes(
500        &self,
501        should_remove: impl Fn(&ScopeEntry) -> bool,
502    ) -> Vec<Rc<RefCell<ScopeEntry>>> {
503        let mut owned_scopes = self.owned_scopes.borrow_mut();
504        let mut retained = HashMap::default();
505        let mut removed = Vec::new();
506        for (key, mut bucket) in owned_scopes.drain() {
507            let mut retained_bucket = OwnedScopeBucket::new();
508            for entry in bucket.drain(..) {
509                if should_remove(&entry.borrow()) {
510                    removed.push(entry);
511                } else {
512                    retained_bucket.push(entry);
513                }
514            }
515            if !retained_bucket.is_empty() {
516                retained.insert(key, retained_bucket);
517            }
518        }
519        *owned_scopes = retained;
520        shrink_map_if_sparse(&mut owned_scopes, Self::MIN_RETAINED_SCOPE_CAPACITY);
521        removed
522    }
523
524    fn debug_stats(&self) -> SnapshotStateObserverDebugStats {
525        let owned_scopes = self.owned_scopes.borrow();
526        let fast_scopes = self.fast_scopes.borrow();
527        let owned_scope_len = owned_scopes.values().map(SmallVec::len).sum::<usize>();
528        let owned_scope_cap =
529            owned_scopes.capacity() + owned_scopes.values().map(SmallVec::capacity).sum::<usize>();
530        let scopes_len = owned_scope_len + fast_scopes.len();
531        let scopes_cap = owned_scope_cap + fast_scopes.capacity();
532        let mut observed_state_count = 0;
533        let mut observed_state_capacity = 0;
534        let mut stateless_scope_count = 0;
535
536        for entry in owned_scopes
537            .values()
538            .flat_map(|bucket| bucket.iter())
539            .chain(fast_scopes.values())
540        {
541            let entry = entry.borrow();
542            observed_state_count += entry.observed.len();
543            observed_state_capacity += entry.observed.capacity();
544            stateless_scope_count += usize::from(entry.observed.is_empty());
545        }
546
547        SnapshotStateObserverDebugStats {
548            scopes_len,
549            scopes_cap,
550            fast_scopes_len: fast_scopes.len(),
551            fast_scopes_cap: fast_scopes.capacity(),
552            stateless_scope_count,
553            observed_state_count,
554            observed_state_capacity,
555        }
556    }
557
558    fn run_with_read_observer<R>(&self, block: impl FnOnce() -> R) -> R {
559        use crate::snapshot_v2::take_transparent_observer_mutable_snapshot_reusing;
560
561        let mut snapshot = take_transparent_observer_mutable_snapshot_reusing(
562            Some(self.read_dispatcher.clone()),
563            None,
564            self.read_snapshot.take(),
565        );
566        let result = snapshot.enter(block);
567        snapshot.dispose();
568        if Arc::get_mut(&mut snapshot).is_some() && !snapshot.has_pending_changes() {
569            self.read_snapshot.replace(Some(snapshot));
570        }
571        result
572    }
573
574    fn handle_apply(&self, modified: &[Arc<dyn StateObject>]) {
575        if modified.is_empty() {
576            return;
577        }
578
579        let mut seen_scope_ids: HashSet<usize> = HashSet::default();
580        let mut to_notify: Vec<Rc<RefCell<ScopeEntry>>> = Vec::new();
581        {
582            let observed_to_scopes = self.observed_to_scopes.borrow();
583            let indexed_scopes = self.indexed_scopes.borrow();
584            for state in modified {
585                if let Some(scope_ids) = observed_to_scopes.get(&state.object_id().as_usize()) {
586                    let mut ordered_scope_ids: SmallVec<[usize; 8]> =
587                        scope_ids.iter().copied().collect();
588                    ordered_scope_ids.sort_unstable();
589                    for scope_id in ordered_scope_ids {
590                        if seen_scope_ids.insert(scope_id)
591                            && let Some(entry) = indexed_scopes.get(&scope_id)
592                        {
593                            to_notify.push(entry.clone());
594                        }
595                    }
596                }
597            }
598        }
599
600        if to_notify.is_empty() {
601            return;
602        }
603
604        for entry in to_notify {
605            let executor = self.executor.clone();
606            executor(Box::new(move || {
607                if let Ok(entry) = entry.try_borrow() {
608                    entry.notify();
609                }
610            }));
611        }
612    }
613
614    fn replace_observed_ids(&self, entry: &Rc<RefCell<ScopeEntry>>, observed: ObservedIds) {
615        let (entry_id, previous) = {
616            let mut entry_mut = entry.borrow_mut();
617            let entry_id = entry_mut.id;
618            let previous = std::mem::replace(&mut entry_mut.observed, observed);
619            (entry_id, previous)
620        };
621        let entry_ref = entry.borrow();
622        if previous.iter().eq(entry_ref.observed.iter()) {
623            return;
624        }
625        self.unregister_observed_ids(entry_id, &previous);
626        self.register_observed_ids(entry_id, &entry_ref.observed);
627    }
628
629    fn register_observed_ids(&self, entry_id: usize, observed: &ObservedIds) {
630        let mut observed_to_scopes = self.observed_to_scopes.borrow_mut();
631        for state_id in observed.iter() {
632            let scope_ids = observed_to_scopes.entry(state_id).or_default();
633            scope_ids.insert(entry_id);
634        }
635    }
636
637    fn unregister_observed_ids(&self, entry_id: usize, observed: &ObservedIds) {
638        let mut observed_to_scopes = self.observed_to_scopes.borrow_mut();
639        let mut emptied = SmallVec::<[StateObjectId; MAX_OBSERVED_STATES]>::new();
640        for state_id in observed.iter() {
641            if let Some(scope_ids) = observed_to_scopes.get_mut(&state_id) {
642                scope_ids.remove(&entry_id);
643                if scope_ids.is_empty() {
644                    emptied.push(state_id);
645                }
646            }
647        }
648        for state_id in emptied {
649            observed_to_scopes.remove(&state_id);
650        }
651        shrink_map_if_sparse(&mut observed_to_scopes, Self::MIN_RETAINED_SCOPE_CAPACITY);
652    }
653
654    fn unregister_entry(&self, entry: &Rc<RefCell<ScopeEntry>>) {
655        let (entry_id, observed) = {
656            let mut entry_mut = entry.borrow_mut();
657            let observed = std::mem::replace(&mut entry_mut.observed, ObservedIds::new());
658            (entry_mut.id, observed)
659        };
660        self.unregister_observed_ids(entry_id, &observed);
661        self.indexed_scopes.borrow_mut().remove(&entry_id);
662    }
663}
664
665fn shrink_map_if_sparse<K, V>(map: &mut HashMap<K, V>, min_retained_capacity: usize)
666where
667    K: Eq + std::hash::Hash,
668{
669    if map.capacity() <= map.len().max(min_retained_capacity).saturating_mul(4) {
670        return;
671    }
672
673    let retained = map.len().max(min_retained_capacity);
674    let mut rebuilt = HashMap::default();
675    rebuilt.reserve(retained);
676    rebuilt.extend(map.drain());
677    *map = rebuilt;
678}
679
680#[derive(Default)]
681struct ReadObservationStack {
682    targets: Vec<Rc<RefCell<ObservedIds>>>,
683    depth: usize,
684}
685
686impl ReadObservationStack {
687    fn push(&mut self) -> Rc<RefCell<ObservedIds>> {
688        if self.depth == self.targets.len() {
689            self.targets.push(Rc::new(RefCell::new(ObservedIds::new())));
690        }
691        let target = Rc::clone(&self.targets[self.depth]);
692        self.depth += 1;
693        target
694    }
695
696    fn last(&self) -> Option<&Rc<RefCell<ObservedIds>>> {
697        self.depth.checked_sub(1).map(|index| &self.targets[index])
698    }
699
700    fn pop(&mut self) -> Rc<RefCell<ObservedIds>> {
701        self.depth -= 1;
702        Rc::clone(&self.targets[self.depth])
703    }
704}
705
706enum ObservedIds {
707    Small(SmallVec<[ObservedState; MAX_OBSERVED_STATES]>),
708    Large(HashMap<StateObjectId, Option<Rc<dyn Any>>>),
709}
710
711struct ObservedState {
712    id: StateObjectId,
713    _lease: Option<Rc<dyn Any>>,
714}
715
716impl ObservedIds {
717    fn new() -> Self {
718        ObservedIds::Small(SmallVec::new())
719    }
720
721    fn insert(&mut self, state: &dyn StateObject) {
722        let id = state.object_id().as_usize();
723        match self {
724            ObservedIds::Small(small) => {
725                if small.iter().any(|observed| observed.id == id) {
726                    return;
727                }
728                if small.len() < MAX_OBSERVED_STATES {
729                    small.push(ObservedState {
730                        id,
731                        _lease: state.observation_lease(),
732                    });
733                } else {
734                    let mut large =
735                        HashMap::with_capacity_and_hasher(small.len() + 1, Default::default());
736                    for observed in small.drain(..) {
737                        large.insert(observed.id, observed._lease);
738                    }
739                    large.insert(id, state.observation_lease());
740                    *self = ObservedIds::Large(large);
741                }
742            }
743            ObservedIds::Large(large) => {
744                large.entry(id).or_insert_with(|| state.observation_lease());
745            }
746        }
747    }
748
749    fn is_empty(&self) -> bool {
750        match self {
751            ObservedIds::Small(small) => small.is_empty(),
752            ObservedIds::Large(large) => large.is_empty(),
753        }
754    }
755
756    fn len(&self) -> usize {
757        match self {
758            ObservedIds::Small(small) => small.len(),
759            ObservedIds::Large(large) => large.len(),
760        }
761    }
762
763    fn capacity(&self) -> usize {
764        match self {
765            ObservedIds::Small(small) => small.capacity(),
766            ObservedIds::Large(large) => large.capacity(),
767        }
768    }
769
770    fn iter(&self) -> impl Iterator<Item = StateObjectId> + '_ {
771        let (small, large) = match self {
772            ObservedIds::Small(small) => (Some(small.as_slice()), None),
773            ObservedIds::Large(large) => (None, Some(large)),
774        };
775        small
776            .into_iter()
777            .flatten()
778            .map(|observed| observed.id)
779            .chain(large.into_iter().flat_map(|states| states.keys().copied()))
780    }
781}
782
783const MAX_OBSERVED_STATES: usize = 8;
784
785enum ScopeStorage {
786    Owned(Box<dyn Any>),
787    RecomposeScope {
788        id: ScopeId,
789        weak: Weak<RecomposeScopeInner>,
790    },
791}
792
793struct ScopeEntry {
794    id: usize,
795    scope: ScopeStorage,
796    on_changed: Rc<dyn ScopeChangedCallback>,
797    observed: ObservedIds,
798    last_seen_version: u64,
799}
800
801impl ScopeEntry {
802    fn new<T>(id: usize, scope: T, on_changed: Rc<dyn ScopeChangedCallback>) -> Self
803    where
804        T: Any + 'static,
805    {
806        Self {
807            id,
808            scope: ScopeStorage::from_value(scope),
809            on_changed,
810            observed: ObservedIds::new(),
811            last_seen_version: u64::MAX,
812        }
813    }
814
815    fn callback_reusing<F: Fn(&dyn Any) + 'static>(
816        &mut self,
817        callback: F,
818    ) -> Rc<dyn ScopeChangedCallback> {
819        if let Some(stored) = Rc::get_mut(&mut self.on_changed)
820            .and_then(|stored| (stored as &mut dyn Any).downcast_mut::<F>())
821        {
822            *stored = callback;
823            Rc::clone(&self.on_changed)
824        } else {
825            Rc::new(callback)
826        }
827    }
828
829    fn update<T>(&mut self, new_scope: T, on_changed: Rc<dyn ScopeChangedCallback>)
830    where
831        T: Any + 'static,
832    {
833        self.update_scope(new_scope);
834        self.on_changed = on_changed;
835    }
836
837    fn update_scope<T>(&mut self, new_scope: T)
838    where
839        T: Any + 'static,
840    {
841        if let ScopeStorage::Owned(stored) = &mut self.scope
842            && let Some(stored) = stored.downcast_mut::<T>()
843        {
844            *stored = new_scope;
845        } else {
846            self.scope = ScopeStorage::from_value(new_scope);
847        }
848    }
849
850    fn matches_scope<T>(&self, scope: &T) -> bool
851    where
852        T: Any + Eq + 'static,
853    {
854        if let Some(scope) = (scope as &dyn Any).downcast_ref::<RecomposeScope>() {
855            return matches!(
856                &self.scope,
857                ScopeStorage::RecomposeScope { id, .. } if *id == scope.id()
858            );
859        }
860
861        match &self.scope {
862            ScopeStorage::Owned(stored) => stored
863                .downcast_ref::<T>()
864                .is_some_and(|stored| stored == scope),
865            ScopeStorage::RecomposeScope { .. } => false,
866        }
867    }
868
869    fn matches_predicate(&self, predicate: &impl Fn(&dyn Any) -> bool) -> bool {
870        match &self.scope {
871            ScopeStorage::Owned(scope) => predicate(scope.as_ref()),
872            ScopeStorage::RecomposeScope { weak, .. } => weak
873                .upgrade()
874                .is_none_or(|inner| predicate(&RecomposeScope { inner })),
875        }
876    }
877
878    fn should_retain(&self) -> bool {
879        match &self.scope {
880            ScopeStorage::Owned(_) => true,
881            ScopeStorage::RecomposeScope { weak, .. } => weak.upgrade().is_some(),
882        }
883    }
884
885    fn notify(&self) {
886        match &self.scope {
887            ScopeStorage::Owned(scope) => (self.on_changed)(scope.as_ref()),
888            ScopeStorage::RecomposeScope { weak, .. } => {
889                if let Some(inner) = weak.upgrade() {
890                    (self.on_changed)(&RecomposeScope { inner });
891                }
892            }
893        }
894    }
895}
896
897impl ScopeStorage {
898    fn from_value<T>(value: T) -> Self
899    where
900        T: Any + 'static,
901    {
902        let any = &value as &dyn Any;
903        if let Some(scope) = any.downcast_ref::<RecomposeScope>() {
904            Self::RecomposeScope {
905                id: scope.id(),
906                weak: scope.downgrade(),
907            }
908        } else {
909            Self::Owned(Box::new(value))
910        }
911    }
912}
913
914#[cfg(test)]
915#[path = "tests/snapshot_state_observer_tests.rs"]
916mod tests;
cranpose_core/snapshot_state_observer.rs

cranpose_core/
snapshot_state_observer.rs
