refactor: pass recompute as arg to flushBatch

pmndrs · Jan 8, 2025 · 99515ad · 99515ad
1 parent 0fc4d6e
commit 99515ad
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 54 deletions.
diff --git a/src/vanilla/store.ts b/src/vanilla/store.ts
@@ -209,7 +209,10 @@ const registerBatchAtom = (
   }
 }
 
-const flushBatch = (batch: Batch) => {
+const flushBatch = (
+  batch: Batch,
+  recomputeDependents: (batch: Batch) => void,
+) => {
   let error: AnyError
   let hasError = false
   const call = (fn: () => void) => {
@@ -223,6 +226,7 @@ const flushBatch = (batch: Batch) => {
     }
   }
   while (batch.C.size || batch.some((channel) => channel.size)) {
+    recomputeDependents(batch)
     batch.C.clear()
     for (const channel of batch) {
       channel.forEach(call)
@@ -377,7 +381,7 @@ const buildStore = (...storeArgs: StoreArgs): Store => {
           const batch = createBatch()
           addDependency(atom, atomState, a, aState)
           mountDependencies(batch, atom, atomState)
-          flushBatch(batch)
+          recomputeAndFlushBatch(batch)
         }
       }
     }
@@ -419,7 +423,7 @@ const buildStore = (...storeArgs: StoreArgs): Store => {
           if (atomState.m) {
             const batch = createBatch()
             mountDependencies(batch, atom, atomState)
-            flushBatch(batch)
+            recomputeAndFlushBatch(batch)
           }
         }
         valueOrPromise.then(complete, complete)
@@ -458,24 +462,20 @@ const buildStore = (...storeArgs: StoreArgs): Store => {
     return dependents
   }
 
-  const dirtyDependents = <Value>(
-    batch: Batch,
-    atom: Atom<Value>,
-    atomState: AtomState<Value>,
-  ) => {
-    const dependents = new Map<AnyAtom, AtomState>([[atom, atomState]])
-    const stack: [AnyAtom, AtomState][] = [[atom, atomState]]
+  const dirtyDependents = <Value>(atomState: AtomState<Value>) => {
+    const dependents = new Set<AtomState>([atomState])
+    const stack: AtomState[] = [atomState]
     while (stack.length > 0) {
-      const [a, aState] = stack.pop()!
+      const aState = stack.pop()!
       if (aState.x) {
         // already dirty
         continue
       }
       aState.x = true
-      for (const [d, s] of getMountedOrBatchDependents(batch, a, aState)) {
-        if (!dependents.has(d)) {
-          dependents.set(d, s)
-          stack.push([d, s])
+      for (const [, s] of getMountedOrBatchDependents(aState)) {
+        if (!dependents.has(s)) {
+          dependents.add(s)
+          stack.push(s)
         }
       }
     }
@@ -497,7 +497,7 @@ const buildStore = (...storeArgs: StoreArgs): Store => {
     // Visit the root atom. This is the only atom in the dependency graph
     // without incoming edges, which is one reason we can simplify the algorithm
     const stack: [a: AnyAtom, aState: AtomState][] = Array.from(
-      batch.D.keys(),
+      batch.C,
       (atom) => [atom, ensureAtomState(atom)],
     )
     while (stack.length > 0) {
@@ -532,7 +532,7 @@ const buildStore = (...storeArgs: StoreArgs): Store => {
       const [a, aState, prevEpochNumber] = topSortedReversed[i]!
       let hasChangedDeps = false
       for (const dep of aState.d.keys()) {
-        if (dep !== a && batch.D.has(dep)) {
+        if (dep !== a && batch.C.has(dep)) {
           hasChangedDeps = true
           break
         }
@@ -548,6 +548,9 @@ const buildStore = (...storeArgs: StoreArgs): Store => {
     }
   }
 
+  const recomputeAndFlushBatch = (batch: Batch) =>
+    flushBatch(batch, recomputeDependents)
+
   const writeAtomState = <Value, Args extends unknown[], Result>(
     batch: Batch,
     atom: WritableAtom<Value, Args, Result>,
@@ -572,8 +575,7 @@ const buildStore = (...storeArgs: StoreArgs): Store => {
           setAtomStateValueOrPromise(a, aState, v)
           mountDependencies(batch, a, aState)
           if (prevEpochNumber !== aState.n) {
-            dirtyDependents(batch, a, aState)
-            batch.R = recomputeDependents
+            dirtyDependents(aState)
             registerBatchAtom(batch, a, aState)
           }
           return undefined as R
@@ -582,7 +584,7 @@ const buildStore = (...storeArgs: StoreArgs): Store => {
         }
       } finally {
         if (!isSync) {
-          flushBatch(batch)
+          recomputeAndFlushBatch(batch)
         }
       }
     }
@@ -601,7 +603,7 @@ const buildStore = (...storeArgs: StoreArgs): Store => {
     try {
       return writeAtomState(batch, atom, ...args)
     } finally {
-      flushBatch(batch)
+      recomputeAndFlushBatch(batch)
     }
   }
 
@@ -658,7 +660,7 @@ const buildStore = (...storeArgs: StoreArgs): Store => {
               return writeAtomState(batch, atom, ...args)
             } finally {
               if (!isSync) {
-                flushBatch(batch)
+                recomputeAndFlushBatch(batch)
               }
             }
           }
@@ -715,12 +717,12 @@ const buildStore = (...storeArgs: StoreArgs): Store => {
     const mounted = mountAtom(batch, atom, atomState)
     const listeners = mounted.l
     listeners.add(listener)
-    flushBatch(batch)
+    recomputeAndFlushBatch(batch)
     return () => {
       listeners.delete(listener)
       const batch = createBatch()
       unmountAtom(batch, atom, atomState)
-      flushBatch(batch)
+      recomputeAndFlushBatch(batch)
     }
   }
 

diff --git a/tests/vanilla/store.test.tsx b/tests/vanilla/store.test.tsx
@@ -1109,36 +1109,6 @@ it('recomputes dependents of unmounted atoms', () => {
   expect(bRead).not.toHaveBeenCalled()
 })
 
-it('runs recomputeDependents on atoms in the correct order', async () => {
-  let i = 0
-  function createHistoryAtoms<T>(initialValue: T) {
-    const historyStackAtom = atom<T[]>([initialValue])
-    historyStackAtom.debugLabel = `${i}:historyStackAtom`
-
-    const valueAtom = atom((get) => {
-      const entry = get(historyStackAtom)[0]
-      return entry
-    })
-    valueAtom.debugLabel = `${i}:valueAtom`
-
-    const resetAtom = atom(null, (_, set, value: T) => {
-      set(historyStackAtom, [value])
-    })
-    resetAtom.debugLabel = `${i}:resetAtom`
-    i++
-    return { valueAtom, resetAtom }
-  }
-
-  const val1Atoms = createHistoryAtoms('foo')
-  const val2Atoms = createHistoryAtoms<string | null>(null)
-
-  const initAtom = atom(null, (_get, set) => {
-    // if comment out this line, the test will pass
-    set(val2Atoms.resetAtom, null)
-    set(val1Atoms.resetAtom, 'bar')
-  }
-})
-
 it('recomputes all changed atom dependents together', async () => {
   const a = atom([0])
   const b = atom([0])