[FIX] utils: fix calls to batched callback from within the callback

Previously, calling the batched function from within the callback being
batched would fail as it would be treated as part of the same batch.
This commit fixes that by scheduling the reset of the "called" flag
before calling the callback. This means that all microtasks that were
already in the microtask queue when a batch is about to run are treated
as part of the batch, and all microtasks that will be added by the
callback are not.

src/utils.ts

@@ -15,11 +15,13 @@ export function batched(callback: Callback): Callback {
     await Promise.resolve();
     if (!called) {
       called = true;
-      callback();
       // wait for all calls in this microtick to fall through before resetting "called"
-      // so that only the first call to the batched function calls the original callback
-      await Promise.resolve();
-      called = false;
+      // so that only the first call to the batched function calls the original callback.
+      // Schedule this before calling the callback so that calls to the batched function
+      // within the callback will proceed only after resetting called to false, and have
+      // a chance to execute the callback again
+      Promise.resolve().then(() => (called = false));
+      callback();
     }
   };
 }

tests/reactivity.test.ts

@@ -208,6 +208,46 @@ describe("Reactivity", () => {
     expect(n).toBe(1); // two operations but only one notification
   });
 
+  test("batched: modifying the reactive in the callback doesn't break reactivity", async () => {
+    let n = 0;
+    let obj = { a: 1 };
+    const state = createReactive(
+      obj,
+      batched(() => {
+        state.a; // subscribe to a
+        state.a = 2;
+        n++;
+      })
+    );
+    expect(n).toBe(0);
+    state.a = 2;
+    expect(n).toBe(0);
+    await nextMicroTick();
+    expect(n).toBe(0); // key has not be read yet
+    state.a = state.a + 5; // key is read and then modified
+    expect(n).toBe(0);
+    await nextMicroTick();
+    expect(n).toBe(1);
+    // the write a = 2 inside the batched callback triggered another notification, wait for it
+    await nextMicroTick();
+    expect(n).toBe(2);
+    // Should now be stable as we're writing the same value again
+    await nextMicroTick();
+    expect(n).toBe(2);
+
+    // Do it again to check it's not broken
+    state.a = state.a + 5; // key is read and then modified
+    expect(n).toBe(2);
+    await nextMicroTick();
+    expect(n).toBe(3);
+    // the write a = 2 inside the batched callback triggered another notification, wait for it
+    await nextMicroTick();
+    expect(n).toBe(4);
+    // Should now be stable as we're writing the same value again
+    await nextMicroTick();
+    expect(n).toBe(4);
+  });
+
   test("setting property to same value does not trigger callback", async () => {
     let n = 0;
     const state = createReactive({ a: 1 }, () => n++);

tests/utils.test.ts

@@ -50,4 +50,24 @@ describe("batched", () => {
     await nextMicroTick();
     expect(n).toBe(1);
   });
+
+  test("calling batched function from within the callback is not treated as part of the original batch", async () => {
+    let n = 0;
+    let fn = batched(() => {
+      n++;
+      if (n === 1) {
+        fn();
+      }
+    });
+
+    expect(n).toBe(0);
+    fn();
+    expect(n).toBe(0);
+    await nextMicroTick(); // First batch
+    expect(n).toBe(1);
+    await nextMicroTick(); // Second batch initiated from within the callback
+    expect(n).toBe(2);
+    await nextMicroTick();
+    expect(n).toBe(2);
+  });
 });