fix: add batching support for updates to state

packages/@lwc/state/src/__tests__/index.test.ts

@@ -1,25 +1,56 @@
-import { describe, expect, test, vi } from 'vitest';
+import { afterEach, describe, expect, test, vi } from 'vitest';
 import { defineState } from '../index.js';
+// biome-ignore lint: test only
+let doubleCountNotifySpy: any;
+// biome-ignore lint: test only
+let fruitNameAndCountNotifySpy: any;
+// biome-ignore lint: test only
+let fruitNameAndCountComputeValueSpy: any;
+
+const state = defineState((atom, computed, update, _fromContext) => (...args) => {
+  const countArg = args[0] as number;
+  const fruitArg = args[1] as string;
+
+  const count = atom(countArg);
+  const fruit = atom(fruitArg);
 
-const state = defineState((atom, computed, update, _fromContext) => (arg: number) => {
-  const count = atom(arg);
   const doubleCount = computed({ count }, ({ count: countValue }) => (countValue as number) * 2);
+  doubleCountNotifySpy = vi.spyOn(doubleCount, 'notify');
+
+  const fruitNameAndCount = computed(
+    { fruit, count },
+    ({ fruit, count }) => `I have ${count} ${fruit}${(count as number) > 1 ? 's' : ''}`,
+  );
+  fruitNameAndCountNotifySpy = vi.spyOn(fruitNameAndCount, 'notify');
+  fruitNameAndCountComputeValueSpy = vi.spyOn(fruitNameAndCount, 'computeValue');
+
   const increment = update({ count }, ({ count: countValue }) => ({
     count: countValue + 1,
   }));
   const incrementBy = update({ count }, ({ count: countValue }, amount: number) => ({
     count: countValue + amount,
   }));
+  const changeFruit = update({ fruit }, ({ fruit }, newFruit: string) => ({
+    fruit: newFruit,
+  }));
 
   return {
     count,
     doubleCount,
     increment,
     incrementBy,
+    fruitNameAndCount,
+    changeFruit,
   };
 });
 
+const flushMicrotasks = () => Promise.resolve();
+
 describe('state manager', () => {
+  afterEach(() => {
+    vi.clearAllMocks();
+  });
+
   test('initial state', () => {
     const s = state(5);
 
@@ -35,14 +66,15 @@ describe('state manager', () => {
     expect(s.value.doubleCount).toBe(4);
   });
 
-  test('subscribing to state manager works', () => {
+  test('subscribing to state manager works', async () => {
     const s = state(1);
     const sub = vi.fn();
 
     s.subscribe(sub);
     s.value.increment();
 
-    expect(sub).toHaveBeenCalledTimes(2);
+    await flushMicrotasks();
+    expect(sub).toHaveBeenCalledTimes(1);
   });
 
   test('incrementBy updates count and doubleCount', () => {
@@ -65,30 +97,32 @@ describe('state manager', () => {
     expect(s.value.doubleCount).toBe(10);
   });
 
-  test('subscription triggers on update', () => {
+  test('subscription triggers on update', async () => {
     const s = state(1);
     const sub = vi.fn();
     s.subscribe(sub);
 
     s.value.increment();
-
-    expect(sub).toHaveBeenCalledTimes(2);
+    await flushMicrotasks();
+    expect(sub).toHaveBeenCalledTimes(1);
     s.value.incrementBy(2);
-    expect(sub).toHaveBeenCalledTimes(4);
+    await flushMicrotasks();
+    expect(sub).toHaveBeenCalledTimes(2);
   });
 
-  test('unsubscribe stops notifications', () => {
+  test('unsubscribe stops notifications', async () => {
     const s = state(1);
     const sub = vi.fn();
 
     const unsubscribe = s.subscribe(sub);
     s.value.increment();
-    expect(sub).toHaveBeenCalledTimes(2);
+    await flushMicrotasks();
+    expect(sub).toHaveBeenCalledTimes(1);
     unsubscribe();
 
     s.value.increment();
-
-    expect(sub).toHaveBeenCalledTimes(2);
+    await flushMicrotasks();
+    expect(sub).toHaveBeenCalledTimes(1);
   });
 
   test('state is immutable', () => {
@@ -98,4 +132,67 @@ describe('state manager', () => {
       s.value.count = 10;
     }).toThrow();
   });
+
+  test('should call computeds notify only after an update and not at init', () => {
+    const s = state(1, 'Apple');
+    expect(doubleCountNotifySpy).not.toHaveBeenCalled();
+    expect(fruitNameAndCountNotifySpy).not.toHaveBeenCalled();
+
+    s.value.increment();
+    expect(doubleCountNotifySpy).toHaveBeenCalledOnce();
+    expect(fruitNameAndCountNotifySpy).toHaveBeenCalledOnce();
+  });
+
+  test('should only notify computeds that depends on the atom being updated', () => {
+    const s = state(5, 'Apple');
+
+    s.value.increment();
+    expect(doubleCountNotifySpy).toHaveBeenCalledOnce();
+    expect(fruitNameAndCountNotifySpy).toHaveBeenCalledOnce();
+
+    s.value.changeFruit('Banana');
+    // should not be notified
+    expect(doubleCountNotifySpy).toHaveBeenCalledOnce();
+    expect(fruitNameAndCountNotifySpy).toHaveBeenCalledTimes(2);
+  });
+
+  test('should compute computeds value lazily', () => {
+    const s = state(5, 'Apple');
+    expect(fruitNameAndCountComputeValueSpy).not.toHaveBeenCalled();
+
+    // every time we access state managers value
+    // all stale values are re-computed
+    s.value.increment();
+    expect(fruitNameAndCountComputeValueSpy).toHaveBeenCalledOnce();
+    s.value.changeFruit('Banana');
+    expect(fruitNameAndCountComputeValueSpy).toHaveBeenCalledTimes(2);
+  });
+
+  test('should not compute value until value is not accessed through getter', () => {
+    const s = state(5, 'Apple');
+    // getting references to the updaters and computed
+    const { increment, changeFruit, fruitNameAndCount } = s.value;
+
+    expect(fruitNameAndCountComputeValueSpy).toHaveBeenCalledOnce();
+
+    increment();
+    changeFruit('Banana');
+    // not re-computed since value is not accessed
+    expect(fruitNameAndCountComputeValueSpy).toHaveBeenCalledOnce();
+    expect(fruitNameAndCount).toBe('I have 5 Apples');
+    // computeds value is re-computed since accessed through state managers getter
+    expect(s.value.fruitNameAndCount).toBe('I have 6 Bananas');
+    expect(fruitNameAndCountComputeValueSpy).toHaveBeenCalledTimes(2);
+  });
+
+  test('should compute values correctly after multiple updates', async () => {
+    const s = state(2, 'Orange');
+    s.value.increment();
+    s.value.incrementBy(3);
+    s.value.changeFruit('Grape');
+    await flushMicrotasks();
+
+    expect(s.value.count).toBe(6);
+    expect(s.value.fruitNameAndCount).toBe('I have 6 Grapes');
+  });
 });

packages/@lwc/state/src/index.ts

@@ -114,6 +114,7 @@ export const defineState: DefineState = (defineStateCallback) => {
       private internalStateShape: Record<string, Signal<unknown> | ExposedUpdater>;
       private _value: OuterStateShape;
       private isStale = true;
+      private isNotifyScheduled = false;
 
       constructor() {
         super();
@@ -148,7 +149,16 @@ export const defineState: DefineState = (defineStateCallback) => {
 
       private scheduledNotify() {
         this.isStale = true;
-        super.notify();
+        // super.notify();
+
+        if (!this.isNotifyScheduled) {
+          queueMicrotask(() => {
+            this.isNotifyScheduled = false;
+            super.notify();
+          });
+
+          this.isNotifyScheduled = true;
+        }
       }
 
       get value() {