improve signal copy runtime

packages/signals/src/copy.rs

@@ -0,0 +1,301 @@
+use std::{
+    cell::{Cell, Ref, RefCell, RefMut},
+    marker::PhantomData,
+    rc::Rc,
+};
+
+use bumpalo::Bump;
+
+/// # Example
+///
+/// ```compile_fail
+/// let data = String::from("hello world");
+/// let store = Store::default();
+/// let owner = store.owner();
+/// let key = owner.insert(&data);
+/// drop(data);
+/// assert_eq!(*key.read(), "hello world");
+/// ```
+#[allow(unused)]
+fn compile_fail() {}
+
+#[test]
+fn reused() {
+    let store = Store::default();
+    let first_ptr;
+    {
+        let owner = store.owner();
+        first_ptr = owner.insert(1).raw.data.as_ptr();
+        drop(owner);
+    }
+    {
+        let owner = store.owner();
+        let second_ptr = owner.insert(1234).raw.data.as_ptr();
+        assert_eq!(first_ptr, second_ptr);
+        drop(owner);
+    }
+}
+
+#[test]
+fn leaking_is_ok() {
+    let data = String::from("hello world");
+    let store = Store::default();
+    let key;
+    {
+        // create an owner
+        let owner = store.owner();
+        // insert data into the store
+        key = owner.insert(data);
+        // don't drop the owner
+        std::mem::forget(owner);
+    }
+    assert_eq!(key.try_read().as_deref(), Some(&"hello world".to_string()));
+}
+
+#[test]
+fn drops() {
+    let data = String::from("hello world");
+    let store = Store::default();
+    let key;
+    {
+        // create an owner
+        let owner = store.owner();
+        // insert data into the store
+        key = owner.insert(data);
+        // drop the owner
+    }
+    assert!(key.try_read().is_none());
+}
+
+#[test]
+fn works() {
+    let store = Store::default();
+    let owner = store.owner();
+    let key = owner.insert(1);
+
+    assert_eq!(*key.read(), 1);
+}
+
+#[test]
+fn insert_while_reading() {
+    let store = Store::default();
+    let owner = store.owner();
+    let key;
+    {
+        let data: String = "hello world".to_string();
+        key = owner.insert(data);
+    }
+    let value = key.read();
+    owner.insert(&1);
+    assert_eq!(*value, "hello world");
+}
+
+#[test]
+#[should_panic]
+fn panics() {
+    let store = Store::default();
+    let owner = store.owner();
+    let key = owner.insert(1);
+    drop(owner);
+
+    assert_eq!(*key.read(), 1);
+}
+
+#[test]
+fn fuzz() {
+    fn maybe_owner_scope(
+        store: &Store,
+        valid_keys: &mut Vec<CopyHandle<String>>,
+        invalid_keys: &mut Vec<CopyHandle<String>>,
+        path: &mut Vec<u8>,
+    ) {
+        let branch_cutoff = 5;
+        let children = if path.len() < branch_cutoff {
+            rand::random::<u8>() % 4
+        } else {
+            rand::random::<u8>() % 2
+        };
+
+        for i in 0..children {
+            let owner = store.owner();
+            let key = owner.insert(format!("hello world {path:?}"));
+            valid_keys.push(key);
+            path.push(i);
+            // read all keys
+            println!("{:?}", path);
+            for key in valid_keys.iter() {
+                let value = key.read();
+                println!("{:?}", value);
+                assert!(value.starts_with("hello world"));
+            }
+            #[cfg(debug_assertions)]
+            for key in invalid_keys.iter() {
+                assert!(!key.validate());
+            }
+            maybe_owner_scope(store, valid_keys, invalid_keys, path);
+            invalid_keys.push(valid_keys.pop().unwrap());
+            path.pop();
+        }
+    }
+
+    for _ in 0..10 {
+        let store = Store::default();
+        maybe_owner_scope(&store, &mut Vec::new(), &mut Vec::new(), &mut Vec::new());
+    }
+}
+
+pub struct CopyHandle<T> {
+    raw: MemoryLocation,
+    #[cfg(debug_assertions)]
+    generation: u32,
+    _marker: PhantomData<T>,
+}
+
+impl<T: 'static> CopyHandle<T> {
+    #[inline(always)]
+    fn validate(&self) -> bool {
+        #[cfg(debug_assertions)]
+        {
+            self.raw.generation.get() == self.generation
+        }
+        #[cfg(not(debug_assertions))]
+        {
+            true
+        }
+    }
+
+    pub fn try_read(&self) -> Option<Ref<'_, T>> {
+        self.validate()
+            .then(|| {
+                Ref::filter_map(self.raw.data.borrow(), |any| {
+                    any.as_ref()?.downcast_ref::<T>()
+                })
+                .ok()
+            })
+            .flatten()
+    }
+
+    pub fn read(&self) -> Ref<'_, T> {
+        self.try_read().unwrap()
+    }
+
+    pub fn try_write(&self) -> Option<RefMut<'_, T>> {
+        self.validate()
+            .then(|| {
+                RefMut::filter_map(self.raw.data.borrow_mut(), |any| {
+                    any.as_mut()?.downcast_mut::<T>()
+                })
+                .ok()
+            })
+            .flatten()
+    }
+
+    pub fn write(&self) -> RefMut<'_, T> {
+        self.try_write().unwrap()
+    }
+}
+
+impl<T> Copy for CopyHandle<T> {}
+
+impl<T> Clone for CopyHandle<T> {
+    fn clone(&self) -> Self {
+        *self
+    }
+}
+
+#[derive(Clone, Copy)]
+struct MemoryLocation {
+    data: &'static RefCell<Option<Box<dyn std::any::Any>>>,
+    #[cfg(debug_assertions)]
+    generation: &'static Cell<u32>,
+}
+
+impl MemoryLocation {
+    #[allow(unused)]
+    fn drop(&self) {
+        let old = self.data.borrow_mut().take();
+        #[cfg(debug_assertions)]
+        if old.is_some() {
+            let new_generation = self.generation.get() + 1;
+            self.generation.set(new_generation);
+        }
+    }
+
+    fn replace<T: 'static>(&mut self, value: T) -> CopyHandle<T> {
+        let mut inner_mut = self.data.borrow_mut();
+
+        let raw = Box::new(value);
+        let old = inner_mut.replace(raw);
+        assert!(old.is_none());
+        CopyHandle {
+            raw: *self,
+            #[cfg(debug_assertions)]
+            generation: self.generation.get(),
+            _marker: PhantomData,
+        }
+    }
+}
+
+#[derive(Clone)]
+pub struct Store {
+    bump: &'static Bump,
+    recycled: Rc<RefCell<Vec<MemoryLocation>>>,
+}
+
+impl Default for Store {
+    fn default() -> Self {
+        Self {
+            bump: Box::leak(Box::new(Bump::new())),
+            recycled: Default::default(),
+        }
+    }
+}
+
+impl Store {
+    fn recycle(&self, location: MemoryLocation) {
+        location.drop();
+        self.recycled.borrow_mut().push(location);
+    }
+
+    fn claim(&self) -> MemoryLocation {
+        if let Some(location) = self.recycled.borrow_mut().pop() {
+            location
+        } else {
+            let data: &'static RefCell<_> = self.bump.alloc(RefCell::new(None));
+            MemoryLocation {
+                data,
+                #[cfg(debug_assertions)]
+                generation: self.bump.alloc(Cell::new(0)),
+            }
+        }
+    }
+
+    pub fn owner(&self) -> Owner {
+        Owner {
+            store: self.clone(),
+            owned: Default::default(),
+        }
+    }
+}
+
+pub struct Owner {
+    store: Store,
+    owned: Rc<RefCell<Vec<MemoryLocation>>>,
+}
+
+impl Owner {
+    pub fn insert<T: 'static>(&self, value: T) -> CopyHandle<T> {
+        let mut location = self.store.claim();
+        let key = location.replace(value);
+        self.owned.borrow_mut().push(location);
+        key
+    }
+}
+
+impl Drop for Owner {
+    fn drop(&mut self) {
+        for location in self.owned.borrow().iter() {
+            self.store.recycle(*location)
+        }
+    }
+}

packages/signals/src/lib.rs

@@ -1,70 +1,65 @@
 use std::{
     cell::{Ref, RefMut},
     fmt::Display,
-    marker::PhantomData,
     ops::{Add, Div, Mul, Sub},
+    sync::Arc,
 };
 
+mod copy;
 mod rt;
-
-use dioxus_core::ScopeState;
 pub use rt::*;
 
-pub fn use_init_signal_rt(cx: &ScopeState) {
-    cx.use_hook(|| {
-        let rt = claim_rt(cx.schedule_update_any());
-        cx.provide_context(rt);
-    });
-}
+use dioxus_core::{
+    prelude::{current_scope_id, schedule_update_any},
+    ScopeId, ScopeState,
+};
 
 pub fn use_signal<T: 'static>(cx: &ScopeState, f: impl FnOnce() -> T) -> Signal<T> {
-    cx.use_hook(|| {
-        let rt: &'static SignalRt = match cx.consume_context() {
-            Some(rt) => rt,
-            None => cx.provide_context(claim_rt(cx.schedule_update_any())),
-        };
-
-        let id = rt.init(f());
-        rt.subscribe(id, cx.scope_id());
-
-        struct SignalHook<T> {
-            signal: Signal<T>,
-        }
-
-        impl<T> Drop for SignalHook<T> {
-            fn drop(&mut self) {
-                self.signal.rt.remove(self.signal.id);
-            }
-        }
-
-        SignalHook {
-            signal: Signal {
-                id,
-                rt,
-                t: PhantomData,
-            },
-        }
-    })
-    .signal
+    *cx.use_hook(|| Signal::new(f()))
 }
 
-pub struct Signal<T> {
-    id: usize,
-    rt: &'static SignalRt,
-    t: PhantomData<T>,
+struct SignalData<T> {
+    subscribers: Vec<ScopeId>,
+    update_any: Arc<dyn Fn(ScopeId)>,
+    value: T,
+}
+
+pub struct Signal<T: 'static> {
+    inner: CopyValue<SignalData<T>>,
 }
 
 impl<T: 'static> Signal<T> {
+    pub fn new(value: T) -> Self {
+        Self {
+            inner: CopyValue::new(SignalData {
+                subscribers: Vec::new(),
+                update_any: schedule_update_any().expect("in a virtual dom"),
+                value,
+            }),
+        }
+    }
+
     pub fn read(&self) -> Ref<T> {
-        self.rt.read(self.id)
+        if let Some(current_scope_id) = current_scope_id() {
+            let mut inner = self.inner.write();
+            if !inner.subscribers.contains(&current_scope_id) {
+                inner.subscribers.push(current_scope_id);
+            }
+        }
+        Ref::map(self.inner.read(), |v| &v.value)
     }
 
     pub fn write(&self) -> RefMut<T> {
-        self.rt.write(self.id)
+        let inner = self.inner.write();
+        for &scope_id in &inner.subscribers {
+            (inner.update_any)(scope_id);
+        }
+
+        RefMut::map(inner, |v| &mut v.value)
     }
 
     pub fn set(&mut self, value: T) {
-        self.rt.set(self.id, value);
+        *self.write() = value;
     }
 
     pub fn with<O>(&self, f: impl FnOnce(&T) -> O) -> O {
@@ -72,23 +67,15 @@ impl<T: 'static> Signal<T> {
         f(&*write)
     }
 
-    pub fn update<O>(&self, _f: impl FnOnce(&mut T) -> O) -> O {
+    pub fn update<O>(&self, f: impl FnOnce(&mut T) -> O) -> O {
         let mut write = self.write();
-        _f(&mut *write)
+        f(&mut *write)
     }
 }
 
 impl<T: Clone + 'static> Signal<T> {
     pub fn get(&self) -> T {
-        self.rt.get(self.id)
-    }
-}
-
-impl<T: Clone + 'static> std::ops::Deref for Signal<T> {
-    type Target = dyn Fn() -> T;
-
-    fn deref(&self) -> &Self::Target {
-        self.rt.getter(self.id)
+        self.read().clone()
     }
 }
 
@@ -102,7 +89,7 @@ impl<T> Copy for Signal<T> {}
 
 impl<T: Display + 'static> Display for Signal<T> {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        self.rt.with::<T, _>(self.id, |v| T::fmt(v, f))
+        self.with(|v| Display::fmt(v, f))
     }
 }
 

packages/signals/src/rt.rs

@@ -1,121 +1,65 @@
-use std::{any::Any, cell::RefCell, sync::Arc};
+use std::cell::{Ref, RefMut};
 
-use dioxus_core::ScopeId;
-use slab::Slab;
+use std::rc::Rc;
 
-thread_local! {
-    // we cannot drop these since any future might be using them
-    static RUNTIMES: RefCell<Vec<&'static SignalRt>> = RefCell::new(Vec::new());
-}
+use dioxus_core::prelude::{consume_context, provide_root_context};
 
-/// Provide the runtime for signals
-///
-/// This will reuse dead runtimes
-pub fn claim_rt(update_any: Arc<dyn Fn(ScopeId)>) -> &'static SignalRt {
-    RUNTIMES.with(|runtimes| {
-        if let Some(rt) = runtimes.borrow_mut().pop() {
-            return rt;
+use crate::copy::{CopyHandle, Owner, Store};
+
+fn current_store() -> Store {
+    match consume_context() {
+        Some(rt) => rt,
+        None => {
+            let store = Store::default();
+            provide_root_context(store).expect("in a virtual dom")
         }
-
-        Box::leak(Box::new(SignalRt {
-            signals: RefCell::new(Slab::new()),
-            update_any,
-        }))
-    })
+    }
 }
 
-/// Push this runtime into the global runtime list
-pub fn reclam_rt(_rt: &'static SignalRt) {
-    RUNTIMES.with(|runtimes| {
-        runtimes.borrow_mut().push(_rt);
-    });
+fn current_owner() -> Rc<Owner> {
+    match consume_context() {
+        Some(rt) => rt,
+        None => {
+            let owner = Rc::new(current_store().owner());
+            provide_root_context(owner).expect("in a virtual dom")
+        }
+    }
 }
 
-pub struct SignalRt {
-    pub(crate) signals: RefCell<Slab<Inner>>,
-    pub(crate) update_any: Arc<dyn Fn(ScopeId)>,
+impl<T> Copy for CopyValue<T> {}
+
+impl<T> Clone for CopyValue<T> {
+    fn clone(&self) -> Self {
+        *self
+    }
 }
 
-impl SignalRt {
-    pub fn init<T: 'static>(&'static self, val: T) -> usize {
-        self.signals.borrow_mut().insert(Inner {
-            value: Box::new(val),
-            subscribers: Vec::new(),
-            getter: None,
-        })
-    }
+pub struct CopyValue<T: 'static> {
+    pub value: CopyHandle<T>,
+}
 
-    pub fn subscribe(&self, id: usize, subscriber: ScopeId) {
-        self.signals.borrow_mut()[id].subscribers.push(subscriber);
-    }
+impl<T: 'static> CopyValue<T> {
+    pub fn new(value: T) -> Self {
+        let owner = current_owner();
 
-    pub fn get<T: Clone + 'static>(&self, id: usize) -> T {
-        self.signals.borrow()[id]
-            .value
-            .downcast_ref::<T>()
-            .cloned()
-            .unwrap()
-    }
-
-    pub fn set<T: 'static>(&self, id: usize, value: T) {
-        let mut signals = self.signals.borrow_mut();
-        let inner = &mut signals[id];
-        inner.value = Box::new(value);
-
-        for subscriber in inner.subscribers.iter() {
-            (self.update_any)(*subscriber);
+        Self {
+            value: owner.insert(value),
         }
     }
 
-    pub fn remove(&self, id: usize) {
-        self.signals.borrow_mut().remove(id);
+    pub fn try_read(&self) -> Option<Ref<'_, T>> {
+        self.value.try_read()
     }
 
-    pub fn with<T: 'static, O>(&self, id: usize, f: impl FnOnce(&T) -> O) -> O {
-        let signals = self.signals.borrow();
-        let inner = &signals[id];
-        let inner = inner.value.downcast_ref::<T>().unwrap();
-        f(inner)
+    pub fn read(&self) -> Ref<'_, T> {
+        self.value.read()
     }
 
-    pub(crate) fn read<T: 'static>(&self, id: usize) -> std::cell::Ref<T> {
-        let signals = self.signals.borrow();
-        std::cell::Ref::map(signals, |signals| {
-            signals[id].value.downcast_ref::<T>().unwrap()
-        })
+    pub fn try_write(&self) -> Option<RefMut<'_, T>> {
+        self.value.try_write()
     }
 
-    pub(crate) fn write<T: 'static>(&self, id: usize) -> std::cell::RefMut<T> {
-        let signals = self.signals.borrow_mut();
-        std::cell::RefMut::map(signals, |signals| {
-            signals[id].value.downcast_mut::<T>().unwrap()
-        })
-    }
-
-    pub(crate) fn getter<T: 'static + Clone>(&self, id: usize) -> &dyn Fn() -> T {
-        let mut signals = self.signals.borrow_mut();
-        let inner = &mut signals[id];
-        let r = inner.getter.as_mut();
-
-        if r.is_none() {
-            let rt = self;
-            let r = move || rt.get::<T>(id);
-            let getter: Box<dyn Fn() -> T> = Box::new(r);
-            let getter: Box<dyn Fn()> = unsafe { std::mem::transmute(getter) };
-
-            inner.getter = Some(getter);
-        }
-
-        let r = inner.getter.as_ref().unwrap();
-
-        unsafe { std::mem::transmute::<&dyn Fn(), &dyn Fn() -> T>(r) }
+    pub fn write(&self) -> RefMut<'_, T> {
+        self.value.write()
     }
 }
-
-pub(crate) struct Inner {
-    pub value: Box<dyn Any>,
-    pub subscribers: Vec<ScopeId>,
-
-    // todo: this has a soundness hole in it that you might not run into
-    pub getter: Option<Box<dyn Fn()>>,
-}