Reduce allocations and dynamic dispatch in effects and memos

packages/sycamore-reactive/src/effect.rs

@@ -1,6 +1,4 @@
-//! Side effects.
-
-use std::panic::Location;
+use std::{panic::Location, ptr};
 
 use hashbrown::HashSet;
 
@@ -17,7 +15,7 @@ thread_local! {
 /// this struct.
 pub(crate) struct EffectState<'a> {
     /// The callback when the effect is re-executed.
-    cb: Rc<RefCell<dyn FnMut() + 'a>>,
+    cb: Rc<RefCell<EffectCallback<'a, dyn FnMut() + 'a>>>,
     /// A list of dependencies that can trigger this effect.
     dependencies: HashSet<EffectDependency>,
 }
@@ -55,6 +53,65 @@ impl<'a> EffectState<'a> {
     }
 }
 
+pub(crate) struct EffectCallback<'a, T: ?Sized>
+where
+    T: FnMut() + 'a,
+{
+    cx: Scope<'a>,
+    effect: &'a RefCell<Option<EffectState<'a>>>,
+    f: T,
+}
+
+impl<'a> EffectCallback<'a, dyn FnMut() + 'a> {
+    pub fn run(&mut self) {
+        let cx = self.cx;
+        let effect = self.effect;
+        let f = &mut self.f;
+        EFFECTS.with(|effects| {
+            // Record initial effect stack length to verify that it is the same after.
+            let initial_effect_stack_len = effects.borrow().len();
+
+            // Take effect out.
+            let mut effect_borrow_mut = effect.borrow_mut();
+            let tmp_effect = effect_borrow_mut.as_mut().unwrap();
+            tmp_effect.clear_dependencies();
+
+            // Push the effect onto the effect stack so that it is visible by signals.
+            effects
+                .borrow_mut()
+                .push((tmp_effect as *mut EffectState<'a>).cast::<EffectState<'static>>());
+            // We want to drop-lock the parent scope during the call to `f` so that the child
+            // scope remains alive during the whole call.
+            cx.raw.inner.borrow_mut().lock_drop = true;
+            // Now we can call the user-provided function.
+            f();
+            cx.raw.inner.borrow_mut().lock_drop = false;
+            // Pop the effect from the effect stack.
+            // This ends the mutable borrow of `tmp_effect`.
+            effects.borrow_mut().pop().unwrap();
+
+            // The raw pointer pushed onto `effects` is dead and can no longer be accessed.
+            // We can now access `effect` directly again.
+
+            // For all the signals collected by the EffectState, we need to add backlinks from
+            // the signal to the effect, so that updating the signal will trigger the effect.
+            for emitter in &tmp_effect.dependencies {
+                // The SignalEmitter might have been destroyed between when the signal was
+                // accessed and now.
+                if let Some(emitter) = emitter.0.upgrade() {
+                    // SAFETY: When the effect is destroyed or when the emitter is dropped,
+                    // this link will be destroyed to prevent dangling references.
+                    emitter.subscribe(Rc::downgrade(unsafe {
+                        std::mem::transmute(&tmp_effect.cb)
+                    }));
+                }
+            }
+
+            debug_assert_eq!(effects.borrow().len(), initial_effect_stack_len);
+        });
+    }
+}
+
 /// Creates an effect on signals used inside the effect closure.
 ///
 /// # Example
@@ -71,68 +128,62 @@ impl<'a> EffectState<'a> {
 /// # });
 /// ```
 pub fn create_effect<'a>(cx: Scope<'a>, f: impl FnMut() + 'a) {
-    _create_effect(cx, Box::new(f))
+    let effect = _make_effect(cx);
+    let cb = Rc::new(RefCell::new(EffectCallback { cx, effect, f }));
+
+    _create_effect(effect, cb)
 }
 
-/// Internal implementation for `create_effect`. Use dynamic dispatch to reduce code-bloat.
-fn _create_effect<'a>(cx: Scope<'a>, mut f: Box<(dyn FnMut() + 'a)>) {
+fn _make_effect<'a>(cx: Scope<'a>) -> &'a RefCell<Option<EffectState<'a>>> {
     // SAFETY: We do not access the scope in the Drop implementation for EffectState.
-    let effect = unsafe { create_ref_unsafe(cx, RefCell::new(None::<EffectState<'a>>)) };
-    let cb = Rc::new(RefCell::new({
-        move || {
-            EFFECTS.with(|effects| {
-                // Record initial effect stack length to verify that it is the same after.
-                let initial_effect_stack_len = effects.borrow().len();
-
-                // Take effect out.
-                let mut effect_borrow_mut = effect.borrow_mut();
-                let tmp_effect = effect_borrow_mut.as_mut().unwrap();
-                tmp_effect.clear_dependencies();
-
-                // Push the effect onto the effect stack so that it is visible by signals.
-                effects
-                    .borrow_mut()
-                    .push((tmp_effect as *mut EffectState<'a>).cast::<EffectState<'static>>());
-                // We want to drop-lock the parent scope during the call to `f` so that the child
-                // scope remains alive during the whole call.
-                cx.raw.inner.borrow_mut().lock_drop = true;
-                // Now we can call the user-provided function.
-                f();
-                cx.raw.inner.borrow_mut().lock_drop = false;
-                // Pop the effect from the effect stack.
-                // This ends the mutable borrow of `tmp_effect`.
-                effects.borrow_mut().pop().unwrap();
-
-                // The raw pointer pushed onto `effects` is dead and can no longer be accessed.
-                // We can now access `effect` directly again.
-
-                // For all the signals collected by the EffectState, we need to add backlinks from
-                // the signal to the effect, so that updating the signal will trigger the effect.
-                for emitter in &tmp_effect.dependencies {
-                    // The SignalEmitter might have been destroyed between when the signal was
-                    // accessed and now.
-                    if let Some(emitter) = emitter.0.upgrade() {
-                        // SAFETY: When the effect is destroyed or when the emitter is dropped,
-                        // this link will be destroyed to prevent dangling references.
-                        emitter.subscribe(Rc::downgrade(unsafe {
-                            std::mem::transmute(&tmp_effect.cb)
-                        }));
-                    }
-                }
+    unsafe { create_ref_unsafe(cx, RefCell::new(None::<EffectState<'a>>)) }
+}
 
-                debug_assert_eq!(effects.borrow().len(), initial_effect_stack_len);
-            });
+/// Creates an effect on signals used inside the effect closure, returning the value from the initial call to the closure.
+///
+/// # Example
+/// ```
+/// # use sycamore_reactive::*;
+/// # create_scope_immediate(|cx| {
+/// let state = create_signal(cx, 0);
+///
+/// println!("Init: {}", create_effect_return_init(cx, || {
+///     let v = state.get();
+///     println!("State changed. New state value = {}", v);
+///     v
+/// })); // Prints "State changed. New state value = 0" and "Init: 0"
+///
+/// state.set(1); // Prints "State changed. New state value = 1"
+/// # });
+/// ```
+pub fn create_effect_return_init<'a, T: 'a>(cx: Scope<'a>, mut f: impl FnMut() -> T + 'a) -> T {
+    // use an option so we don't leak the value in case _create_effect panics after running f()
+    let mut init = None;
+    let mut init_ptr: *mut Option<T> = &mut init as *mut _;
+    create_effect(cx, move || {
+        let value = f();
+        if !init_ptr.is_null() {
+            // SAFETY: since value_ptr is not null and here we set it to null, this must be the initial execution inside _create_effect, so the pointer is valid
+            unsafe { ptr::write(init_ptr, Some(value)) };
+            init_ptr = ptr::null_mut();
         }
-    }));
+    });
+    init.unwrap()
+}
 
+/// Internal implementation for `create_effect`. Use dynamic dispatch to reduce code-bloat.
+fn _create_effect<'a>(
+    effect: &RefCell<Option<EffectState<'a>>>,
+    cb: Rc<RefCell<EffectCallback<'a, dyn FnMut() + 'a>>>,
+) {
     // Initialize initial effect state.
     *effect.borrow_mut() = Some(EffectState {
         cb: cb.clone(),
         dependencies: HashSet::new(),
     });
 
     // Initial callback call to get everything started.
-    cb.borrow_mut()();
+    cb.borrow_mut().run();
 }
 
 /// Creates an effect on signals used inside the effect closure.

packages/sycamore-reactive/src/memo.rs

@@ -1,7 +1,5 @@
 //! Derived and computed data.
 
-use std::cell::Cell;
-
 use crate::*;
 
 /// Creates a memoized computation from some signals.
@@ -90,24 +88,24 @@ pub fn create_selector_with<'a, U: 'static>(
     mut f: impl FnMut() -> U + 'a,
     eq_f: impl Fn(&U, &U) -> bool + 'a,
 ) -> &'a ReadSignal<U> {
-    let signal: Rc<Cell<Option<&Signal<U>>>> = Rc::new(Cell::new(None));
+    let mut signal_opt: Option<&Signal<U>> = None;
 
-    create_effect(cx, {
-        let signal = Rc::clone(&signal);
+    create_effect_return_init(cx, {
         move || {
             let new = f();
-            if let Some(signal) = signal.get() {
+            if let Some(signal) = signal_opt {
                 // Check if new value is different from old value.
                 if !eq_f(&new, &*signal.get_untracked()) {
-                    signal.set(new)
+                    signal.set(new);
                 }
+                signal
             } else {
-                signal.set(Some(create_signal(cx, new)))
+                let signal = create_signal(cx, new);
+                signal_opt = Some(signal);
+                signal
             }
         }
-    });
-
-    signal.get().unwrap()
+    })
 }
 
 /// An alternative to [`create_signal`] that uses a reducer to get the next

packages/sycamore-reactive/src/signal.rs

@@ -4,11 +4,11 @@ use std::fmt::{Debug, Display, Formatter};
 use std::hash::Hash;
 use std::ops::{AddAssign, Deref, DerefMut, DivAssign, MulAssign, SubAssign};
 
-use crate::effect::EFFECTS;
+use crate::effect::{EffectCallback, EFFECTS};
 use crate::*;
 
-type WeakEffectCallback = Weak<RefCell<dyn FnMut()>>;
-type EffectCallbackPtr = *const RefCell<dyn FnMut()>;
+type WeakEffectCallback = Weak<RefCell<EffectCallback<'static, dyn FnMut()>>>;
+type EffectCallbackPtr = *const RefCell<EffectCallback<'static, dyn FnMut()>>;
 
 pub(crate) type SignalEmitterInner = RefCell<IndexMap<EffectCallbackPtr, WeakEffectCallback>>;
 
@@ -77,7 +77,7 @@ impl SignalEmitter {
             // subscriber might have already been destroyed in the case of nested effects.
             if let Some(callback) = subscriber.upgrade() {
                 // Call the callback.
-                callback.borrow_mut()();
+                callback.borrow_mut().run();
             }
         }
     }