From 4b0c13615ce0ee8eb765bfae6219a7e81cb7035c Mon Sep 17 00:00:00 2001
From: Anton Korzunov <thekashey@gmail.com>
Date: Fri, 17 Nov 2023 18:16:39 +1100
Subject: [PATCH 1/2] fix: correct nested batch behavior. Solves #217

---
 src/utils/__tests__/batched-updates.test.js | 11 +++++++++--
 src/utils/batched-updates.js                |  7 ++++---
 2 files changed, 13 insertions(+), 5 deletions(-)

diff --git a/src/utils/__tests__/batched-updates.test.js b/src/utils/__tests__/batched-updates.test.js
index eeade58..82896a4 100644
--- a/src/utils/__tests__/batched-updates.test.js
+++ b/src/utils/__tests__/batched-updates.test.js
@@ -56,13 +56,20 @@ describe('batch', () => {
     const child = jest.fn().mockReturnValue(null);
     render(<TestComponent>{child}</TestComponent>);
     const update = child.mock.calls[0][0];
-    act(() => update());
+    act(() => {
+      update();
+      update();
+      update();
+    });
+
+    // nothing should be yet called
+    expect(child.mock.calls[2]).toEqual(undefined);
 
     // scheduler uses timeouts on non-browser envs
     await act(() => new Promise((r) => setTimeout(r, 10)));
 
     // assertion no longer relevant with React 18+
-    expect(child.mock.calls[2]).toEqual([expect.any(Function), 1, 1]);
+    expect(child.mock.calls[2]).toEqual([expect.any(Function), 3, 1]);
 
     supportsMock.mockRestore();
   });
diff --git a/src/utils/batched-updates.js b/src/utils/batched-updates.js
index 5fcf8b8..df346f7 100644
--- a/src/utils/batched-updates.js
+++ b/src/utils/batched-updates.js
@@ -8,7 +8,7 @@ import {
 import defaults from '../defaults';
 import supports from './supported-features';
 
-let isInsideBatchedSchedule = false;
+let isInsideBatchedSchedule = 0;
 
 export function batch(fn) {
   // if we are in node/tests or nested schedule
@@ -20,11 +20,12 @@ export function batch(fn) {
     return unstable_batchedUpdates(fn);
   }
 
-  isInsideBatchedSchedule = true;
+  isInsideBatchedSchedule = 0;
   // Use ImmediatePriority as it has -1ms timeout
   // https://github.com/facebook/react/blob/main/packages/scheduler/src/forks/Scheduler.js#L65
   return scheduleCallback(ImmediatePriority, function scheduleBatchedUpdates() {
+    isInsideBatchedSchedule++;
     unstable_batchedUpdates(fn);
-    isInsideBatchedSchedule = false;
+    isInsideBatchedSchedule--;
   });
 }

From e89bef94fd6b678777a610f5b9803464aadb3567 Mon Sep 17 00:00:00 2001
From: Anton Korzunov <thekashey@gmail.com>
Date: Mon, 11 Dec 2023 20:33:09 +1100
Subject: [PATCH 2/2] revamp batching implementation to batch in longer
 sequences

---
 src/utils/batched-updates.js | 37 ++++++++++++++++++++++--------------
 1 file changed, 23 insertions(+), 14 deletions(-)

diff --git a/src/utils/batched-updates.js b/src/utils/batched-updates.js
index df346f7..6ac286a 100644
--- a/src/utils/batched-updates.js
+++ b/src/utils/batched-updates.js
@@ -8,24 +8,33 @@ import {
 import defaults from '../defaults';
 import supports from './supported-features';
 
-let isInsideBatchedSchedule = 0;
+let batchedScheduled = false;
+
+let batched = [];
+
+const executeBatched = () => {
+  unstable_batchedUpdates(() => {
+    while (batched.length) {
+      const currentBatched = batched;
+      batched = [];
+      currentBatched.forEach((fn) => fn());
+    }
+    // important to reset it before exiting this function
+    // as React will dump everything right after
+    batchedScheduled = false;
+  });
+};
 
 export function batch(fn) {
   // if we are in node/tests or nested schedule
-  if (
-    !defaults.batchUpdates ||
-    !supports.scheduling() ||
-    isInsideBatchedSchedule
-  ) {
+  if (!defaults.batchUpdates || !supports.scheduling()) {
     return unstable_batchedUpdates(fn);
   }
 
-  isInsideBatchedSchedule = 0;
-  // Use ImmediatePriority as it has -1ms timeout
-  // https://github.com/facebook/react/blob/main/packages/scheduler/src/forks/Scheduler.js#L65
-  return scheduleCallback(ImmediatePriority, function scheduleBatchedUpdates() {
-    isInsideBatchedSchedule++;
-    unstable_batchedUpdates(fn);
-    isInsideBatchedSchedule--;
-  });
+  batched.push(fn);
+
+  if (!batchedScheduled) {
+    batchedScheduled = true;
+    return scheduleCallback(ImmediatePriority, executeBatched);
+  }
 }