Skip to main content

telar_reactive_core/
memo.rs

1use std::cell::RefCell;
2use std::rc::Rc;
3
4use smallvec::SmallVec;
5
6use crate::runtime::{self, EffectId};
7
8enum MemoState<T> {
9    Computing, // reading while Computing means the closure re-entered itself: a cycle
10    Clean(T),
11    Dirty,
12}
13
14struct MemoInner<T> {
15    state: MemoState<T>,
16    subscribers: SmallVec<[EffectId; 4]>,
17    effect_id: EffectId,
18}
19
20pub struct Memo<T: 'static> {
21    inner: Rc<RefCell<MemoInner<T>>>,
22}
23
24impl<T: 'static> Clone for Memo<T> {
25    fn clone(&self) -> Self {
26        Memo {
27            inner: Rc::clone(&self.inner),
28        }
29    }
30}
31
32impl<T: 'static> Drop for Memo<T> {
33    fn drop(&mut self) {
34        if Rc::strong_count(&self.inner) == 1 {
35            let id = self.inner.borrow().effect_id;
36            runtime::deregister_effect(id);
37        }
38    }
39}
40
41impl<T: Clone + 'static> Memo<T> {
42    pub fn get(&self) -> T {
43        self.track();
44        match &self.inner.borrow().state {
45            MemoState::Clean(v) => v.clone(),
46            MemoState::Dirty => panic!("memo read while Dirty — flush ordering issue"),
47            MemoState::Computing => panic!("reactive cycle detected in memo"),
48        }
49    }
50}
51
52impl<T: 'static> Memo<T> {
53    pub fn with<R>(&self, f: impl FnOnce(&T) -> R) -> R {
54        self.track();
55        let borrow = self.inner.borrow();
56        match &borrow.state {
57            MemoState::Clean(v) => f(v),
58            MemoState::Dirty => panic!("memo read while Dirty — flush ordering issue"),
59            MemoState::Computing => panic!("reactive cycle detected in memo"),
60        }
61    }
62
63    fn track(&self) {
64        if let Some(id) = runtime::current_observer() {
65            let mut borrow = self.inner.borrow_mut();
66            if !borrow.subscribers.contains(&id) {
67                borrow.subscribers.push(id);
68            }
69        }
70    }
71}
72
73pub fn memo<T: PartialEq + 'static>(f: impl Fn() -> T + 'static) -> Memo<T> {
74    use std::rc::Weak;
75
76    let inner: Rc<RefCell<MemoInner<T>>> = Rc::new(RefCell::new(MemoInner {
77        state: MemoState::Dirty,
78        subscribers: SmallVec::new(),
79        effect_id: 0,
80    }));
81
82    let weak: Weak<RefCell<MemoInner<T>>> = Rc::downgrade(&inner);
83
84    let effect_f: Box<dyn Fn()> = Box::new(move || {
85        let Some(inner) = weak.upgrade() else {
86            return;
87        };
88        inner.borrow_mut().state = MemoState::Computing;
89        let new_value = f();
90        let subs: SmallVec<[EffectId; 8]> = {
91            let mut memo = inner.borrow_mut();
92            let changed = match &memo.state {
93                MemoState::Clean(old) => old != &new_value,
94                _ => true,
95            };
96            memo.state = MemoState::Clean(new_value);
97            if changed {
98                memo.subscribers.iter().copied().collect()
99            } else {
100                SmallVec::new()
101            }
102        };
103        let mut dead: Option<Vec<EffectId>> = None;
104        for id in subs {
105            if runtime::is_alive(id) {
106                runtime::schedule(id);
107            } else {
108                dead.get_or_insert_with(Vec::new).push(id);
109            }
110        }
111        if let Some(dead) = dead {
112            let mut memo = inner.borrow_mut();
113            for id in dead {
114                memo.subscribers.retain(|x| x != &id);
115            }
116        }
117    });
118
119    // Register as a pure effect so it runs before user effects during flush
120    let effect_id = runtime::register_pure_effect(effect_f);
121    inner.borrow_mut().effect_id = effect_id;
122
123    runtime::run_effect(effect_id);
124
125    Memo { inner }
126}