use crate::runtime::RuntimeInner;
use std::cell::RefCell;
use std::rc::{Rc, Weak};
#[derive(Clone)]
pub struct Memo<T> {
cached_value: Rc<RefCell<Option<T>>>,
compute: Rc<dyn Fn() -> T>,
id: usize,
runtime: Weak<RefCell<RuntimeInner>>,
}
impl<T: Clone + 'static> Memo<T> {
pub(crate) fn new_in_scope<F>(compute: F, id: usize, runtime: Rc<RefCell<RuntimeInner>>) -> Self
where
F: Fn() -> T + 'static,
{
runtime.borrow().register_memo(id);
Self {
cached_value: Rc::new(RefCell::new(None)),
compute: Rc::new(compute),
id,
runtime: Rc::downgrade(&runtime),
}
}
pub fn get(&self) -> T {
if let Some(runtime) = self.runtime.upgrade() {
runtime.borrow().track_read(self.id);
if runtime.borrow().is_memo_dirty(self.id) {
let value =
RuntimeInner::with_observer_scoped(&runtime, self.id, || (self.compute)());
*self.cached_value.borrow_mut() = Some(value.clone());
runtime.borrow().mark_memo_clean(self.id);
value
} else {
self.cached_value.borrow().as_ref().unwrap().clone()
}
} else {
(self.compute)()
}
}
pub fn with<R>(&self, f: impl FnOnce(&T) -> R) -> R {
if let Some(runtime) = self.runtime.upgrade() {
runtime.borrow().track_read(self.id);
if runtime.borrow().is_memo_dirty(self.id) {
let value =
RuntimeInner::with_observer_scoped(&runtime, self.id, || (self.compute)());
*self.cached_value.borrow_mut() = Some(value.clone());
runtime.borrow().mark_memo_clean(self.id);
let cached = self.cached_value.borrow();
f(cached.as_ref().unwrap())
} else {
let cached = self.cached_value.borrow();
f(cached.as_ref().unwrap())
}
} else {
let value = (self.compute)();
f(&value)
}
}
}
impl<T> Drop for Memo<T> {
fn drop(&mut self) {
if Rc::strong_count(&self.cached_value) == 1 {
if let Some(runtime) = self.runtime.upgrade() {
runtime.borrow().remove_memo(self.id);
}
}
}
}