use std::cell::{Cell, RefCell};
use std::collections::{HashMap, HashSet, VecDeque};
use std::rc::{Rc, Weak};
use std::sync::atomic::{AtomicU64, Ordering};
struct EffectInner {
id: u64,
body: RefCell<Box<dyn FnMut()>>,
}
type Effect = Rc<EffectInner>;
type WeakEffect = Weak<EffectInner>;
static NEXT_EFFECT_ID: AtomicU64 = AtomicU64::new(1);
fn next_effect_id() -> u64 {
NEXT_EFFECT_ID.fetch_add(1, Ordering::Relaxed)
}
thread_local! {
static CURRENT_EFFECT: RefCell<Option<Effect>> = const { RefCell::new(None) };
static EFFECT_QUEUE: RefCell<VecDeque<Effect>> = const { RefCell::new(VecDeque::new()) };
static QUEUED: RefCell<HashSet<u64>> = RefCell::new(HashSet::new());
static IS_FLUSHING: Cell<bool> = const { Cell::new(false) };
static EFFECT_STORAGE: RefCell<Vec<Effect>> = const { RefCell::new(Vec::new()) };
static STOPPED_EFFECTS: RefCell<HashMap<u64, WeakEffect>> = RefCell::new(HashMap::new());
static BATCH_DEPTH: Cell<usize> = const { Cell::new(0) };
}
const STOPPED_EFFECTS_PRUNE_AT: usize = 1024;
const SUBSCRIBER_SWEEP_AT: usize = 64;
fn mark_stopped(eff: &Effect) {
STOPPED_EFFECTS.with(|stopped| {
let mut stopped = stopped.borrow_mut();
stopped.insert(eff.id, Rc::downgrade(eff));
if stopped.len() > STOPPED_EFFECTS_PRUNE_AT {
stopped.retain(|_, weak| weak.strong_count() > 0);
}
});
}
fn enqueue_effect(eff: Effect) {
EFFECT_QUEUE.with(|q| {
QUEUED.with(|set| {
let id = eff.id;
let mut set_b = set.borrow_mut();
if set_b.insert(id) {
q.borrow_mut().push_back(eff);
}
});
});
}
pub fn batch(f: impl FnOnce()) {
struct BatchGuard;
impl Drop for BatchGuard {
fn drop(&mut self) {
BATCH_DEPTH.with(|d| d.set(d.get().saturating_sub(1)));
}
}
BATCH_DEPTH.with(|d| d.set(d.get().saturating_add(1)));
let _guard = BatchGuard;
f();
drop(_guard);
flush_if_not_batching();
}
fn flush_if_not_batching() {
if BATCH_DEPTH.with(|d| d.get()) == 0 {
flush_queue();
}
}
struct FlushResetGuard;
impl Drop for FlushResetGuard {
fn drop(&mut self) {
IS_FLUSHING.with(|f| f.set(false));
}
}
struct EffectScopeGuard {
was_flushing: bool,
previous_effect: Option<Effect>,
}
impl EffectScopeGuard {
fn enter(eff: &Effect) -> Self {
let was_flushing = IS_FLUSHING.with(|f| f.replace(true));
let previous_effect = CURRENT_EFFECT.with(|cur| cur.borrow_mut().replace(eff.clone()));
Self {
was_flushing,
previous_effect,
}
}
}
impl Drop for EffectScopeGuard {
fn drop(&mut self) {
let previous = self.previous_effect.take();
CURRENT_EFFECT.with(|cur| *cur.borrow_mut() = previous);
IS_FLUSHING.with(|f| f.set(self.was_flushing));
}
}
fn flush_queue() {
if IS_FLUSHING.with(|f| f.replace(true)) {
return;
}
let _flush_guard = FlushResetGuard;
loop {
let next = EFFECT_QUEUE.with(|q| q.borrow_mut().pop_front());
let Some(eff) = next else { break };
let id = eff.id;
QUEUED.with(|set| {
set.borrow_mut().remove(&id);
});
let is_stopped = STOPPED_EFFECTS.with(|stopped| stopped.borrow().contains_key(&id));
if is_stopped {
continue;
}
let _scope = EffectScopeGuard::enter(&eff);
match eff.body.try_borrow_mut() {
Ok(mut body) => body(),
Err(_) => enqueue_effect(eff.clone()),
}
}
}
pub struct EffectHandle {
effect: Effect,
active: Cell<bool>,
}
impl EffectHandle {
pub fn stop(&self) {
if self.active.get() {
self.active.set(false);
let id = self.effect.id;
mark_stopped(&self.effect);
EFFECT_STORAGE.with(|storage| {
storage.borrow_mut().retain(|e| e.id != id);
});
}
}
pub fn is_active(&self) -> bool {
self.active.get()
}
}
impl Drop for EffectHandle {
fn drop(&mut self) {
let id = self.effect.id;
mark_stopped(&self.effect);
EFFECT_STORAGE.with(|storage| {
storage.borrow_mut().retain(|e| e.id != id);
});
}
}
pub struct Signal<T> {
value: RefCell<T>,
pending: RefCell<Option<T>>,
subscribers: RefCell<Vec<WeakEffect>>,
_effect_handle: RefCell<Option<EffectHandle>>,
}
impl<T> Signal<T>
where
T: Clone,
{
pub fn new(initial: T) -> Self {
Self {
value: RefCell::new(initial),
pending: RefCell::new(None),
subscribers: RefCell::new(Vec::new()),
_effect_handle: RefCell::new(None),
}
}
pub fn get(&self) -> T {
if let Some(pending) = self.pending.borrow_mut().take() {
if let Ok(mut v) = self.value.try_borrow_mut() {
*v = pending;
} else {
*self.pending.borrow_mut() = Some(pending);
}
}
CURRENT_EFFECT.with(|current| {
if let Some(effect_rc) = current.borrow().as_ref() {
let mut subs = self.subscribers.borrow_mut();
if subs.len() >= SUBSCRIBER_SWEEP_AT {
subs.retain(|w| w.strong_count() > 0);
}
let already_subscribed = subs
.iter()
.any(|w| w.upgrade().is_some_and(|rc| Rc::ptr_eq(&rc, effect_rc)));
if !already_subscribed {
subs.push(Rc::downgrade(effect_rc));
}
}
});
self.value
.try_borrow()
.ok()
.map(|v| v.clone())
.or_else(|| self.pending.try_borrow().ok().and_then(|slot| slot.clone()))
.expect(
"Signal::get: value is mutably borrowed and no deferred update \
is available; RefCell borrows never escape this module, so \
this indicates a borrow held across user code",
)
}
pub fn set(&self, new: T) {
match self.value.try_borrow_mut() {
Ok(mut v) => *v = new,
Err(_) => {
*self.pending.try_borrow_mut().expect(
"Signal::set: deferred slot is already borrowed; \
RefCell borrows never escape this module, so this \
indicates a borrow held across user code",
) = Some(new);
}
}
let subscribers: Vec<Effect> = {
let mut subs = self.subscribers.borrow_mut();
let living: Vec<Effect> = subs.drain(..).filter_map(|w| w.upgrade()).collect();
for eff in &living {
subs.push(Rc::downgrade(eff));
}
living
};
for subscriber in subscribers {
enqueue_effect(subscriber);
}
flush_if_not_batching();
}
pub fn update<F>(&self, f: F)
where
F: FnOnce(T) -> T,
{
if let Some(pending) = self
.pending
.try_borrow_mut()
.ok()
.and_then(|mut slot| slot.take())
{
match self.value.try_borrow_mut() {
Ok(mut v) => *v = pending,
Err(_) => {
if let Ok(mut slot) = self.pending.try_borrow_mut() {
*slot = Some(pending);
}
}
}
}
let current = self
.value
.try_borrow()
.ok()
.map(|v| v.clone())
.or_else(|| self.pending.try_borrow().ok().and_then(|slot| slot.clone()))
.expect(
"Signal::update: value is mutably borrowed and no deferred \
update is available; RefCell borrows never escape this module",
);
let new = f(current);
self.set(new);
}
pub fn set_if_changed(&self, new: T)
where
T: PartialEq,
{
let changed = self.value.try_borrow().map(|v| *v != new).unwrap_or(true);
if changed {
self.set(new);
}
}
}
pub fn effect<F>(f: F) -> EffectHandle
where
F: FnMut() + 'static,
{
let eff = Rc::new(EffectInner {
id: next_effect_id(),
body: RefCell::new(Box::new(f) as Box<dyn FnMut()>),
});
let handle = EffectHandle {
effect: eff.clone(),
active: Cell::new(true),
};
EFFECT_STORAGE.with(|storage| {
storage.borrow_mut().push(eff.clone());
});
{
let _scope = EffectScopeGuard::enter(&eff);
match eff.body.try_borrow_mut() {
Ok(mut body) => body(),
Err(_) => {
drop(_scope);
enqueue_effect(eff.clone());
flush_if_not_batching();
}
}
}
handle
}
pub fn computed<T, F>(compute: F) -> Rc<Signal<T>>
where
T: Clone + PartialEq + 'static,
F: Fn() -> T + 'static,
{
use std::cell::RefCell;
let compute_rc = Rc::new(RefCell::new(compute));
let signal = Rc::new(Signal::new((compute_rc.borrow())()));
let signal_weak = Rc::downgrade(&signal);
let handle = effect({
let compute_rc = compute_rc.clone();
move || {
if let Some(sig) = signal_weak.upgrade() {
let new_value = (compute_rc.borrow())();
let changed = sig
.value
.try_borrow()
.map(|v| *v != new_value)
.unwrap_or(true);
if changed {
sig.set(new_value);
}
}
}
});
*signal._effect_handle.borrow_mut() = Some(handle);
signal
}
#[cfg(test)]
mod tests {
use super::*;
use std::cell::Cell;
use std::rc::Rc;
#[test]
fn stopped_effects_store_stays_bounded() {
const CHURN: usize = STOPPED_EFFECTS_PRUNE_AT * 4;
let mut handles = Vec::with_capacity(CHURN);
for _ in 0..CHURN {
let sig = Rc::new(Signal::new(0u32));
let s = sig.clone();
handles.push(effect(move || {
s.get();
}));
}
drop(handles);
let len = STOPPED_EFFECTS.with(|stopped| stopped.borrow().len());
assert!(
len <= STOPPED_EFFECTS_PRUNE_AT,
"STOPPED_EFFECTS held {len} entries after stopping {CHURN} effects; \
the prune threshold is {STOPPED_EFFECTS_PRUNE_AT}"
);
}
#[test]
fn effect_ids_are_unique_and_increasing() {
let sig = Rc::new(Signal::new(0u32));
let s1 = sig.clone();
let first = effect(move || {
s1.get();
});
let first_id = first.effect.id;
let s2 = sig.clone();
let second = effect(move || {
s2.get();
});
let second_id = second.effect.id;
assert_ne!(first_id, second_id, "two live effects must not share an id");
assert!(
second_id > first_id,
"ids should increase: got {first_id} then {second_id}"
);
drop((first, second));
}
#[test]
fn a_stopped_effect_whose_handle_is_still_held_does_not_run_again() {
let sig = Rc::new(Signal::new(0u32));
let runs = Rc::new(Cell::new(0u32));
let s = sig.clone();
let r = runs.clone();
let handle = effect(move || {
s.get();
r.set(r.get() + 1);
});
assert_eq!(runs.get(), 1, "the effect runs once on registration");
handle.stop();
assert!(!handle.is_active());
sig.set(1);
assert_eq!(
runs.get(),
1,
"a stopped effect ran again: {} runs",
runs.get()
);
drop(handle);
}
#[test]
fn a_prune_keeps_the_tombstone_of_a_stopped_but_live_effect() {
let sig = Rc::new(Signal::new(0u32));
let runs = Rc::new(Cell::new(0u32));
let s = sig.clone();
let r = runs.clone();
let sentinel = effect(move || {
s.get();
r.set(r.get() + 1);
});
assert_eq!(runs.get(), 1, "the effect runs once on registration");
sentinel.stop();
for _ in 0..=STOPPED_EFFECTS_PRUNE_AT {
let churn_sig = Rc::new(Signal::new(0u32));
let c = churn_sig.clone();
let handle = effect(move || {
c.get();
});
drop(handle);
}
let len = STOPPED_EFFECTS.with(|stopped| stopped.borrow().len());
assert!(
len <= STOPPED_EFFECTS_PRUNE_AT,
"expected a prune to have fired, but the store still holds {len} entries \
after {} stops (threshold {STOPPED_EFFECTS_PRUNE_AT})",
STOPPED_EFFECTS_PRUNE_AT + 2
);
sig.set(1);
assert_eq!(
runs.get(),
1,
"the prune discarded the tombstone of a stopped effect that was still \
alive, so it ran again: {} runs",
runs.get()
);
drop(sentinel);
}
fn panicking_on_second_run() -> (Rc<Signal<u32>>, Rc<Cell<u32>>, EffectHandle) {
let sig = Rc::new(Signal::new(0u32));
let runs = Rc::new(Cell::new(0u32));
let panicked = Rc::new(Cell::new(false));
let s = sig.clone();
let r = runs.clone();
let p = panicked.clone();
let handle = effect(move || {
r.set(r.get() + 1);
s.get();
if r.get() >= 2 && !p.replace(true) {
panic!("effect body panicked on purpose");
}
});
assert_eq!(runs.get(), 1, "the effect runs once on registration");
(sig, runs, handle)
}
#[test]
fn a_panicking_effect_body_does_not_strand_the_scheduler() {
let (sig, _runs, _panicking) = panicking_on_second_run();
let outcome = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
sig.set(1);
}));
assert!(
outcome.is_err(),
"the body was supposed to panic during the flush"
);
let other = Rc::new(Signal::new(0u32));
let other_runs = Rc::new(Cell::new(0u32));
let o = other.clone();
let or = other_runs.clone();
let _other_handle = effect(move || {
or.set(or.get() + 1);
o.get();
});
assert_eq!(other_runs.get(), 1, "the fresh effect ran on registration");
other.set(1);
assert_eq!(
other_runs.get(),
2,
"the scheduler is still dead: `IS_FLUSHING` was stranded by the unwind, \
so this `set` enqueued the effect and nothing ever flushed it"
);
}
#[test]
fn a_panicking_effect_body_does_not_strand_current_effect() {
let (sig, runs, _panicking) = panicking_on_second_run();
let outcome = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
sig.set(1);
}));
assert!(
outcome.is_err(),
"the body was supposed to panic during the flush"
);
IS_FLUSHING.with(|f| f.set(false));
let fresh = Rc::new(Signal::new(0u32));
assert_eq!(
fresh.get(),
0,
"a read outside any effect must not be attributed to a dead effect"
);
fresh.set(1);
assert_eq!(
runs.get(),
2,
"a stranded `CURRENT_EFFECT` attributed an out-of-effect read to the \
dead effect, which then ran again: {} runs",
runs.get()
);
}
#[test]
fn an_initial_body_that_writes_the_signal_it_reads_does_not_panic() {
let sig = Rc::new(Signal::new(0u32));
let runs = Rc::new(Cell::new(0u32));
let wrote = Rc::new(Cell::new(false));
let s = sig.clone();
let r = runs.clone();
let w = wrote.clone();
let handle = effect(move || {
r.set(r.get() + 1);
let v = s.get();
if !w.replace(true) {
s.set(v + 1);
}
});
assert_eq!(runs.get(), 1, "the initial run completed without panicking");
assert_eq!(sig.get(), 1, "the initial run's write landed");
sig.set(5);
assert_eq!(runs.get(), 2, "the effect must still be subscribed");
assert_eq!(sig.get(), 5);
drop(handle);
}
#[test]
fn dead_subscribers_are_reclaimed_when_a_signal_is_never_set_again() {
const BATCH: usize = SUBSCRIBER_SWEEP_AT * 2;
const ROUNDS: usize = 5;
let sig = Rc::new(Signal::new(0u32));
for round in 0..ROUNDS {
{
let mut handles = Vec::with_capacity(BATCH);
for _ in 0..BATCH {
let s = sig.clone();
handles.push(effect(move || {
s.get();
}));
}
drop(handles);
}
{
let s = sig.clone();
let last = effect(move || {
s.get();
});
drop(last);
}
let dead = sig
.subscribers
.borrow()
.iter()
.filter(|w| w.strong_count() == 0)
.count();
assert!(
dead <= SUBSCRIBER_SWEEP_AT,
"round {round}: {dead} dead subscribers accumulated with no bound, \
so the bounded sweep is not reclaiming them (threshold {SUBSCRIBER_SWEEP_AT})"
);
}
}
#[test]
fn batch_defers_the_flush_until_the_outermost_scope_closes() {
let sig = Rc::new(Signal::new(0u32));
let runs = Rc::new(Cell::new(0u32));
let s = sig.clone();
let r = runs.clone();
let _handle = effect(move || {
r.set(r.get() + 1);
s.get();
});
assert_eq!(runs.get(), 1, "the effect runs once on registration");
batch(|| {
sig.set(1);
sig.set(2);
assert_eq!(
runs.get(),
1,
"an effect ran mid-batch: the flush was not deferred"
);
batch(|| {
sig.set(3);
});
assert_eq!(
runs.get(),
1,
"a nested batch flushed early: only the outermost close may flush"
);
});
assert_eq!(sig.get(), 3, "batched sets must all land, last one wins");
assert_eq!(
runs.get(),
2,
"three batched sets must cause exactly one re-run, got {}",
runs.get()
);
sig.set(4);
assert_eq!(runs.get(), 3);
}
}