add single recomputeDependents to batch

Author: dmaskasky

src/vanilla/store.ts

@@ -173,11 +173,11 @@ type BatchPriority = 0 | 1 | 2
 
 type Batch = [
   /** high priority listeners */
-  priority0: Set<() => void>,
+  priority0: Set<(batch: Batch) => void>,
   /** atom listeners */
-  priority1: Set<() => void>,
+  priority1: Set<(batch: Batch) => void>,
   /** atom mount hooks */
-  priority2: Set<() => void>,
+  priority2: Set<(batch: Batch) => void>,
 ] & {
   /** changed Atoms */
   C: Set<AnyAtom>
@@ -189,7 +189,7 @@ const createBatch = (): Batch =>
 const addBatchFunc = (
   batch: Batch,
   priority: BatchPriority,
-  fn: () => void,
+  fn: (batch: Batch) => void,
 ) => {
   batch[priority].add(fn)
 }
@@ -203,21 +203,20 @@ const registerBatchAtom = (
     batch.C.add(atom)
     atomState.u?.(batch)
     const scheduleListeners = () => {
-      atomState.m?.l.forEach((listener) => addBatchFunc(batch, 1, listener))
+      atomState.m?.l.forEach((listener) =>
+        addBatchFunc(batch, 1, () => listener()),
+      )
     }
     addBatchFunc(batch, 1, scheduleListeners)
   }
 }
 
-const flushBatch = (
-  batch: Batch,
-  recomputeDependents: (batch: Batch) => void,
-) => {
+const flushBatch = (batch: Batch) => {
   let error: AnyError
   let hasError = false
-  const call = (fn: () => void) => {
+  const call = (fn: (batch: Batch) => void) => {
     try {
-      fn()
+      fn(batch)
     } catch (e) {
       if (!hasError) {
         error = e
@@ -226,8 +225,6 @@ const flushBatch = (
     }
   }
   while (batch.C.size || batch.some((channel) => channel.size)) {
-    recomputeDependents(batch)
-    batch.C.clear()
     for (const channel of batch) {
       channel.forEach(call)
       channel.clear()
@@ -380,7 +377,7 @@ const buildStore = (...storeArgs: StoreArgs): Store => {
           const batch = createBatch()
           addDependency(atom, atomState, a, aState)
           mountDependencies(batch, atom, atomState)
-          recomputeAndFlushBatch(batch)
+          flushBatch(batch)
         }
       }
     }
@@ -422,7 +419,7 @@ const buildStore = (...storeArgs: StoreArgs): Store => {
           if (atomState.m) {
             const batch = createBatch()
             mountDependencies(batch, atom, atomState)
-            recomputeAndFlushBatch(batch)
+            flushBatch(batch)
           }
         }
         valueOrPromise.then(complete, complete)
@@ -441,18 +438,23 @@ const buildStore = (...storeArgs: StoreArgs): Store => {
   const readAtom = <Value>(atom: Atom<Value>): Value =>
     returnAtomValue(readAtomState(undefined, atom))
 
-  const getMountedDependents = <Value>(
+  const getMountedOrBatchDependents = <Value>(
     atomState: AtomState<Value>,
   ): Map<AnyAtom, AtomState> => {
-    const mountedDependents = new Map<AnyAtom, AtomState>()
-    const dependents = new Set([...(atomState.m?.t || []), ...atomState.p])
-    for (const a of dependents) {
+    const dependents = new Map<AnyAtom, AtomState>()
+    for (const a of atomState.m?.t || []) {
       const aState = ensureAtomState(a)
       if (aState.m) {
-        mountedDependents.set(a, aState)
+        dependents.set(a, aState)
       }
     }
-    return mountedDependents
+    for (const atomWithPendingPromise of atomState.p) {
+      dependents.set(
+        atomWithPendingPromise,
+        ensureAtomState(atomWithPendingPromise),
+      )
+    }
+    return dependents
   }
 
   const dirtyDependents = <Value>(atomState: AtomState<Value>) => {
@@ -465,7 +467,7 @@ const buildStore = (...storeArgs: StoreArgs): Store => {
         continue
       }
       aState.x = true
-      for (const [, s] of getMountedDependents(aState)) {
+      for (const [, s] of getMountedOrBatchDependents(aState)) {
         if (!dependents.has(s)) {
           dependents.add(s)
           stack.push(s)
@@ -512,7 +514,7 @@ const buildStore = (...storeArgs: StoreArgs): Store => {
       }
       visiting.add(a)
       // Push unvisited dependents onto the stack
-      for (const [d, s] of getMountedDependents(aState)) {
+      for (const [d, s] of getMountedOrBatchDependents(aState)) {
         if (a !== d && !visiting.has(d)) {
           stack.push([d, s])
         }
@@ -539,11 +541,9 @@ const buildStore = (...storeArgs: StoreArgs): Store => {
       }
       delete aState.x
     }
+    batch.C.clear()
   }
 
-  const recomputeAndFlushBatch = (batch: Batch) =>
-    flushBatch(batch, recomputeDependents)
-
   const writeAtomState = <Value, Args extends unknown[], Result>(
     batch: Batch,
     atom: WritableAtom<Value, Args, Result>,
@@ -569,6 +569,7 @@ const buildStore = (...storeArgs: StoreArgs): Store => {
           mountDependencies(batch, a, aState)
           if (prevEpochNumber !== aState.n) {
             dirtyDependents(aState)
+            addBatchFunc(batch, 0, recomputeDependents)
             registerBatchAtom(batch, a, aState)
           }
           return undefined as R
@@ -577,7 +578,7 @@ const buildStore = (...storeArgs: StoreArgs): Store => {
         }
       } finally {
         if (!isSync) {
-          recomputeAndFlushBatch(batch)
+          flushBatch(batch)
         }
       }
     }
@@ -596,7 +597,7 @@ const buildStore = (...storeArgs: StoreArgs): Store => {
     try {
       return writeAtomState(batch, atom, ...args)
     } finally {
-      recomputeAndFlushBatch(batch)
+      flushBatch(batch)
     }
   }
 
@@ -653,7 +654,7 @@ const buildStore = (...storeArgs: StoreArgs): Store => {
               return writeAtomState(batch, atom, ...args)
             } finally {
               if (!isSync) {
-                recomputeAndFlushBatch(batch)
+                flushBatch(batch)
               }
             }
           }
@@ -690,7 +691,7 @@ const buildStore = (...storeArgs: StoreArgs): Store => {
       // unmount self
       const onUnmount = atomState.m.u
       if (onUnmount) {
-        addBatchFunc(batch, 2, () => onUnmount(batch))
+        addBatchFunc(batch, 2, onUnmount)
       }
       delete atomState.m
       atomState.h?.(batch)
@@ -710,12 +711,12 @@ const buildStore = (...storeArgs: StoreArgs): Store => {
     const mounted = mountAtom(batch, atom, atomState)
     const listeners = mounted.l
     listeners.add(listener)
-    recomputeAndFlushBatch(batch)
+    flushBatch(batch)
     return () => {
       listeners.delete(listener)
       const batch = createBatch()
       unmountAtom(batch, atom, atomState)
-      recomputeAndFlushBatch(batch)
+      flushBatch(batch)
     }
   }
 

tests/vanilla/effect.test.ts

@@ -1,10 +1,11 @@
 import { expect, it, vi } from 'vitest'
-import type { Atom, Getter, Setter } from 'jotai/vanilla'
+import type { Atom, Getter, PrimitiveAtom, Setter } from 'jotai/vanilla'
 import { atom, createStore } from 'jotai/vanilla'
 
 type Store = ReturnType<typeof createStore>
 type GetAtomState = Parameters<Parameters<Store['unstable_derive']>[0]>[0]
 type AtomState = NonNullable<ReturnType<GetAtomState>>
+type Batch = Parameters<NonNullable<AtomState['u']>>[0]
 type AnyAtom = Atom<unknown>
 
 type Cleanup = () => void
@@ -59,7 +60,7 @@ function syncEffect(effect: Effect): Atom<void> {
         store.set(refreshAtom, (v) => v + 1)
       } else {
         // unmount
-        batch[0].add(() => {
+        scheduleListener(batch, () => {
           ref.cleanup?.()
           delete ref.cleanup
         })
@@ -69,7 +70,17 @@ function syncEffect(effect: Effect): Atom<void> {
     internalAtomState.u = (batch) => {
       originalUpdateHook?.(batch)
       // update
-      batch[0].add(runEffect)
+      scheduleListener(batch, runEffect)
+    }
+    function scheduleListener(batch: Batch, listener: () => void) {
+      if (batch[0].size === 0) {
+        store.set(
+          atom(0, function (this: PrimitiveAtom<number>, _, set) {
+            set(this, 1)
+          }),
+        )
+      }
+      batch[0].add(listener)
     }
   }
   return atom((get) => {