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 recompose_scopes_len: usize,
49    pub stateless_scope_count: usize,
50    pub observed_state_count: usize,
51    pub observed_state_capacity: usize,
52}
53
54impl SnapshotStateObserver {
55    /// Create a new observer that schedules callbacks using `on_changed_executor`.
56    pub fn new(on_changed_executor: impl Fn(Box<dyn FnOnce() + 'static>) + 'static) -> Self {
57        let inner = Rc::new(SnapshotStateObserverInner::new(on_changed_executor));
58        inner.set_self(Rc::downgrade(&inner));
59        Self { inner }
60    }
61
62    /// Observe state object reads performed while executing `block`.
63    ///
64    /// Subsequent calls to `observe_reads` replace any previously recorded
65    /// observations for the provided `scope`. When one of the observed objects
66    /// mutates, `on_value_changed_for_scope` will be invoked on the executor.
67    pub fn observe_reads<T, R>(
68        &self,
69        scope: T,
70        on_value_changed_for_scope: impl Fn(&T) + 'static,
71        block: impl FnOnce() -> R,
72    ) -> R
73    where
74        T: Any + Clone + Eq + Hash + 'static,
75    {
76        self.inner
77            .observe_reads(scope, on_value_changed_for_scope, block)
78    }
79
80    pub(crate) fn ptr_eq(&self, other: &Self) -> bool {
81        Rc::ptr_eq(&self.inner, &other.inner)
82    }
83
84    /// Notify the observer that a new composition frame is starting.
85    pub fn begin_frame(&self) {
86        self.inner.begin_frame();
87    }
88
89    /// Drop bookkeeping for scopes that were released during the current frame.
90    pub fn prune_dead_scopes(&self) {
91        self.inner.prune_dead_scopes();
92    }
93
94    /// Temporarily pause read observation while executing `block`.
95    pub fn with_no_observations<R>(&self, block: impl FnOnce() -> R) -> R {
96        self.inner.with_no_observations(block)
97    }
98
99    /// Remove any recorded reads for `scope`.
100    pub fn clear<T>(&self, scope: &T)
101    where
102        T: Any + Eq + Hash + 'static,
103    {
104        self.inner.clear(scope);
105    }
106
107    /// Remove recorded reads for scopes that satisfy `predicate`.
108    pub fn clear_if(&self, predicate: impl Fn(&dyn Any) -> bool) {
109        self.inner.clear_if(predicate);
110    }
111
112    /// Remove all recorded observations.
113    pub fn clear_all(&self) {
114        self.inner.clear_all();
115    }
116
117    /// Begin listening for snapshot apply notifications.
118    pub fn start(&self) {
119        let weak = Rc::downgrade(&self.inner);
120        self.inner.start(weak);
121    }
122
123    /// Stop listening for snapshot apply notifications.
124    pub fn stop(&self) {
125        self.inner.stop();
126    }
127
128    pub fn debug_stats(&self) -> SnapshotStateObserverDebugStats {
129        self.inner.debug_stats()
130    }
131
132    #[cfg(test)]
133    pub fn notify_changes(&self, modified: &[Arc<dyn StateObject>]) {
134        self.inner.handle_apply(modified);
135    }
136}
137
138struct SnapshotStateObserverInner {
139    executor: Rc<Executor>,
140    owned_scopes: RefCell<HashMap<OwnedScopeIndexKey, OwnedScopeBucket>>,
141    indexed_scopes: RefCell<HashMap<usize, Rc<RefCell<ScopeEntry>>>>,
142    observed_to_scopes: RefCell<HashMap<StateObjectId, HashSet<usize>>>,
143    pause_count: Rc<Cell<usize>>,
144    active_read_targets: Rc<RefCell<ReadObservationStack>>,
145    read_dispatcher: ReadObserver,
146    read_snapshot: RefCell<Option<Arc<TransparentObserverMutableSnapshot>>>,
147    apply_handle: RefCell<Option<crate::snapshot_v2::ObserverHandle>>,
148    weak_self: RefCell<Weak<SnapshotStateObserverInner>>,
149    frame_version: Cell<u64>,
150    dead_entries: RefCell<Vec<Rc<RefCell<ScopeEntry>>>>,
151    next_entry_id: Cell<usize>,
152    /// One `Rc` per type of callback that captures nothing: see
153    /// [`SnapshotStateObserverInner::capture_free_callback`].
154    capture_free_callbacks: RefCell<CaptureFreeCallbacks>,
155}
156
157#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
158struct OwnedScopeIndexKey {
159    type_id: TypeId,
160    value_hash: u64,
161}
162
163type OwnedScopeBucket = SmallVec<[Rc<RefCell<ScopeEntry>>; 1]>;
164
165/// The shared `Rc` of each type of callback that captures nothing.
166type CaptureFreeCallbacks = SmallVec<[(TypeId, Rc<dyn ScopeChangedCallback>); 2]>;
167
168fn owned_scope_index_key<T>(scope: &T) -> OwnedScopeIndexKey
169where
170    T: Any + Hash + 'static,
171{
172    let mut hasher = default_hash::new();
173    scope.hash(&mut hasher);
174    OwnedScopeIndexKey {
175        type_id: TypeId::of::<T>(),
176        value_hash: hasher.finish(),
177    }
178}
179
180impl SnapshotStateObserverInner {
181    const MIN_RETAINED_SCOPE_CAPACITY: usize = 256;
182
183    fn new(on_changed_executor: impl Fn(Box<dyn FnOnce() + 'static>) + 'static) -> Self {
184        let pause_count = Rc::new(Cell::new(0));
185        let active_read_targets = Rc::new(RefCell::new(ReadObservationStack::default()));
186        let dispatcher_pause_count = Rc::clone(&pause_count);
187        let dispatcher_targets = Rc::clone(&active_read_targets);
188        let read_dispatcher: ReadObserver = Arc::new(move |state| {
189            if dispatcher_pause_count.get() > 0 {
190                return;
191            }
192            let observed = dispatcher_targets.borrow().last().cloned();
193            if let Some(observed) = observed {
194                observed.borrow_mut().insert(state);
195            }
196        });
197
198        Self {
199            executor: Rc::new(on_changed_executor),
200            owned_scopes: RefCell::new(HashMap::default()),
201            indexed_scopes: RefCell::new(HashMap::default()),
202            observed_to_scopes: RefCell::new(HashMap::default()),
203            pause_count,
204            active_read_targets,
205            read_dispatcher,
206            read_snapshot: RefCell::new(None),
207            apply_handle: RefCell::new(None),
208            weak_self: RefCell::new(Weak::new()),
209            frame_version: Cell::new(0),
210            dead_entries: RefCell::new(Vec::new()),
211            next_entry_id: Cell::new(0),
212            capture_free_callbacks: RefCell::new(SmallVec::new()),
213        }
214    }
215
216    fn set_self(&self, weak: Weak<SnapshotStateObserverInner>) {
217        self.weak_self.replace(weak);
218    }
219
220    fn begin_frame(&self) {
221        let next = self.frame_version.get().wrapping_add(1);
222        self.frame_version.set(next);
223        self.prune_dead_scopes();
224    }
225
226    fn observe_reads<T, R>(
227        &self,
228        scope: T,
229        on_value_changed_for_scope: impl Fn(&T) + 'static,
230        block: impl FnOnce() -> R,
231    ) -> R
232    where
233        T: Any + Clone + Eq + Hash + 'static,
234    {
235        let frame_version = self.frame_version.get();
236        let has_frame_version = frame_version != 0;
237
238        let existing_entry = self.find_scope_entry(&scope);
239        let on_changed = std::cell::LazyCell::new(|| {
240            let callback = move |scope_any: &dyn Any| {
241                if let Some(typed) = scope_any.downcast_ref::<T>() {
242                    on_value_changed_for_scope(typed);
243                }
244            };
245            if std::mem::size_of_val(&callback) == 0 {
246                return self.capture_free_callback(callback);
247            }
248            match existing_entry.as_ref() {
249                Some(entry) => entry.borrow_mut().callback_reusing(callback),
250                None => Rc::new(callback),
251            }
252        });
253
254        if let Some(entry) = existing_entry.as_ref() {
255            let already_observed = {
256                let mut entry_mut = entry.borrow_mut();
257                entry_mut.update_scope(&scope);
258                has_frame_version && entry_mut.last_seen_version == frame_version
259            };
260            let callback = on_changed.clone();
261            entry.borrow_mut().on_changed = callback;
262            if already_observed {
263                return block();
264            }
265        }
266
267        let observed = self.active_read_targets.borrow_mut().push();
268        struct ActiveObservationGuard<'a> {
269            stack: &'a RefCell<ReadObservationStack>,
270        }
271        impl Drop for ActiveObservationGuard<'_> {
272            fn drop(&mut self) {
273                let target = self.stack.borrow_mut().pop();
274                target.borrow_mut().clear();
275            }
276        }
277        let _guard = ActiveObservationGuard {
278            stack: &self.active_read_targets,
279        };
280
281        let result = self.run_with_read_observer(block);
282
283        if observed.borrow().is_empty() {
284            if existing_entry.is_some() {
285                self.clear(&scope);
286            }
287            return result;
288        }
289
290        let callback = Rc::clone(&on_changed);
291        drop(on_changed);
292        let entry = existing_entry
293            .unwrap_or_else(|| self.insert_scope_entry(scope.clone(), Rc::clone(&callback)));
294        {
295            let mut entry_mut = entry.borrow_mut();
296            entry_mut.update(&scope, callback);
297            entry_mut.last_seen_version = if has_frame_version {
298                frame_version
299            } else {
300                u64::MAX
301            };
302        }
303        self.replace_observed_ids(&entry, &mut observed.borrow_mut());
304
305        result
306    }
307
308    fn with_no_observations<R>(&self, block: impl FnOnce() -> R) -> R {
309        self.pause_count.set(self.pause_count.get() + 1);
310        let result = block();
311        self.pause_count
312            .set(self.pause_count.get().saturating_sub(1));
313        result
314    }
315
316    fn clear<T>(&self, scope: &T)
317    where
318        T: Any + Eq + Hash + 'static,
319    {
320        if let Some(rc_scope) = (scope as &dyn Any).downcast_ref::<RecomposeScope>() {
321            if let Some(entry) = rc_scope.inner.observers.unlink(self) {
322                self.unregister_entry(&entry);
323            }
324            return;
325        }
326
327        let removed = self.remove_owned_scope_entry(scope);
328        if let Some(entry) = removed {
329            self.unregister_entry(&entry);
330        }
331    }
332
333    fn clear_if(&self, predicate: impl Fn(&dyn Any) -> bool) {
334        let removed_recompose: Vec<_> = self
335            .indexed_scopes
336            .borrow()
337            .values()
338            .filter(|entry| {
339                let entry = entry.borrow();
340                entry.holds_recompose_scope() && entry.matches_predicate(&predicate)
341            })
342            .cloned()
343            .collect();
344        for entry in &removed_recompose {
345            self.unlink_entry(entry);
346        }
347        let removed_owned =
348            { self.partition_owned_scopes(|entry| entry.matches_predicate(&predicate)) };
349
350        for entry in removed_recompose.into_iter().chain(removed_owned) {
351            self.unregister_entry(&entry);
352        }
353    }
354
355    fn clear_all(&self) {
356        let entries = std::mem::take(&mut *self.indexed_scopes.borrow_mut());
357        for entry in entries.values() {
358            self.unlink_entry(entry);
359        }
360        let owned = std::mem::take(&mut *self.owned_scopes.borrow_mut());
361        self.observed_to_scopes.borrow_mut().clear();
362        let dead = std::mem::take(&mut *self.dead_entries.borrow_mut());
363        drop((entries, owned, dead));
364    }
365
366    fn unlink_entry(&self, entry: &Rc<RefCell<ScopeEntry>>) {
367        let scope = entry.borrow().live_recompose_scope();
368        if let Some(scope) = scope {
369            scope.inner.observers.unlink(self);
370        }
371    }
372
373    fn start(&self, weak_self: Weak<SnapshotStateObserverInner>) {
374        if self.apply_handle.borrow().is_some() {
375            return;
376        }
377
378        let handle = register_apply_observer(Rc::new(move |modified, _snapshot_id| {
379            if let Some(inner) = weak_self.upgrade() {
380                inner.handle_apply(modified);
381            }
382        }));
383        self.apply_handle.replace(Some(handle));
384    }
385
386    fn stop(&self) {
387        if let Some(handle) = self.apply_handle.borrow_mut().take() {
388            drop(handle);
389        }
390    }
391
392    fn find_scope_entry<T>(&self, scope: &T) -> Option<Rc<RefCell<ScopeEntry>>>
393    where
394        T: Any + Eq + Hash + 'static,
395    {
396        if let Some(scope) = (scope as &dyn Any).downcast_ref::<RecomposeScope>() {
397            return scope.inner.observers.find(self);
398        }
399
400        self.find_owned_scope_entry(scope)
401    }
402
403    /// The `Rc` of a callback that captures nothing. Every closure of such a
404    /// type does the same thing, as `RecomposeScope::invalidate` does for
405    /// every group, so one `Rc` serves all its scopes instead of one each.
406    fn capture_free_callback<F: Fn(&dyn Any) + 'static>(
407        &self,
408        callback: F,
409    ) -> Rc<dyn ScopeChangedCallback> {
410        let type_id = TypeId::of::<F>();
411        let mut shared = self.capture_free_callbacks.borrow_mut();
412        if let Some((_, callback)) = shared.iter().find(|(id, _)| *id == type_id) {
413            return Rc::clone(callback);
414        }
415        let callback: Rc<dyn ScopeChangedCallback> = Rc::new(callback);
416        shared.push((type_id, Rc::clone(&callback)));
417        callback
418    }
419
420    fn insert_scope_entry(
421        &self,
422        scope: impl Any + Clone + Eq + Hash + 'static,
423        on_changed: Rc<dyn ScopeChangedCallback>,
424    ) -> Rc<RefCell<ScopeEntry>> {
425        let entry_id = self.next_entry_id.get();
426        self.next_entry_id.set(entry_id.wrapping_add(1));
427        let recompose_scope = (&scope as &dyn Any)
428            .downcast_ref::<RecomposeScope>()
429            .cloned();
430        let owned_scope_key = recompose_scope
431            .is_none()
432            .then(|| owned_scope_index_key(&scope));
433        let entry = Rc::new(RefCell::new(ScopeEntry::new(entry_id, scope, on_changed)));
434        self.indexed_scopes
435            .borrow_mut()
436            .insert(entry_id, Rc::clone(&entry));
437        if let Some(recompose_scope) = recompose_scope {
438            recompose_scope
439                .inner
440                .observers
441                .link(self.weak_self.borrow().clone(), Rc::clone(&entry));
442        } else if let Some(scope_key) = owned_scope_key {
443            self.owned_scopes
444                .borrow_mut()
445                .entry(scope_key)
446                .or_default()
447                .push(Rc::clone(&entry));
448        }
449        entry
450    }
451
452    fn prune_dead_scopes(&self) {
453        while let Some(entry) = self.pop_dead_entry() {
454            self.unregister_entry(&entry);
455        }
456    }
457
458    fn pop_dead_entry(&self) -> Option<Rc<RefCell<ScopeEntry>>> {
459        self.dead_entries.borrow_mut().pop()
460    }
461
462    fn find_owned_scope_entry<T>(&self, scope: &T) -> Option<Rc<RefCell<ScopeEntry>>>
463    where
464        T: Any + Eq + Hash + 'static,
465    {
466        let key = owned_scope_index_key(scope);
467        self.owned_scopes.borrow().get(&key).and_then(|bucket| {
468            bucket
469                .iter()
470                .find(|entry| entry.borrow().matches_scope(scope))
471                .cloned()
472        })
473    }
474
475    fn remove_owned_scope_entry<T>(&self, scope: &T) -> Option<Rc<RefCell<ScopeEntry>>>
476    where
477        T: Any + Eq + Hash + 'static,
478    {
479        let key = owned_scope_index_key(scope);
480        let mut owned_scopes = self.owned_scopes.borrow_mut();
481        let mut removed = None;
482        let mut remove_bucket = false;
483        if let Some(bucket) = owned_scopes.get_mut(&key)
484            && let Some(index) = bucket
485                .iter()
486                .position(|entry| entry.borrow().matches_scope(scope))
487        {
488            removed = Some(bucket.remove(index));
489            remove_bucket = bucket.is_empty();
490        }
491        if remove_bucket {
492            owned_scopes.remove(&key);
493        }
494        shrink_map_if_sparse(&mut owned_scopes, Self::MIN_RETAINED_SCOPE_CAPACITY);
495        removed
496    }
497
498    fn partition_owned_scopes(
499        &self,
500        should_remove: impl Fn(&ScopeEntry) -> bool,
501    ) -> Vec<Rc<RefCell<ScopeEntry>>> {
502        let mut owned_scopes = self.owned_scopes.borrow_mut();
503        let mut retained = HashMap::default();
504        let mut removed = Vec::new();
505        for (key, mut bucket) in owned_scopes.drain() {
506            let mut retained_bucket = OwnedScopeBucket::new();
507            for entry in bucket.drain(..) {
508                if should_remove(&entry.borrow()) {
509                    removed.push(entry);
510                } else {
511                    retained_bucket.push(entry);
512                }
513            }
514            if !retained_bucket.is_empty() {
515                retained.insert(key, retained_bucket);
516            }
517        }
518        *owned_scopes = retained;
519        shrink_map_if_sparse(&mut owned_scopes, Self::MIN_RETAINED_SCOPE_CAPACITY);
520        removed
521    }
522
523    fn debug_stats(&self) -> SnapshotStateObserverDebugStats {
524        let owned_scopes = self.owned_scopes.borrow();
525        let indexed_scopes = self.indexed_scopes.borrow();
526        let owned_scope_cap =
527            owned_scopes.capacity() + owned_scopes.values().map(SmallVec::capacity).sum::<usize>();
528        let scopes_cap = owned_scope_cap + indexed_scopes.capacity();
529        let mut observed_state_count = 0;
530        let mut observed_state_capacity = 0;
531        let mut stateless_scope_count = 0;
532        let mut recompose_scopes_len = 0;
533
534        for entry in indexed_scopes.values() {
535            let entry = entry.borrow();
536            observed_state_count += entry.observed.len();
537            observed_state_capacity += entry.observed.capacity();
538            stateless_scope_count += usize::from(entry.observed.is_empty());
539            recompose_scopes_len += usize::from(entry.holds_recompose_scope());
540        }
541
542        SnapshotStateObserverDebugStats {
543            scopes_len: indexed_scopes.len(),
544            scopes_cap,
545            recompose_scopes_len,
546            stateless_scope_count,
547            observed_state_count,
548            observed_state_capacity,
549        }
550    }
551
552    fn run_with_read_observer<R>(&self, block: impl FnOnce() -> R) -> R {
553        use crate::snapshot_v2::{
554            current_snapshot_reads_into, take_transparent_observer_mutable_snapshot_reusing,
555        };
556
557        if current_snapshot_reads_into(&self.read_dispatcher) {
558            return block();
559        }
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>>, collected: &mut ObservedIds) {
615        let (entry_id, previous) = {
616            let mut entry_mut = entry.borrow_mut();
617            if entry_mut.observed.iter().eq(collected.iter()) {
618                entry_mut.observed.take_leases(collected);
619                collected.clear();
620                return;
621            }
622            let entry_id = entry_mut.id;
623            let previous = std::mem::replace(&mut entry_mut.observed, collected.take_sized());
624            (entry_id, previous)
625        };
626        let entry_ref = entry.borrow();
627        self.unregister_observed_ids(entry_id, &previous);
628        self.register_observed_ids(entry_id, &entry_ref.observed);
629    }
630
631    fn register_observed_ids(&self, entry_id: usize, observed: &ObservedIds) {
632        let mut observed_to_scopes = self.observed_to_scopes.borrow_mut();
633        for state_id in observed.iter() {
634            let scope_ids = observed_to_scopes.entry(state_id).or_default();
635            scope_ids.insert(entry_id);
636        }
637    }
638
639    fn unregister_observed_ids(&self, entry_id: usize, observed: &ObservedIds) {
640        let mut observed_to_scopes = self.observed_to_scopes.borrow_mut();
641        let mut emptied = SmallVec::<[StateObjectId; MAX_OBSERVED_STATES]>::new();
642        for state_id in observed.iter() {
643            if let Some(scope_ids) = observed_to_scopes.get_mut(&state_id) {
644                scope_ids.remove(&entry_id);
645                if scope_ids.is_empty() {
646                    emptied.push(state_id);
647                }
648            }
649        }
650        for state_id in emptied {
651            observed_to_scopes.remove(&state_id);
652        }
653        shrink_map_if_sparse(&mut observed_to_scopes, Self::MIN_RETAINED_SCOPE_CAPACITY);
654    }
655
656    fn unregister_entry(&self, entry: &Rc<RefCell<ScopeEntry>>) {
657        let (entry_id, observed) = {
658            let mut entry_mut = entry.borrow_mut();
659            let observed = std::mem::replace(&mut entry_mut.observed, ObservedIds::new());
660            (entry_mut.id, observed)
661        };
662        self.unregister_observed_ids(entry_id, &observed);
663        self.indexed_scopes.borrow_mut().remove(&entry_id);
664    }
665}
666
667fn shrink_map_if_sparse<K, V>(map: &mut HashMap<K, V>, min_retained_capacity: usize)
668where
669    K: Eq + std::hash::Hash,
670{
671    if map.capacity() <= map.len().max(min_retained_capacity).saturating_mul(4) {
672        return;
673    }
674
675    let retained = map.len().max(min_retained_capacity);
676    let mut rebuilt = HashMap::default();
677    rebuilt.reserve(retained);
678    rebuilt.extend(map.drain());
679    *map = rebuilt;
680}
681
682#[derive(Default)]
683struct ReadObservationStack {
684    targets: Vec<Rc<RefCell<ObservedIds>>>,
685    depth: usize,
686}
687
688impl ReadObservationStack {
689    fn push(&mut self) -> Rc<RefCell<ObservedIds>> {
690        if self.depth == self.targets.len() {
691            self.targets.push(Rc::new(RefCell::new(ObservedIds::new())));
692        }
693        let target = Rc::clone(&self.targets[self.depth]);
694        self.depth += 1;
695        target
696    }
697
698    fn last(&self) -> Option<&Rc<RefCell<ObservedIds>>> {
699        self.depth.checked_sub(1).map(|index| &self.targets[index])
700    }
701
702    fn pop(&mut self) -> Rc<RefCell<ObservedIds>> {
703        self.depth -= 1;
704        Rc::clone(&self.targets[self.depth])
705    }
706}
707
708enum ObservedIds {
709    Small(SmallVec<[ObservedState; 1]>),
710    Large(Box<HashMap<StateObjectId, Option<Rc<dyn Any>>>>),
711}
712
713struct ObservedState {
714    id: StateObjectId,
715    _lease: Option<Rc<dyn Any>>,
716}
717
718impl ObservedIds {
719    fn new() -> Self {
720        ObservedIds::Small(SmallVec::new())
721    }
722
723    fn insert(&mut self, state: &dyn StateObject) {
724        let id = state.object_id().as_usize();
725        match self {
726            ObservedIds::Small(small) => {
727                if small.iter().any(|observed| observed.id == id) {
728                    return;
729                }
730                if small.len() < MAX_OBSERVED_STATES {
731                    small.push(ObservedState {
732                        id,
733                        _lease: state.observation_lease(),
734                    });
735                } else {
736                    let mut large =
737                        HashMap::with_capacity_and_hasher(small.len() + 1, Default::default());
738                    for observed in small.drain(..) {
739                        large.insert(observed.id, observed._lease);
740                    }
741                    large.insert(id, state.observation_lease());
742                    *self = ObservedIds::Large(Box::new(large));
743                }
744            }
745            ObservedIds::Large(large) => {
746                large.entry(id).or_insert_with(|| state.observation_lease());
747            }
748        }
749    }
750
751    fn clear(&mut self) {
752        match self {
753            ObservedIds::Small(small) => small.clear(),
754            ObservedIds::Large(_) => *self = ObservedIds::new(),
755        }
756    }
757
758    /// Keeps `fresh`'s observation leases for the states both name, leaving
759    /// the old ones in `fresh` to drop. An id is a state's address, which a
760    /// new state can take after the old one is dropped, and only its own
761    /// lease keeps the new one observed.
762    fn take_leases(&mut self, fresh: &mut ObservedIds) {
763        match (self, fresh) {
764            (ObservedIds::Small(kept), ObservedIds::Small(fresh)) => {
765                for (kept, fresh) in kept.iter_mut().zip(fresh.iter_mut()) {
766                    std::mem::swap(&mut kept._lease, &mut fresh._lease);
767                }
768            }
769            (kept, fresh) => std::mem::swap(kept, fresh),
770        }
771    }
772
773    fn take_sized(&mut self) -> ObservedIds {
774        match self {
775            ObservedIds::Small(small) => ObservedIds::Small(small.drain(..).collect()),
776            ObservedIds::Large(_) => std::mem::replace(self, ObservedIds::new()),
777        }
778    }
779
780    fn is_empty(&self) -> bool {
781        match self {
782            ObservedIds::Small(small) => small.is_empty(),
783            ObservedIds::Large(large) => large.is_empty(),
784        }
785    }
786
787    fn len(&self) -> usize {
788        match self {
789            ObservedIds::Small(small) => small.len(),
790            ObservedIds::Large(large) => large.len(),
791        }
792    }
793
794    fn capacity(&self) -> usize {
795        match self {
796            ObservedIds::Small(small) => small.capacity(),
797            ObservedIds::Large(large) => large.capacity(),
798        }
799    }
800
801    fn iter(&self) -> impl Iterator<Item = StateObjectId> + '_ {
802        let (small, large) = match self {
803            ObservedIds::Small(small) => (Some(small.as_slice()), None),
804            ObservedIds::Large(large) => (None, Some(large)),
805        };
806        small
807            .into_iter()
808            .flatten()
809            .map(|observed| observed.id)
810            .chain(large.into_iter().flat_map(|states| states.keys().copied()))
811    }
812}
813
814const MAX_OBSERVED_STATES: usize = 8;
815
816enum ScopeStorage {
817    Owned(Box<dyn Any>),
818    RecomposeScope {
819        id: ScopeId,
820        weak: Weak<RecomposeScopeInner>,
821    },
822}
823
824struct ScopeEntry {
825    id: usize,
826    scope: ScopeStorage,
827    on_changed: Rc<dyn ScopeChangedCallback>,
828    observed: ObservedIds,
829    last_seen_version: u64,
830}
831
832impl ScopeEntry {
833    fn new<T>(id: usize, scope: T, on_changed: Rc<dyn ScopeChangedCallback>) -> Self
834    where
835        T: Any + 'static,
836    {
837        Self {
838            id,
839            scope: ScopeStorage::from_value(scope),
840            on_changed,
841            observed: ObservedIds::new(),
842            last_seen_version: u64::MAX,
843        }
844    }
845
846    fn callback_reusing<F: Fn(&dyn Any) + 'static>(
847        &mut self,
848        callback: F,
849    ) -> Rc<dyn ScopeChangedCallback> {
850        if let Some(stored) = Rc::get_mut(&mut self.on_changed)
851            .and_then(|stored| (stored as &mut dyn Any).downcast_mut::<F>())
852        {
853            *stored = callback;
854            Rc::clone(&self.on_changed)
855        } else {
856            Rc::new(callback)
857        }
858    }
859
860    fn update<T>(&mut self, new_scope: &T, on_changed: Rc<dyn ScopeChangedCallback>)
861    where
862        T: Any + Clone + 'static,
863    {
864        self.update_scope(new_scope);
865        self.on_changed = on_changed;
866    }
867
868    fn update_scope<T>(&mut self, new_scope: &T)
869    where
870        T: Any + Clone + 'static,
871    {
872        match &mut self.scope {
873            ScopeStorage::Owned(stored) => {
874                if let Some(stored) = stored.downcast_mut::<T>() {
875                    stored.clone_from(new_scope);
876                    return;
877                }
878            }
879            ScopeStorage::RecomposeScope { weak, .. } => {
880                if (new_scope as &dyn Any)
881                    .downcast_ref::<RecomposeScope>()
882                    .is_some_and(|scope| Weak::as_ptr(weak) == Rc::as_ptr(&scope.inner))
883                {
884                    return;
885                }
886            }
887        }
888        self.scope = ScopeStorage::from_value(new_scope.clone());
889    }
890
891    fn matches_scope<T>(&self, scope: &T) -> bool
892    where
893        T: Any + Eq + 'static,
894    {
895        if let Some(scope) = (scope as &dyn Any).downcast_ref::<RecomposeScope>() {
896            return matches!(
897                &self.scope,
898                ScopeStorage::RecomposeScope { id, .. } if *id == scope.id()
899            );
900        }
901
902        match &self.scope {
903            ScopeStorage::Owned(stored) => stored
904                .downcast_ref::<T>()
905                .is_some_and(|stored| stored == scope),
906            ScopeStorage::RecomposeScope { .. } => false,
907        }
908    }
909
910    fn matches_predicate(&self, predicate: &impl Fn(&dyn Any) -> bool) -> bool {
911        match &self.scope {
912            ScopeStorage::Owned(scope) => predicate(scope.as_ref()),
913            ScopeStorage::RecomposeScope { weak, .. } => weak
914                .upgrade()
915                .is_none_or(|inner| predicate(&RecomposeScope { inner })),
916        }
917    }
918
919    fn holds_recompose_scope(&self) -> bool {
920        matches!(self.scope, ScopeStorage::RecomposeScope { .. })
921    }
922
923    fn live_recompose_scope(&self) -> Option<RecomposeScope> {
924        match &self.scope {
925            ScopeStorage::Owned(_) => None,
926            ScopeStorage::RecomposeScope { weak, .. } => {
927                weak.upgrade().map(|inner| RecomposeScope { inner })
928            }
929        }
930    }
931
932    fn notify(&self) {
933        match &self.scope {
934            ScopeStorage::Owned(scope) => (self.on_changed)(scope.as_ref()),
935            ScopeStorage::RecomposeScope { weak, .. } => {
936                if let Some(inner) = weak.upgrade() {
937                    (self.on_changed)(&RecomposeScope { inner });
938                }
939            }
940        }
941    }
942}
943
944impl ScopeStorage {
945    fn from_value<T>(value: T) -> Self
946    where
947        T: Any + 'static,
948    {
949        let any = &value as &dyn Any;
950        if let Some(scope) = any.downcast_ref::<RecomposeScope>() {
951            Self::RecomposeScope {
952                id: scope.id(),
953                weak: scope.downgrade(),
954            }
955        } else {
956            Self::Owned(Box::new(value))
957        }
958    }
959}
960
961#[derive(Default)]
962pub(crate) struct ScopeObservers {
963    head: Cell<Option<Box<ObserverLink>>>,
964}
965
966struct ObserverLink {
967    observer: Weak<SnapshotStateObserverInner>,
968    entry: Rc<RefCell<ScopeEntry>>,
969    next: Option<Box<ObserverLink>>,
970}
971
972impl ScopeObservers {
973    fn find(&self, observer: &SnapshotStateObserverInner) -> Option<Rc<RefCell<ScopeEntry>>> {
974        let head = self.head.take();
975        let mut link = head.as_deref();
976        let mut found = None;
977        while let Some(current) = link {
978            if std::ptr::eq(current.observer.as_ptr(), observer) {
979                found = Some(Rc::clone(&current.entry));
980                break;
981            }
982            link = current.next.as_deref();
983        }
984        self.head.set(head);
985        found
986    }
987
988    fn link(&self, observer: Weak<SnapshotStateObserverInner>, entry: Rc<RefCell<ScopeEntry>>) {
989        let next = self.head.take();
990        self.head.set(Some(Box::new(ObserverLink {
991            observer,
992            entry,
993            next,
994        })));
995    }
996
997    fn unlink(&self, observer: &SnapshotStateObserverInner) -> Option<Rc<RefCell<ScopeEntry>>> {
998        let mut rest = self.head.take();
999        let mut kept = None;
1000        let mut removed = None;
1001        while let Some(mut link) = rest {
1002            rest = link.next.take();
1003            if removed.is_none() && std::ptr::eq(link.observer.as_ptr(), observer) {
1004                removed = Some(link.entry);
1005            } else {
1006                link.next = kept;
1007                kept = Some(link);
1008            }
1009        }
1010        self.head.set(kept);
1011        removed
1012    }
1013
1014    pub(crate) fn release(&mut self) {
1015        let mut rest = self.head.get_mut().take();
1016        while let Some(link) = rest {
1017            let ObserverLink {
1018                observer,
1019                entry,
1020                next,
1021            } = *link;
1022            rest = next;
1023            if let Some(observer) = observer.upgrade() {
1024                observer.dead_entries.borrow_mut().push(entry);
1025            }
1026        }
1027    }
1028}
1029
1030impl Drop for SnapshotStateObserverInner {
1031    fn drop(&mut self) {
1032        for entry in self.indexed_scopes.borrow().values() {
1033            self.unlink_entry(entry);
1034        }
1035    }
1036}
1037
1038#[cfg(test)]
1039#[path = "tests/snapshot_state_observer_tests.rs"]
1040mod tests;
cranpose_core/snapshot_state_observer.rs

cranpose_core/
snapshot_state_observer.rs
