From 6908f9c7181692f814d0171f2f3a20d9f423741f 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/24] 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 f6fdeab7c6..39a613f452 100644
--- a/tests/vanilla/store.test.tsx
+++ b/tests/vanilla/store.test.tsx
@@ -1086,3 +1086,64 @@ it('should pass store and atomState to the atom initializer', () => {
   }
   store.get(a)
 })
+
+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 ade5a39077ec403d5d59f183a5bc901e11fded8c 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/24] 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 6f1a0862e5..334d55f2d0 100644
--- a/src/vanilla/store.ts
+++ b/src/vanilla/store.ts
@@ -481,11 +481,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
@@ -500,7 +496,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)) {
@@ -531,27 +530,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)
   }
@@ -581,7 +582,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 39a613f452..16c4f82a57 100644
--- a/tests/vanilla/store.test.tsx
+++ b/tests/vanilla/store.test.tsx
@@ -1096,15 +1096,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++
@@ -1115,35 +1115,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 c985ca878951e63089c4099ee4f4ddbe29a2146a 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/24] centralize finishRecompute

---
 src/vanilla/store.ts         | 46 +++++++++++++++++++++++++-----------
 tests/vanilla/effect.test.ts | 42 ++------------------------------
 2 files changed, 34 insertions(+), 54 deletions(-)

diff --git a/src/vanilla/store.ts b/src/vanilla/store.ts
index 334d55f2d0..931c38cc71 100644
--- a/src/vanilla/store.ts
+++ b/src/vanilla/store.ts
@@ -58,9 +58,9 @@ const patchPromiseForCancelability = <T>(promise: PromiseLike<T>) => {
 }
 
 const isPromiseLike = (
-  x: unknown,
-): x is PromiseLike<unknown> & { onCancel?: (fn: CancelHandler) => void } =>
-  typeof (x as any)?.then === 'function'
+  p: unknown,
+): p is PromiseLike<unknown> & { onCancel?: (fn: CancelHandler) => void } =>
+  typeof (p as any)?.then === 'function'
 
 /**
  * State tracked for mounted atoms. An atom is considered "mounted" if it has a
@@ -185,6 +185,8 @@ type Batch = [
 ] & {
   /** Atom dependents map */
   D: Map<AnyAtom, Set<AnyAtom>>
+  /** Recompute dependents function */
+  R?: (batch: Batch) => void
 }
 
 const createBatch = (): Batch =>
@@ -241,6 +243,7 @@ const flushBatch = (batch: Batch) => {
     }
   }
   while (batch.some((channel) => channel.size)) {
+    batch.R?.(batch)
     batch.D.clear()
     for (const channel of batch) {
       channel.forEach(call)
@@ -481,6 +484,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.
@@ -498,7 +524,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]!
@@ -514,8 +540,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
       }
@@ -530,10 +554,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]!
@@ -554,7 +574,6 @@ const buildStore = (...storeArgs: StoreArgs): Store => {
       }
       delete aState.x
     }
-    addBatchFunc(batch, 0, finishRecompute)
   }
 
   const writeAtomState = <Value, Args extends unknown[], Result>(
@@ -581,10 +600,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 ba3075aaa63848508022626f83eb7d76ee3d63c1 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/24] 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 16c4f82a57..7bac63cf48 100644
--- a/tests/vanilla/store.test.tsx
+++ b/tests/vanilla/store.test.tsx
@@ -1115,35 +1115,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 3e0e3243d9991c57a57621703a5e52275ab46552 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/24] 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 7bac63cf48..01baeda4cb 100644
--- a/tests/vanilla/store.test.tsx
+++ b/tests/vanilla/store.test.tsx
@@ -1088,10 +1088,6 @@ it('should pass store and atomState to the atom initializer', () => {
 })
 
 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 7a66fd039c7c5f296beec1bbf738ed50eb58eada 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/24] simplify test

---
 tests/vanilla/store.test.tsx | 54 ++++++++++--------------------------
 1 file changed, 14 insertions(+), 40 deletions(-)

diff --git a/tests/vanilla/store.test.tsx b/tests/vanilla/store.test.tsx
index 01baeda4cb..145f629cd1 100644
--- a/tests/vanilla/store.test.tsx
+++ b/tests/vanilla/store.test.tsx
@@ -1087,46 +1087,20 @@ it('should pass store and atomState to the atom initializer', () => {
   store.get(a)
 })
 
-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')
-  })
-  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 c6e4d81a4dd4192a77cab80db3e47606d6330edb 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/24] 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 931c38cc71..bc33ed1427 100644
--- a/src/vanilla/store.ts
+++ b/src/vanilla/store.ts
@@ -554,12 +554,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
         }
@@ -569,7 +568,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 ea0ed26c664d000a59d74133387ca451ce81aef3 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/24] 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 bc33ed1427..a8ed1d43a6 100644
--- a/src/vanilla/store.ts
+++ b/src/vanilla/store.ts
@@ -489,7 +489,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 3b435d791326033deeb88f1e77050cb6b8993b51 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/24] flushBatch while condition includes dependents size

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

diff --git a/src/vanilla/store.ts b/src/vanilla/store.ts
index a8ed1d43a6..8fca45cb4f 100644
--- a/src/vanilla/store.ts
+++ b/src/vanilla/store.ts
@@ -176,7 +176,7 @@ const addDependency = <Value>(
 type BatchPriority = 0 | 1 | 2
 
 type Batch = [
-  /** finish recompute */
+  /** high priority listeners */
   priority0: Set<() => void>,
   /** atom listeners */
   priority1: Set<() => void>,
@@ -242,7 +242,7 @@ const flushBatch = (batch: Batch) => {
       }
     }
   }
-  while (batch.some((channel) => channel.size)) {
+  while (batch.D.size || batch.some((channel) => channel.size)) {
     batch.R?.(batch)
     batch.D.clear()
     for (const channel of batch) {

From d8b73b530a985c63082ed41f88194f3f2f3dac36 Mon Sep 17 00:00:00 2001
From: daishi <daishi@axlight.com>
Date: Tue, 7 Jan 2025 21:27:13 +0900
Subject: [PATCH 10/24] recover the test

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

diff --git a/tests/vanilla/store.test.tsx b/tests/vanilla/store.test.tsx
index 425063617e..6ee45b2680 100644
--- a/tests/vanilla/store.test.tsx
+++ b/tests/vanilla/store.test.tsx
@@ -1108,3 +1108,21 @@ it('recomputes dependents of unmounted atoms', () => {
   store.set(w)
   expect(bRead).not.toHaveBeenCalled()
 })
+
+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])
+  })
+  const store = createStore()
+  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 e5a43d98219367f02291f4dd64ba94566135a440 Mon Sep 17 00:00:00 2001
From: David Maskasky <dmaskasky@gmail.com>
Date: Sat, 4 Jan 2025 10:07:17 -0700
Subject: [PATCH 11/24] 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 15a1f137a5b17feb504e4d14e779756aca447da4 Mon Sep 17 00:00:00 2001
From: David Maskasky <dmaskasky@gmail.com>
Date: Sun, 5 Jan 2025 20:18:17 -0800
Subject: [PATCH 12/24] 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 bbfaee5c296cf3467d2ca25686a9220e4ae0fe7b Mon Sep 17 00:00:00 2001
From: David Maskasky <dmaskasky@gmail.com>
Date: Sun, 5 Jan 2025 20:53:26 -0800
Subject: [PATCH 13/24] 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 27ce27553f180587c9eb7c28c13bf388286ce5e7 Mon Sep 17 00:00:00 2001
From: David Maskasky <dmaskasky@gmail.com>
Date: Sun, 5 Jan 2025 21:02:31 -0800
Subject: [PATCH 14/24] 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 cca2b1823fd216bbe4d67feef276401e22641fd8 Mon Sep 17 00:00:00 2001
From: David Maskasky <dmaskasky@gmail.com>
Date: Sun, 5 Jan 2025 21:23:01 -0800
Subject: [PATCH 15/24] 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 6a892f5d49bf342169876b92bb9e34bef3fd98f8 Mon Sep 17 00:00:00 2001
From: David Maskasky <dmaskasky@gmail.com>
Date: Sun, 5 Jan 2025 22:21:00 -0800
Subject: [PATCH 16/24] 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 c9209466b54a087eca56c6cf3591dc6979056af9 Mon Sep 17 00:00:00 2001
From: David Maskasky <dmaskasky@gmail.com>
Date: Sun, 5 Jan 2025 22:26:17 -0800
Subject: [PATCH 17/24] 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 048359a26ee21fb6a59889ec41b10e3f37e6f252 Mon Sep 17 00:00:00 2001
From: David Maskasky <dmaskasky@gmail.com>
Date: Sun, 5 Jan 2025 22:29:10 -0800
Subject: [PATCH 18/24] 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 ba84b1898a291eca98d21da2c4944597ba4b59ea Mon Sep 17 00:00:00 2001
From: David Maskasky <dmaskasky@gmail.com>
Date: Mon, 6 Jan 2025 00:08:45 -0800
Subject: [PATCH 19/24] 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 a488c0af354a0f6ee21d301404543987c1fb574d Mon Sep 17 00:00:00 2001
From: David Maskasky <dmaskasky@gmail.com>
Date: Tue, 7 Jan 2025 09:16:51 -0800
Subject: [PATCH 20/24] refactor: move flushBatch in buildStore and call
 recomputeDependents directly

---
 src/vanilla/store.ts         | 78 +++++++++++++++++-------------------
 tests/vanilla/store.test.tsx | 30 --------------
 2 files changed, 37 insertions(+), 71 deletions(-)

diff --git a/src/vanilla/store.ts b/src/vanilla/store.ts
index 34c0e11522..c9b35bc5e0 100644
--- a/src/vanilla/store.ts
+++ b/src/vanilla/store.ts
@@ -209,31 +209,6 @@ const registerBatchAtom = (
   }
 }
 
-const flushBatch = (batch: Batch) => {
-  let error: AnyError
-  let hasError = false
-  const call = (fn: () => void) => {
-    try {
-      fn()
-    } catch (e) {
-      if (!hasError) {
-        error = e
-        hasError = true
-      }
-    }
-  }
-  while (batch.C.size || batch.some((channel) => channel.size)) {
-    batch.C.clear()
-    for (const channel of batch) {
-      channel.forEach(call)
-      channel.clear()
-    }
-  }
-  if (hasError) {
-    throw error
-  }
-}
-
 // internal & unstable type
 type StoreArgs = readonly [
   getAtomState: <Value>(atom: Atom<Value>) => AtomState<Value> | undefined,
@@ -458,24 +433,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 +468,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 +503,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 +519,32 @@ const buildStore = (...storeArgs: StoreArgs): Store => {
     }
   }
 
+  const flushBatch = (batch: Batch) => {
+    let error: AnyError
+    let hasError = false
+    const call = (fn: () => void) => {
+      try {
+        fn()
+      } catch (e) {
+        if (!hasError) {
+          error = e
+          hasError = true
+        }
+      }
+    }
+    while (batch.C.size || batch.some((channel) => channel.size)) {
+      recomputeDependents(batch)
+      batch.C.clear()
+      for (const channel of batch) {
+        channel.forEach(call)
+        channel.clear()
+      }
+    }
+    if (hasError) {
+      throw error
+    }
+  }
+
   const writeAtomState = <Value, Args extends unknown[], Result>(
     batch: Batch,
     atom: WritableAtom<Value, Args, Result>,
@@ -572,8 +569,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
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 ef52454d802b32e88751b2274b4dce512450ffdb Mon Sep 17 00:00:00 2001
From: daishi <daishi@axlight.com>
Date: Wed, 8 Jan 2025 07:24:23 +0900
Subject: [PATCH 21/24] remove log

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

diff --git a/tests/vanilla/store.test.tsx b/tests/vanilla/store.test.tsx
index 5f15378093..dd81545d5a 100644
--- a/tests/vanilla/store.test.tsx
+++ b/tests/vanilla/store.test.tsx
@@ -1091,7 +1091,6 @@ it('recomputes dependents of unmounted atoms', () => {
   const a = atom(0)
   a.debugLabel = 'a'
   const bRead = vi.fn((get: Getter) => {
-    console.log('bRead')
     return get(a)
   })
   const b = atom(bRead)
@@ -1147,38 +1146,29 @@ it('runs recomputeDependents on atoms in the correct order', async () => {
     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 5c503d38ef0780f51021e13a28bb1815a078b82d Mon Sep 17 00:00:00 2001
From: daishi <daishi@axlight.com>
Date: Wed, 8 Jan 2025 08:20:37 +0900
Subject: [PATCH 22/24] use map for batch

---
 src/vanilla/store.ts         | 40 ++++++++++++++++--------------------
 tests/vanilla/effect.test.ts | 15 +++++++-------
 2 files changed, 26 insertions(+), 29 deletions(-)

diff --git a/src/vanilla/store.ts b/src/vanilla/store.ts
index 6c20762127..8a36876f8d 100644
--- a/src/vanilla/store.ts
+++ b/src/vanilla/store.ts
@@ -173,25 +173,22 @@ type BatchPriority = 0 | 1 | 2
 
 type Batch = [
   /** finish recompute */
-  priority0: Set<() => void>,
+  priority0: Map<object, () => void>,
   /** atom listeners */
-  priority1: Set<() => void>,
+  priority1: Map<object, () => void>,
   /** atom mount hooks */
-  priority2: Set<() => void>,
-] & {
-  /** changed Atoms */
-  C: Set<AnyAtom>
-}
+  priority2: Map<object, () => void>,
+]
 
-const createBatch = (): Batch =>
-  Object.assign([new Set(), new Set(), new Set()], { C: new Set() }) as Batch
+const createBatch = (): Batch => [new Map(), new Map(), new Map()]
 
 const addBatchFunc = (
   batch: Batch,
   priority: BatchPriority,
+  key: object,
   fn: () => void,
 ) => {
-  batch[priority].add(fn)
+  batch[priority].set(key, fn)
 }
 
 const registerBatchAtom = (
@@ -199,14 +196,13 @@ const registerBatchAtom = (
   atom: AnyAtom,
   atomState: AtomState,
 ) => {
-  if (!batch.C.has(atom)) {
-    batch.C.add(atom)
-    atomState.u?.(batch)
-    const scheduleListeners = () => {
-      atomState.m?.l.forEach((listener) => addBatchFunc(batch, 1, listener))
-    }
-    addBatchFunc(batch, 1, scheduleListeners)
+  atomState.u?.(batch)
+  const scheduleListeners = () => {
+    atomState.m?.l.forEach((listener) =>
+      addBatchFunc(batch, 1, listener, listener),
+    )
   }
+  addBatchFunc(batch, 1, atom, scheduleListeners)
 }
 
 const flushBatch = (batch: Batch) => {
@@ -222,8 +218,7 @@ const flushBatch = (batch: Batch) => {
       }
     }
   }
-  while (batch.C.size || batch.some((channel) => channel.size)) {
-    batch.C.clear()
+  while (batch.some((channel) => channel.size)) {
     for (const channel of batch) {
       channel.forEach(call)
       channel.clear()
@@ -530,7 +525,7 @@ const buildStore = (...storeArgs: StoreArgs): Store => {
         delete aState.x
       }
     }
-    addBatchFunc(batch, 0, finishRecompute)
+    addBatchFunc(batch, 0, finishRecompute, finishRecompute)
   }
 
   const writeAtomState = <Value, Args extends unknown[], Result>(
@@ -660,7 +655,7 @@ const buildStore = (...storeArgs: StoreArgs): Store => {
             mounted.u = (batch) => createInvocationContext(batch, onUnmount)
           }
         }
-        addBatchFunc(batch, 2, processOnMount)
+        addBatchFunc(batch, 2, processOnMount, processOnMount)
       }
     }
     return atomState.m
@@ -679,7 +674,8 @@ const buildStore = (...storeArgs: StoreArgs): Store => {
       // unmount self
       const onUnmount = atomState.m.u
       if (onUnmount) {
-        addBatchFunc(batch, 2, () => onUnmount(batch))
+        const processOnUnmount = () => onUnmount(batch)
+        addBatchFunc(batch, 2, processOnUnmount, processOnUnmount)
       }
       delete atomState.m
       atomState.h?.(batch)
diff --git a/tests/vanilla/effect.test.ts b/tests/vanilla/effect.test.ts
index 697d6a6733..678c9f88a8 100644
--- a/tests/vanilla/effect.test.ts
+++ b/tests/vanilla/effect.test.ts
@@ -63,10 +63,11 @@ function syncEffect(effect: Effect): Atom<void> {
       } else {
         // unmount
         const syncEffectChannel = ensureBatchChannel(batch)
-        syncEffectChannel.add(() => {
+        const processCleanup = () => {
           ref.cleanup?.()
           delete ref.cleanup
-        })
+        }
+        syncEffectChannel.set(processCleanup, processCleanup)
       }
     }
     const originalUpdateHook = internalAtomState.u
@@ -74,7 +75,7 @@ function syncEffect(effect: Effect): Atom<void> {
       originalUpdateHook?.(batch)
       // update
       const syncEffectChannel = ensureBatchChannel(batch)
-      syncEffectChannel.add(runEffect)
+      syncEffectChannel.set(runEffect, runEffect)
     }
   }
   return atom((get) => {
@@ -83,7 +84,7 @@ function syncEffect(effect: Effect): Atom<void> {
 }
 
 type BatchWithSyncEffect = Batch & {
-  [syncEffectChannelSymbol]?: Set<() => void>
+  [syncEffectChannelSymbol]?: Map<object, () => void>
 }
 function ensureBatchChannel(batch: BatchWithSyncEffect) {
   // ensure continuation of the flushBatch while loop
@@ -92,11 +93,11 @@ function ensureBatchChannel(batch: BatchWithSyncEffect) {
     throw new Error('batch[1] must be present')
   }
   if (!batch[syncEffectChannelSymbol]) {
-    batch[syncEffectChannelSymbol] = new Set<() => void>()
+    batch[syncEffectChannelSymbol] = new Map<object, () => void>()
     batch[1] = {
       ...originalQueue,
-      add(item) {
-        originalQueue.add(item)
+      set(key, item) {
+        originalQueue.set(key, item)
         return this
       },
       clear() {

From c08c09cb676bec1cf2d299f878e9eae20eca1aca Mon Sep 17 00:00:00 2001
From: daishi <daishi@axlight.com>
Date: Wed, 8 Jan 2025 08:29:23 +0900
Subject: [PATCH 23/24] make dirty flag a number

---
 src/vanilla/store.ts | 10 ++++------
 1 file changed, 4 insertions(+), 6 deletions(-)

diff --git a/src/vanilla/store.ts b/src/vanilla/store.ts
index 8a36876f8d..25affd5eba 100644
--- a/src/vanilla/store.ts
+++ b/src/vanilla/store.ts
@@ -115,7 +115,7 @@ type AtomState<Value = AnyValue> = {
   /** Atom error */
   e?: AnyError
   /** Indicates that the atom value has been changed */
-  x?: true
+  x: number
 }
 
 const isAtomStateInitialized = <Value>(atomState: AtomState<Value>) =>
@@ -285,7 +285,7 @@ const buildStore = (...storeArgs: StoreArgs): Store => {
     }
     let atomState = getAtomState(atom)
     if (!atomState) {
-      atomState = { d: new Map(), p: new Set(), n: 0 }
+      atomState = { d: new Map(), p: new Set(), n: 0, x: 0 }
       setAtomState(atom, atomState)
       atomOnInit?.(atom, store)
     }
@@ -310,7 +310,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) {
@@ -423,7 +422,6 @@ const buildStore = (...storeArgs: StoreArgs): Store => {
     } catch (error) {
       delete atomState.v
       atomState.e = error
-      delete atomState.x
       ++atomState.n
       return atomState
     } finally {
@@ -488,7 +486,7 @@ const buildStore = (...storeArgs: StoreArgs): Store => {
         // Atom has been visited but not yet processed
         visited.add(a)
         // Mark atom dirty
-        aState.x = true
+        ++aState.x
         stack.pop()
         continue
       }
@@ -522,7 +520,7 @@ const buildStore = (...storeArgs: StoreArgs): Store => {
             changedAtoms.add(a)
           }
         }
-        delete aState.x
+        --aState.x
       }
     }
     addBatchFunc(batch, 0, finishRecompute, finishRecompute)

From 628f5692c519072ebeaeb4d3321795c9807dcd7e Mon Sep 17 00:00:00 2001
From: daishi <daishi@axlight.com>
Date: Wed, 8 Jan 2025 08:35:45 +0900
Subject: [PATCH 24/24] Revert "use map for batch"

This reverts commit 5c503d38ef0780f51021e13a28bb1815a078b82d.
---
 src/vanilla/store.ts         | 40 ++++++++++++++++++++----------------
 tests/vanilla/effect.test.ts | 15 +++++++-------
 2 files changed, 29 insertions(+), 26 deletions(-)

diff --git a/src/vanilla/store.ts b/src/vanilla/store.ts
index 25affd5eba..fcc2fb7ea7 100644
--- a/src/vanilla/store.ts
+++ b/src/vanilla/store.ts
@@ -173,22 +173,25 @@ type BatchPriority = 0 | 1 | 2
 
 type Batch = [
   /** finish recompute */
-  priority0: Map<object, () => void>,
+  priority0: Set<() => void>,
   /** atom listeners */
-  priority1: Map<object, () => void>,
+  priority1: Set<() => void>,
   /** atom mount hooks */
-  priority2: Map<object, () => void>,
-]
+  priority2: Set<() => void>,
+] & {
+  /** changed Atoms */
+  C: Set<AnyAtom>
+}
 
-const createBatch = (): Batch => [new Map(), new Map(), new Map()]
+const createBatch = (): Batch =>
+  Object.assign([new Set(), new Set(), new Set()], { C: new Set() }) as Batch
 
 const addBatchFunc = (
   batch: Batch,
   priority: BatchPriority,
-  key: object,
   fn: () => void,
 ) => {
-  batch[priority].set(key, fn)
+  batch[priority].add(fn)
 }
 
 const registerBatchAtom = (
@@ -196,13 +199,14 @@ const registerBatchAtom = (
   atom: AnyAtom,
   atomState: AtomState,
 ) => {
-  atomState.u?.(batch)
-  const scheduleListeners = () => {
-    atomState.m?.l.forEach((listener) =>
-      addBatchFunc(batch, 1, listener, listener),
-    )
+  if (!batch.C.has(atom)) {
+    batch.C.add(atom)
+    atomState.u?.(batch)
+    const scheduleListeners = () => {
+      atomState.m?.l.forEach((listener) => addBatchFunc(batch, 1, listener))
+    }
+    addBatchFunc(batch, 1, scheduleListeners)
   }
-  addBatchFunc(batch, 1, atom, scheduleListeners)
 }
 
 const flushBatch = (batch: Batch) => {
@@ -218,7 +222,8 @@ const flushBatch = (batch: Batch) => {
       }
     }
   }
-  while (batch.some((channel) => channel.size)) {
+  while (batch.C.size || batch.some((channel) => channel.size)) {
+    batch.C.clear()
     for (const channel of batch) {
       channel.forEach(call)
       channel.clear()
@@ -523,7 +528,7 @@ const buildStore = (...storeArgs: StoreArgs): Store => {
         --aState.x
       }
     }
-    addBatchFunc(batch, 0, finishRecompute, finishRecompute)
+    addBatchFunc(batch, 0, finishRecompute)
   }
 
   const writeAtomState = <Value, Args extends unknown[], Result>(
@@ -653,7 +658,7 @@ const buildStore = (...storeArgs: StoreArgs): Store => {
             mounted.u = (batch) => createInvocationContext(batch, onUnmount)
           }
         }
-        addBatchFunc(batch, 2, processOnMount, processOnMount)
+        addBatchFunc(batch, 2, processOnMount)
       }
     }
     return atomState.m
@@ -672,8 +677,7 @@ const buildStore = (...storeArgs: StoreArgs): Store => {
       // unmount self
       const onUnmount = atomState.m.u
       if (onUnmount) {
-        const processOnUnmount = () => onUnmount(batch)
-        addBatchFunc(batch, 2, processOnUnmount, processOnUnmount)
+        addBatchFunc(batch, 2, () => onUnmount(batch))
       }
       delete atomState.m
       atomState.h?.(batch)
diff --git a/tests/vanilla/effect.test.ts b/tests/vanilla/effect.test.ts
index 678c9f88a8..697d6a6733 100644
--- a/tests/vanilla/effect.test.ts
+++ b/tests/vanilla/effect.test.ts
@@ -63,11 +63,10 @@ function syncEffect(effect: Effect): Atom<void> {
       } else {
         // unmount
         const syncEffectChannel = ensureBatchChannel(batch)
-        const processCleanup = () => {
+        syncEffectChannel.add(() => {
           ref.cleanup?.()
           delete ref.cleanup
-        }
-        syncEffectChannel.set(processCleanup, processCleanup)
+        })
       }
     }
     const originalUpdateHook = internalAtomState.u
@@ -75,7 +74,7 @@ function syncEffect(effect: Effect): Atom<void> {
       originalUpdateHook?.(batch)
       // update
       const syncEffectChannel = ensureBatchChannel(batch)
-      syncEffectChannel.set(runEffect, runEffect)
+      syncEffectChannel.add(runEffect)
     }
   }
   return atom((get) => {
@@ -84,7 +83,7 @@ function syncEffect(effect: Effect): Atom<void> {
 }
 
 type BatchWithSyncEffect = Batch & {
-  [syncEffectChannelSymbol]?: Map<object, () => void>
+  [syncEffectChannelSymbol]?: Set<() => void>
 }
 function ensureBatchChannel(batch: BatchWithSyncEffect) {
   // ensure continuation of the flushBatch while loop
@@ -93,11 +92,11 @@ function ensureBatchChannel(batch: BatchWithSyncEffect) {
     throw new Error('batch[1] must be present')
   }
   if (!batch[syncEffectChannelSymbol]) {
-    batch[syncEffectChannelSymbol] = new Map<object, () => void>()
+    batch[syncEffectChannelSymbol] = new Set<() => void>()
     batch[1] = {
       ...originalQueue,
-      set(key, item) {
-        originalQueue.set(key, item)
+      add(item) {
+        originalQueue.add(item)
         return this
       },
       clear() {