From d552d7e87bfce082fbdbe39118925ae265bc94cf Mon Sep 17 00:00:00 2001
From: MinusGix <MinusGix@gmail.com>
Date: Wed, 1 Nov 2023 03:23:53 -0500
Subject: [PATCH] Add batch (#130)

---
 reactive/Cargo.toml     |  1 +
 reactive/src/effect.rs  | 13 +++++++++++++
 reactive/src/lib.rs     |  2 +-
 reactive/src/runtime.rs | 32 ++++++++++++++++++++++++++++++--
 reactive/src/signal.rs  | 10 ++++++++++
 5 files changed, 55 insertions(+), 3 deletions(-)

diff --git a/reactive/Cargo.toml b/reactive/Cargo.toml
index 9bbf2bd7..ac8a3d17 100644
--- a/reactive/Cargo.toml
+++ b/reactive/Cargo.toml
@@ -5,3 +5,4 @@ edition = "2021"
 license.workspace = true
 
 [dependencies]
+smallvec = "1.10.0"
diff --git a/reactive/src/effect.rs b/reactive/src/effect.rs
index 2d92b8ff..4557e9fc 100644
--- a/reactive/src/effect.rs
+++ b/reactive/src/effect.rs
@@ -69,6 +69,19 @@ pub fn untrack<T>(f: impl FnOnce() -> T) -> T {
     result
 }
 
+pub fn batch<T>(f: impl FnOnce() -> T) -> T {
+    RUNTIME.with(|runtime| {
+        runtime.batching.set(true);
+    });
+    let result = f();
+    RUNTIME.with(|runtime| {
+        runtime.batching.set(false);
+        runtime.run_pending_effects();
+    });
+
+    result
+}
+
 pub(crate) fn run_effect(effect: Rc<dyn EffectTrait>) {
     let effect_id = effect.id();
     effect_id.dispose();
diff --git a/reactive/src/lib.rs b/reactive/src/lib.rs
index 7b715bd6..6f961fc3 100644
--- a/reactive/src/lib.rs
+++ b/reactive/src/lib.rs
@@ -8,7 +8,7 @@ mod signal;
 mod trigger;
 
 pub use context::{provide_context, use_context};
-pub use effect::{create_effect, untrack};
+pub use effect::{batch, create_effect, untrack};
 pub use memo::{create_memo, Memo};
 pub use scope::{as_child_of_current_scope, with_scope, Scope};
 pub use signal::{create_rw_signal, create_signal, ReadSignal, RwSignal, WriteSignal};
diff --git a/reactive/src/runtime.rs b/reactive/src/runtime.rs
index 2e03ed2c..e9cebe5e 100644
--- a/reactive/src/runtime.rs
+++ b/reactive/src/runtime.rs
@@ -1,11 +1,17 @@
 use std::{
     any::{Any, TypeId},
-    cell::RefCell,
+    cell::{Cell, RefCell},
     collections::{HashMap, HashSet},
     rc::Rc,
 };
 
-use crate::{effect::EffectTrait, id::Id, signal::Signal};
+use smallvec::SmallVec;
+
+use crate::{
+    effect::{run_effect, EffectTrait},
+    id::Id,
+    signal::Signal,
+};
 
 thread_local! {
     pub(crate) static RUNTIME: Runtime = Runtime::new();
@@ -19,6 +25,8 @@ pub(crate) struct Runtime {
     pub(crate) children: RefCell<HashMap<Id, HashSet<Id>>>,
     pub(crate) signals: RefCell<HashMap<Id, Signal>>,
     pub(crate) contexts: RefCell<HashMap<TypeId, Box<dyn Any>>>,
+    pub(crate) batching: Cell<bool>,
+    pub(crate) pending_effects: RefCell<SmallVec<[Rc<dyn EffectTrait>; 10]>>,
 }
 
 impl Default for Runtime {
@@ -35,6 +43,26 @@ impl Runtime {
             children: RefCell::new(HashMap::new()),
             signals: Default::default(),
             contexts: Default::default(),
+            batching: Cell::new(false),
+            pending_effects: RefCell::new(SmallVec::new()),
+        }
+    }
+
+    pub(crate) fn add_pending_effect(&self, effect: Rc<dyn EffectTrait>) {
+        let has_effect = self
+            .pending_effects
+            .borrow()
+            .iter()
+            .any(|e| e.id() == effect.id());
+        if !has_effect {
+            self.pending_effects.borrow_mut().push(effect);
+        }
+    }
+
+    pub(crate) fn run_pending_effects(&self) {
+        let pending_effects = self.pending_effects.take();
+        for effect in pending_effects {
+            run_effect(effect);
         }
     }
 }
diff --git a/reactive/src/signal.rs b/reactive/src/signal.rs
index 3f85da69..1eeca89d 100644
--- a/reactive/src/signal.rs
+++ b/reactive/src/signal.rs
@@ -347,6 +347,16 @@ impl Signal {
     }
 
     pub(crate) fn run_effects(&self) {
+        // If we are batching then add it as a pending effect
+        if RUNTIME.with(|r| r.batching.get()) {
+            RUNTIME.with(|r| {
+                for (_, subscriber) in self.subscribers() {
+                    r.add_pending_effect(subscriber);
+                }
+            });
+            return;
+        }
+
         for (_, subscriber) in self.subscribers() {
             run_effect(subscriber);
         }