From 7949d4a98c5e85b745eb91a3b2ee92a1cb0796f9 Mon Sep 17 00:00:00 2001
From: David Maskasky <dmaskasky@gmail.com>
Date: Sat, 4 Jan 2025 10:07:17 -0700
Subject: [PATCH 01/13] add failing test

---
 tests/vanilla/store.test.tsx | 61 ++++++++++++++++++++++++++++++++++++
 1 file changed, 61 insertions(+)

diff --git a/tests/vanilla/store.test.tsx b/tests/vanilla/store.test.tsx
index 425063617e..0ea0d19fab 100644
--- a/tests/vanilla/store.test.tsx
+++ b/tests/vanilla/store.test.tsx
@@ -1108,3 +1108,64 @@ it('recomputes dependents of unmounted atoms', () => {
   store.set(w)
   expect(bRead).not.toHaveBeenCalled()
 })
+
+it('runs recomputeDependents on atoms in the correct order', async () => {
+  const store = createStore().unstable_derive((getAtomState, ...storeArgs) => [
+    (a) => Object.assign(getAtomState(a), { label: a.debugLabel }),
+    ...storeArgs,
+  ])
+  let i = 0
+  function createHistoryAtoms<T>(initialValue: T) {
+    const historyStackAtom = atom<T[]>([initialValue])
+    historyStackAtom.debugLabel = `${i}:historyStackAtom`
+    const historyIndexAtom = atom(0)
+    historyIndexAtom.debugLabel = `${i}:historyIndexAtom`
+    const valueAtom = atom(
+      (get) => get(historyStackAtom)[get(historyIndexAtom)]!,
+    )
+    valueAtom.debugLabel = `${i}:valueAtom`
+    const resetAtom = atom(null, (_, set, value: T) => {
+      set(historyStackAtom, [value])
+      set(historyIndexAtom, 0)
+    })
+    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
+    console.log('initAtom write val2Atoms')
+    set(val2Atoms.resetAtom, null)
+    console.log('initAtom write val1Atoms')
+    set(val1Atoms.resetAtom, 'bar')
+  })
+  initAtom.debugLabel = 'initAtom'
+
+  const computedValAtom = atom((get) => {
+    const v2Value = get(val2Atoms.valueAtom)
+    if (v2Value !== null) {
+      console.log('computedValAtom read val2Atoms', v2Value)
+      return v2Value
+    }
+    const v1Value = get(val1Atoms.valueAtom)
+    console.log('computedValAtom read val2Atoms', v1Value)
+    return v1Value
+  })
+  computedValAtom.debugLabel = 'computedValAtom'
+
+  const asyncInitAtom = atom(null, async (_get, set) => {
+    // if comment out this line, the test will pass [DOES NOT WORK]
+    await new Promise((resolve) => setTimeout(resolve, 0))
+
+    set(initAtom)
+  })
+  store.sub(computedValAtom, () => {})
+  console.log('set asyncInitAtom')
+  await store.set(asyncInitAtom)
+  const result = store.get(computedValAtom)
+  expect(result).toBe('bar')
+})

From 7930a8dfc311171fe2bfe4bab748c049a966f454 Mon Sep 17 00:00:00 2001
From: David Maskasky <dmaskasky@gmail.com>
Date: Sun, 5 Jan 2025 20:18:17 -0800
Subject: [PATCH 02/13] a single recomputeDependents in flushBatch

---
 src/vanilla/store.ts         | 53 +++++++++++++++++++-----------------
 tests/vanilla/store.test.tsx | 48 ++++++++++++++++----------------
 2 files changed, 52 insertions(+), 49 deletions(-)

diff --git a/src/vanilla/store.ts b/src/vanilla/store.ts
index 6c20762127..2ac868091b 100644
--- a/src/vanilla/store.ts
+++ b/src/vanilla/store.ts
@@ -458,11 +458,7 @@ const buildStore = (...storeArgs: StoreArgs): Store => {
     return dependents
   }
 
-  const recomputeDependents = <Value>(
-    batch: Batch,
-    atom: Atom<Value>,
-    atomState: AtomState<Value>,
-  ) => {
+  const recomputeDependents = (batch: Batch) => {
     // Step 1: traverse the dependency graph to build the topsorted atom list
     // We don't bother to check for cycles, which simplifies the algorithm.
     // This is a topological sort via depth-first search, slightly modified from
@@ -477,7 +473,10 @@ const buildStore = (...storeArgs: StoreArgs): Store => {
     const visited = new Set<AnyAtom>()
     // 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][] = [[atom, atomState]]
+    const stack: [a: AnyAtom, aState: AtomState][] = Array.from(
+      batch.D.keys(),
+      (atom) => [atom, getAtomState(atom)],
+    )
     while (stack.length > 0) {
       const [a, aState] = stack[stack.length - 1]!
       if (visited.has(a)) {
@@ -508,27 +507,29 @@ const buildStore = (...storeArgs: StoreArgs): Store => {
 
     // Step 2: use the topSortedReversed atom list to recompute all affected atoms
     // Track what's changed, so that we can short circuit when possible
-    const finishRecompute = () => {
-      const changedAtoms = new Set<AnyAtom>([atom])
-      for (let i = topSortedReversed.length - 1; i >= 0; --i) {
-        const [a, aState, prevEpochNumber] = topSortedReversed[i]!
-        let hasChangedDeps = false
-        for (const dep of aState.d.keys()) {
-          if (dep !== a && changedAtoms.has(dep)) {
-            hasChangedDeps = true
-            break
-          }
+    console.log(
+      '    finishRecompute',
+      topSortedReversed.map(([a]) => a.debugLabel),
+    )
+    const changedAtoms = new Set<AnyAtom>(batch.D.keys())
+    for (let i = topSortedReversed.length - 1; i >= 0; --i) {
+      const [a, aState, prevEpochNumber] = topSortedReversed[i]!
+      let hasChangedDeps = false
+      for (const dep of aState.d.keys()) {
+        if (dep !== a && changedAtoms.has(dep)) {
+          hasChangedDeps = true
+          break
         }
-        if (hasChangedDeps) {
-          readAtomState(batch, a)
-          mountDependencies(batch, a, aState)
-          if (prevEpochNumber !== aState.n) {
-            registerBatchAtom(batch, a, aState)
-            changedAtoms.add(a)
-          }
+      }
+      if (hasChangedDeps) {
+        readAtomState(batch, a)
+        mountDependencies(batch, a, aState)
+        if (prevEpochNumber !== aState.n) {
+          registerBatchAtom(batch, a, aState)
+          changedAtoms.add(a)
         }
-        delete aState.x
       }
+      delete aState.x
     }
     addBatchFunc(batch, 0, finishRecompute)
   }
@@ -558,7 +559,9 @@ const buildStore = (...storeArgs: StoreArgs): Store => {
           mountDependencies(batch, a, aState)
           if (prevEpochNumber !== aState.n) {
             registerBatchAtom(batch, a, aState)
-            recomputeDependents(batch, a, aState)
+            addBatchFunc(batch, BATCH_PRIORITY_HIGH, () =>
+              recomputeDependents(batch),
+            )
           }
           return undefined as R
         } else {
diff --git a/tests/vanilla/store.test.tsx b/tests/vanilla/store.test.tsx
index 0ea0d19fab..3597a04ab4 100644
--- a/tests/vanilla/store.test.tsx
+++ b/tests/vanilla/store.test.tsx
@@ -1118,15 +1118,15 @@ it('runs recomputeDependents on atoms in the correct order', async () => {
   function createHistoryAtoms<T>(initialValue: T) {
     const historyStackAtom = atom<T[]>([initialValue])
     historyStackAtom.debugLabel = `${i}:historyStackAtom`
-    const historyIndexAtom = atom(0)
-    historyIndexAtom.debugLabel = `${i}:historyIndexAtom`
-    const valueAtom = atom(
-      (get) => get(historyStackAtom)[get(historyIndexAtom)]!,
-    )
+
+    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])
-      set(historyIndexAtom, 0)
     })
     resetAtom.debugLabel = `${i}:resetAtom`
     i++
@@ -1137,35 +1137,35 @@ it('runs recomputeDependents on atoms in the correct order', async () => {
   const val2Atoms = createHistoryAtoms<string | null>(null)
 
   const initAtom = atom(null, (_get, set) => {
+    console.log('  initAtom write val2Atoms')
     // if comment out this line, the test will pass
-    console.log('initAtom write val2Atoms')
     set(val2Atoms.resetAtom, null)
-    console.log('initAtom write val1Atoms')
+    console.log('  initAtom write val1Atoms')
     set(val1Atoms.resetAtom, 'bar')
   })
   initAtom.debugLabel = 'initAtom'
 
   const computedValAtom = atom((get) => {
     const v2Value = get(val2Atoms.valueAtom)
-    if (v2Value !== null) {
-      console.log('computedValAtom read val2Atoms', v2Value)
-      return v2Value
-    }
     const v1Value = get(val1Atoms.valueAtom)
-    console.log('computedValAtom read val2Atoms', v1Value)
+    console.log('  computedValAtom read val1Atoms', v1Value, v2Value)
     return v1Value
   })
   computedValAtom.debugLabel = 'computedValAtom'
 
-  const asyncInitAtom = atom(null, async (_get, set) => {
-    // if comment out this line, the test will pass [DOES NOT WORK]
-    await new Promise((resolve) => setTimeout(resolve, 0))
-
-    set(initAtom)
-  })
-  store.sub(computedValAtom, () => {})
-  console.log('set asyncInitAtom')
-  await store.set(asyncInitAtom)
-  const result = store.get(computedValAtom)
-  expect(result).toBe('bar')
+  type Store = ReturnType<typeof createStore>
+  function testStore(store: Store) {
+    console.log('sub computedValAtom ----')
+    store.sub(computedValAtom, () => {})
+    console.log('set initAtom ----')
+    store.set(initAtom)
+    const result = store.get(computedValAtom)
+    expect(result).toBe('bar')
+  }
+  // console.log('\n2.10.0')
+  // testStore(createStores['2.10.0']!())
+  // console.log('\n2.10.4')
+  // testStore(createStores['2.10.4']!())
+  console.log('\n2.11.0')
+  testStore(createStore())
 })

From cc4ed04216a1c32363590c7cd9eca50190d5e3df Mon Sep 17 00:00:00 2001
From: David Maskasky <dmaskasky@gmail.com>
Date: Sun, 5 Jan 2025 20:53:26 -0800
Subject: [PATCH 03/13] centralize finishRecompute

---
 src/vanilla/store.ts         | 37 +++++++++++++++++++++----------
 tests/vanilla/effect.test.ts | 42 ++----------------------------------
 2 files changed, 28 insertions(+), 51 deletions(-)

diff --git a/src/vanilla/store.ts b/src/vanilla/store.ts
index 2ac868091b..17a0c18a16 100644
--- a/src/vanilla/store.ts
+++ b/src/vanilla/store.ts
@@ -458,6 +458,29 @@ const buildStore = (...storeArgs: StoreArgs): Store => {
     return dependents
   }
 
+  const dirtyDependents = <Value>(
+    batch: Batch,
+    atom: Atom<Value>,
+    atomState: AtomState<Value>,
+  ) => {
+    const dependents = new Map<AnyAtom, AtomState>()
+    const stack: [AnyAtom, AtomState][] = [[atom, atomState]]
+    while (stack.length > 0) {
+      const [a, 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])
+        }
+      }
+    }
+  }
+
   const recomputeDependents = (batch: Batch) => {
     // Step 1: traverse the dependency graph to build the topsorted atom list
     // We don't bother to check for cycles, which simplifies the algorithm.
@@ -475,7 +498,7 @@ const buildStore = (...storeArgs: StoreArgs): Store => {
     // without incoming edges, which is one reason we can simplify the algorithm
     const stack: [a: AnyAtom, aState: AtomState][] = Array.from(
       batch.D.keys(),
-      (atom) => [atom, getAtomState(atom)],
+      (atom) => [atom, ensureAtomState(atom)],
     )
     while (stack.length > 0) {
       const [a, aState] = stack[stack.length - 1]!
@@ -491,8 +514,6 @@ const buildStore = (...storeArgs: StoreArgs): Store => {
         topSortedReversed.push([a, aState, aState.n])
         // Atom has been visited but not yet processed
         visited.add(a)
-        // Mark atom dirty
-        aState.x = true
         stack.pop()
         continue
       }
@@ -507,10 +528,6 @@ const buildStore = (...storeArgs: StoreArgs): Store => {
 
     // Step 2: use the topSortedReversed atom list to recompute all affected atoms
     // Track what's changed, so that we can short circuit when possible
-    console.log(
-      '    finishRecompute',
-      topSortedReversed.map(([a]) => a.debugLabel),
-    )
     const changedAtoms = new Set<AnyAtom>(batch.D.keys())
     for (let i = topSortedReversed.length - 1; i >= 0; --i) {
       const [a, aState, prevEpochNumber] = topSortedReversed[i]!
@@ -531,7 +548,6 @@ const buildStore = (...storeArgs: StoreArgs): Store => {
       }
       delete aState.x
     }
-    addBatchFunc(batch, 0, finishRecompute)
   }
 
   const writeAtomState = <Value, Args extends unknown[], Result>(
@@ -558,10 +574,9 @@ const buildStore = (...storeArgs: StoreArgs): Store => {
           setAtomStateValueOrPromise(a, aState, v)
           mountDependencies(batch, a, aState)
           if (prevEpochNumber !== aState.n) {
+            dirtyDependents(batch, a, aState)
+            batch.R = recomputeDependents
             registerBatchAtom(batch, a, aState)
-            addBatchFunc(batch, BATCH_PRIORITY_HIGH, () =>
-              recomputeDependents(batch),
-            )
           }
           return undefined as R
         } else {
diff --git a/tests/vanilla/effect.test.ts b/tests/vanilla/effect.test.ts
index 697d6a6733..3c1d158484 100644
--- a/tests/vanilla/effect.test.ts
+++ b/tests/vanilla/effect.test.ts
@@ -6,7 +6,6 @@ type Store = ReturnType<typeof createStore>
 type GetAtomState = Parameters<Parameters<Store['unstable_derive']>[0]>[0]
 type AtomState = NonNullable<ReturnType<GetAtomState>>
 type AnyAtom = Atom<unknown>
-type Batch = Parameters<NonNullable<AtomState['u']>>[0]
 
 type Cleanup = () => void
 type Effect = (get: Getter, set: Setter) => Cleanup | void
@@ -17,8 +16,6 @@ type Ref = {
   cleanup?: Cleanup | undefined
 }
 
-const syncEffectChannelSymbol = Symbol()
-
 function syncEffect(effect: Effect): Atom<void> {
   const refAtom = atom<Ref>(() => ({ inProgress: 0, epoch: 0 }))
   const refreshAtom = atom(0)
@@ -62,8 +59,7 @@ function syncEffect(effect: Effect): Atom<void> {
         store.set(refreshAtom, (v) => v + 1)
       } else {
         // unmount
-        const syncEffectChannel = ensureBatchChannel(batch)
-        syncEffectChannel.add(() => {
+        batch[0].add(() => {
           ref.cleanup?.()
           delete ref.cleanup
         })
@@ -73,8 +69,7 @@ function syncEffect(effect: Effect): Atom<void> {
     internalAtomState.u = (batch) => {
       originalUpdateHook?.(batch)
       // update
-      const syncEffectChannel = ensureBatchChannel(batch)
-      syncEffectChannel.add(runEffect)
+      batch[0].add(runEffect)
     }
   }
   return atom((get) => {
@@ -82,39 +77,6 @@ function syncEffect(effect: Effect): Atom<void> {
   })
 }
 
-type BatchWithSyncEffect = Batch & {
-  [syncEffectChannelSymbol]?: Set<() => void>
-}
-function ensureBatchChannel(batch: BatchWithSyncEffect) {
-  // ensure continuation of the flushBatch while loop
-  const originalQueue = batch[1]
-  if (!originalQueue) {
-    throw new Error('batch[1] must be present')
-  }
-  if (!batch[syncEffectChannelSymbol]) {
-    batch[syncEffectChannelSymbol] = new Set<() => void>()
-    batch[1] = {
-      ...originalQueue,
-      add(item) {
-        originalQueue.add(item)
-        return this
-      },
-      clear() {
-        batch[syncEffectChannelSymbol]!.clear()
-        originalQueue.clear()
-      },
-      forEach(callback) {
-        batch[syncEffectChannelSymbol]!.forEach(callback)
-        originalQueue.forEach(callback)
-      },
-      get size() {
-        return batch[syncEffectChannelSymbol]!.size + originalQueue.size
-      },
-    }
-  }
-  return batch[syncEffectChannelSymbol]!
-}
-
 const getAtomStateMap = new WeakMap<Store, GetAtomState>()
 
 /**

From e7434a0ede0acc633a0ade992c0eef5a1f65be16 Mon Sep 17 00:00:00 2001
From: David Maskasky <dmaskasky@gmail.com>
Date: Sun, 5 Jan 2025 21:02:31 -0800
Subject: [PATCH 04/13] cleanup

---
 tests/vanilla/store.test.tsx | 25 ++++++-------------------
 1 file changed, 6 insertions(+), 19 deletions(-)

diff --git a/tests/vanilla/store.test.tsx b/tests/vanilla/store.test.tsx
index 3597a04ab4..0eb4f51379 100644
--- a/tests/vanilla/store.test.tsx
+++ b/tests/vanilla/store.test.tsx
@@ -1137,35 +1137,22 @@ it('runs recomputeDependents on atoms in the correct order', async () => {
   const val2Atoms = createHistoryAtoms<string | null>(null)
 
   const initAtom = atom(null, (_get, set) => {
-    console.log('  initAtom write val2Atoms')
     // if comment out this line, the test will pass
     set(val2Atoms.resetAtom, null)
-    console.log('  initAtom write val1Atoms')
     set(val1Atoms.resetAtom, 'bar')
   })
   initAtom.debugLabel = 'initAtom'
 
   const computedValAtom = atom((get) => {
-    const v2Value = get(val2Atoms.valueAtom)
+    get(val2Atoms.valueAtom)
     const v1Value = get(val1Atoms.valueAtom)
-    console.log('  computedValAtom read val1Atoms', v1Value, v2Value)
     return v1Value
   })
   computedValAtom.debugLabel = 'computedValAtom'
 
-  type Store = ReturnType<typeof createStore>
-  function testStore(store: Store) {
-    console.log('sub computedValAtom ----')
-    store.sub(computedValAtom, () => {})
-    console.log('set initAtom ----')
-    store.set(initAtom)
-    const result = store.get(computedValAtom)
-    expect(result).toBe('bar')
-  }
-  // console.log('\n2.10.0')
-  // testStore(createStores['2.10.0']!())
-  // console.log('\n2.10.4')
-  // testStore(createStores['2.10.4']!())
-  console.log('\n2.11.0')
-  testStore(createStore())
+  const store = createStore()
+  store.sub(computedValAtom, () => {})
+  store.set(initAtom)
+  const result = store.get(computedValAtom)
+  expect(result).toBe('bar')
 })

From a5fd3d9d22f15d182bd321ab51cb2ef908771adc Mon Sep 17 00:00:00 2001
From: David Maskasky <dmaskasky@gmail.com>
Date: Sun, 5 Jan 2025 21:23:01 -0800
Subject: [PATCH 05/13] fix test

---
 tests/vanilla/store.test.tsx | 4 ----
 1 file changed, 4 deletions(-)

diff --git a/tests/vanilla/store.test.tsx b/tests/vanilla/store.test.tsx
index 0eb4f51379..3ea1cb89a1 100644
--- a/tests/vanilla/store.test.tsx
+++ b/tests/vanilla/store.test.tsx
@@ -1110,10 +1110,6 @@ it('recomputes dependents of unmounted atoms', () => {
 })
 
 it('runs recomputeDependents on atoms in the correct order', async () => {
-  const store = createStore().unstable_derive((getAtomState, ...storeArgs) => [
-    (a) => Object.assign(getAtomState(a), { label: a.debugLabel }),
-    ...storeArgs,
-  ])
   let i = 0
   function createHistoryAtoms<T>(initialValue: T) {
     const historyStackAtom = atom<T[]>([initialValue])

From 4522c6f128c88db5bc597f1268f7dd40fec60d39 Mon Sep 17 00:00:00 2001
From: David Maskasky <dmaskasky@gmail.com>
Date: Sun, 5 Jan 2025 22:21:00 -0800
Subject: [PATCH 06/13] simplify test

---
 tests/vanilla/store.test.tsx | 28 ++++++++++++++++------------
 1 file changed, 16 insertions(+), 12 deletions(-)

diff --git a/tests/vanilla/store.test.tsx b/tests/vanilla/store.test.tsx
index 3ea1cb89a1..a8d2279274 100644
--- a/tests/vanilla/store.test.tsx
+++ b/tests/vanilla/store.test.tsx
@@ -1136,19 +1136,23 @@ it('runs recomputeDependents on atoms in the correct order', async () => {
     // if comment out this line, the test will pass
     set(val2Atoms.resetAtom, null)
     set(val1Atoms.resetAtom, 'bar')
-  })
-  initAtom.debugLabel = 'initAtom'
+  }
+})
 
-  const computedValAtom = atom((get) => {
-    get(val2Atoms.valueAtom)
-    const v1Value = get(val1Atoms.valueAtom)
-    return v1Value
+it('recomputes all changed atom dependents together', async () => {
+  const a = atom([0])
+  const b = atom([0])
+  const a0 = atom((get) => get(a)[0]!)
+  const b0 = atom((get) => get(b)[0]!)
+  const a0b0 = atom((get) => [get(a0), get(b0)])
+  const w = atom(null, (_, set) => {
+    set(a, [0])
+    set(b, [1])
   })
-  computedValAtom.debugLabel = 'computedValAtom'
-
   const store = createStore()
-  store.sub(computedValAtom, () => {})
-  store.set(initAtom)
-  const result = store.get(computedValAtom)
-  expect(result).toBe('bar')
+  store.sub(a0b0, () => {})
+  store.set(w)
+  expect(store.get(a0)).toBe(0)
+  expect(store.get(b0)).toBe(1)
+  expect(store.get(a0b0)).toEqual([0, 1])
 })

From 6ef2164e3a95fa78926dcfa7a6e2154e165c0082 Mon Sep 17 00:00:00 2001
From: David Maskasky <dmaskasky@gmail.com>
Date: Sun, 5 Jan 2025 22:26:17 -0800
Subject: [PATCH 07/13] remove redundant changedAtoms set

---
 src/vanilla/store.ts | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/src/vanilla/store.ts b/src/vanilla/store.ts
index 17a0c18a16..d5943846cb 100644
--- a/src/vanilla/store.ts
+++ b/src/vanilla/store.ts
@@ -528,12 +528,11 @@ const buildStore = (...storeArgs: StoreArgs): Store => {
 
     // Step 2: use the topSortedReversed atom list to recompute all affected atoms
     // Track what's changed, so that we can short circuit when possible
-    const changedAtoms = new Set<AnyAtom>(batch.D.keys())
     for (let i = topSortedReversed.length - 1; i >= 0; --i) {
       const [a, aState, prevEpochNumber] = topSortedReversed[i]!
       let hasChangedDeps = false
       for (const dep of aState.d.keys()) {
-        if (dep !== a && changedAtoms.has(dep)) {
+        if (dep !== a && batch.D.has(dep)) {
           hasChangedDeps = true
           break
         }
@@ -543,7 +542,6 @@ const buildStore = (...storeArgs: StoreArgs): Store => {
         mountDependencies(batch, a, aState)
         if (prevEpochNumber !== aState.n) {
           registerBatchAtom(batch, a, aState)
-          changedAtoms.add(a)
         }
       }
       delete aState.x

From a736376f3cdedb4777f51afde8474ae3a00d39a1 Mon Sep 17 00:00:00 2001
From: David Maskasky <dmaskasky@gmail.com>
Date: Sun, 5 Jan 2025 22:29:10 -0800
Subject: [PATCH 08/13] mark original atom changed

---
 src/vanilla/store.ts | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/vanilla/store.ts b/src/vanilla/store.ts
index d5943846cb..fc66990a33 100644
--- a/src/vanilla/store.ts
+++ b/src/vanilla/store.ts
@@ -463,7 +463,7 @@ const buildStore = (...storeArgs: StoreArgs): Store => {
     atom: Atom<Value>,
     atomState: AtomState<Value>,
   ) => {
-    const dependents = new Map<AnyAtom, AtomState>()
+    const dependents = new Map<AnyAtom, AtomState>([[atom, atomState]])
     const stack: [AnyAtom, AtomState][] = [[atom, atomState]]
     while (stack.length > 0) {
       const [a, aState] = stack.pop()!

From 0fc4d6e7c3aac2060bfc114db26303ade8963e29 Mon Sep 17 00:00:00 2001
From: David Maskasky <dmaskasky@gmail.com>
Date: Mon, 6 Jan 2025 00:08:45 -0800
Subject: [PATCH 09/13] flushBatch while condition includes dependents size

---
 src/vanilla/store.ts | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/vanilla/store.ts b/src/vanilla/store.ts
index fc66990a33..34c0e11522 100644
--- a/src/vanilla/store.ts
+++ b/src/vanilla/store.ts
@@ -172,7 +172,7 @@ const addDependency = <Value>(
 type BatchPriority = 0 | 1 | 2
 
 type Batch = [
-  /** finish recompute */
+  /** high priority listeners */
   priority0: Set<() => void>,
   /** atom listeners */
   priority1: Set<() => void>,

From 99515ad96786cbf271f45a750d07d9ba43d3e9db Mon Sep 17 00:00:00 2001
From: David Maskasky <dmaskasky@gmail.com>
Date: Tue, 7 Jan 2025 18:38:07 -0800
Subject: [PATCH 10/13] refactor: pass recompute as arg to flushBatch

---
 src/vanilla/store.ts         | 50 +++++++++++++++++++-----------------
 tests/vanilla/store.test.tsx | 30 ----------------------
 2 files changed, 26 insertions(+), 54 deletions(-)

diff --git a/src/vanilla/store.ts b/src/vanilla/store.ts
index 34c0e11522..1f48ad04f1 100644
--- 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
index a8d2279274..6ee45b2680 100644
--- 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])

From 0dc8a477be4cce212c30e72b8438d550ca948a2c Mon Sep 17 00:00:00 2001
From: David Maskasky <dmaskasky@gmail.com>
Date: Tue, 7 Jan 2025 19:22:34 -0800
Subject: [PATCH 11/13] rename getMountedDependents and simplify

---
 src/vanilla/store.ts | 21 ++++++++-------------
 1 file changed, 8 insertions(+), 13 deletions(-)

diff --git a/src/vanilla/store.ts b/src/vanilla/store.ts
index 1f48ad04f1..6e2822d8e3 100644
--- a/src/vanilla/store.ts
+++ b/src/vanilla/store.ts
@@ -443,23 +443,18 @@ const buildStore = (...storeArgs: StoreArgs): Store => {
   const readAtom = <Value>(atom: Atom<Value>): Value =>
     returnAtomValue(readAtomState(undefined, atom))
 
-  const getMountedOrBatchDependents = <Value>(
+  const getMountedDependents = <Value>(
     atomState: AtomState<Value>,
   ): Map<AnyAtom, AtomState> => {
-    const dependents = new Map<AnyAtom, AtomState>()
-    for (const a of atomState.m?.t || []) {
+    const mountedDependents = new Map<AnyAtom, AtomState>()
+    const dependents = new Set([...(atomState.m?.t || []), ...atomState.p])
+    for (const a of dependents) {
       const aState = ensureAtomState(a)
       if (aState.m) {
-        dependents.set(a, aState)
+        mountedDependents.set(a, aState)
       }
     }
-    for (const atomWithPendingPromise of atomState.p) {
-      dependents.set(
-        atomWithPendingPromise,
-        ensureAtomState(atomWithPendingPromise),
-      )
-    }
-    return dependents
+    return mountedDependents
   }
 
   const dirtyDependents = <Value>(atomState: AtomState<Value>) => {
@@ -472,7 +467,7 @@ const buildStore = (...storeArgs: StoreArgs): Store => {
         continue
       }
       aState.x = true
-      for (const [, s] of getMountedOrBatchDependents(aState)) {
+      for (const [, s] of getMountedDependents(aState)) {
         if (!dependents.has(s)) {
           dependents.add(s)
           stack.push(s)
@@ -519,7 +514,7 @@ const buildStore = (...storeArgs: StoreArgs): Store => {
       }
       visiting.add(a)
       // Push unvisited dependents onto the stack
-      for (const [d, s] of getMountedOrBatchDependents(aState)) {
+      for (const [d, s] of getMountedDependents(aState)) {
         if (a !== d && !visiting.has(d)) {
           stack.push([d, s])
         }

From 023480ce0ef5583140e3a0bb545534df57ba3993 Mon Sep 17 00:00:00 2001
From: David Maskasky <dmaskasky@gmail.com>
Date: Tue, 7 Jan 2025 19:24:21 -0800
Subject: [PATCH 12/13] remove unnecessary delete dirtybit

---
 src/vanilla/store.ts | 2 --
 1 file changed, 2 deletions(-)

diff --git a/src/vanilla/store.ts b/src/vanilla/store.ts
index 6e2822d8e3..3b3636a5ad 100644
--- a/src/vanilla/store.ts
+++ b/src/vanilla/store.ts
@@ -319,7 +319,6 @@ const buildStore = (...storeArgs: StoreArgs): Store => {
       atomState.v = valueOrPromise
     }
     delete atomState.e
-    delete atomState.x
     if (!hasPrevValue || !Object.is(prevValue, atomState.v)) {
       ++atomState.n
       if (pendingPromise) {
@@ -432,7 +431,6 @@ const buildStore = (...storeArgs: StoreArgs): Store => {
     } catch (error) {
       delete atomState.v
       atomState.e = error
-      delete atomState.x
       ++atomState.n
       return atomState
     } finally {

From 8ebb8f0a878f5ad4db0fb80665d07a4d024780eb Mon Sep 17 00:00:00 2001
From: David Maskasky <dmaskasky@gmail.com>
Date: Tue, 7 Jan 2025 20:55:46 -0800
Subject: [PATCH 13/13] add single recomputeDependents to batch

---
 src/vanilla/store.ts         | 67 ++++++++++++++++++------------------
 tests/vanilla/effect.test.ts | 20 +++++++++--
 2 files changed, 52 insertions(+), 35 deletions(-)

diff --git a/src/vanilla/store.ts b/src/vanilla/store.ts
index 3b3636a5ad..3af9793916 100644
--- a/src/vanilla/store.ts
+++ b/src/vanilla/store.ts
@@ -172,12 +172,12 @@ const addDependency = <Value>(
 type BatchPriority = 0 | 1 | 2
 
 type Batch = [
-  /** high priority listeners */
-  priority0: Set<() => void>,
+  /** finish recompute */
+  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>,
@@ -570,6 +570,7 @@ const buildStore = (...storeArgs: StoreArgs): Store => {
           if (prevEpochNumber !== aState.n) {
             dirtyDependents(aState)
             registerBatchAtom(batch, a, aState)
+            addBatchFunc(batch, 0, recomputeDependents)
           }
           return undefined as R
         } else {
@@ -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)
     }
   }
 
diff --git a/tests/vanilla/effect.test.ts b/tests/vanilla/effect.test.ts
index 3c1d158484..7ebe1b8d77 100644
--- a/tests/vanilla/effect.test.ts
+++ b/tests/vanilla/effect.test.ts
@@ -6,6 +6,7 @@ type Store = ReturnType<typeof createStore>
 type GetAtomState = Parameters<Parameters<Store['unstable_derive']>[0]>[0]
 type AtomState = NonNullable<ReturnType<GetAtomState>>
 type AnyAtom = Atom<unknown>
+type Batch = Parameters<NonNullable<AtomState['u']>>[0]
 
 type Cleanup = () => void
 type Effect = (get: Getter, set: Setter) => Cleanup | void
@@ -59,7 +60,7 @@ function syncEffect(effect: Effect): Atom<void> {
         store.set(refreshAtom, (v) => v + 1)
       } else {
         // unmount
-        batch[0].add(() => {
+        scheduleListenerAfterRecompute(batch, () => {
           ref.cleanup?.()
           delete ref.cleanup
         })
@@ -69,7 +70,22 @@ function syncEffect(effect: Effect): Atom<void> {
     internalAtomState.u = (batch) => {
       originalUpdateHook?.(batch)
       // update
-      batch[0].add(runEffect)
+      scheduleListenerAfterRecompute(batch, runEffect)
+    }
+    function scheduleListenerAfterRecompute(
+      batch: Batch,
+      listener: () => void,
+    ) {
+      const scheduleListener = () => {
+        batch[0].add(listener)
+      }
+      if (batch[0].size === 0) {
+        // no other listeners
+        // schedule after recomputeDependents
+        batch[0].add(scheduleListener)
+      } else {
+        scheduleListener()
+      }
     }
   }
   return atom((get) => {