From b7dac4ccac69e159395016c5d2d5d28a06dc9f99 Mon Sep 17 00:00:00 2001
From: David Maskasky <dmaskasky@gmail.com>
Date: Mon, 30 Mar 2026 00:10:50 -0700
Subject: [PATCH] refactor: invalidateDependents supports batching

---
 src/vanilla/internals.ts         | 32 ++++++++++++++++++++++----------
 tests/vanilla/internals.test.tsx |  2 +-
 2 files changed, 23 insertions(+), 11 deletions(-)

diff --git a/src/vanilla/internals.ts b/src/vanilla/internals.ts
index a3bac1a448..bee54d35dc 100644
--- a/src/vanilla/internals.ts
+++ b/src/vanilla/internals.ts
@@ -109,7 +109,7 @@ type ReadAtomState = <Value>(
   store: Store,
   atom: Atom<Value>,
 ) => AtomState<Value>
-type InvalidateDependents = (store: Store, atom: AnyAtom) => void
+type InvalidateDependents = (store: Store, atoms: Iterable<AnyAtom>) => void
 type WriteAtomState = <Value, Args extends unknown[], Result>(
   store: Store,
   atom: WritableAtom<Value, Args, Result>,
@@ -690,21 +690,27 @@ const BUILDING_BLOCK_readAtomState: ReadAtomState = (store, atom) => {
 
 const BUILDING_BLOCK_invalidateDependents: InvalidateDependents = (
   store,
-  atom,
+  atoms,
 ) => {
   const buildingBlocks = getInternalBuildingBlocks(store)
   const mountedMap = buildingBlocks[1]
   const invalidatedAtoms = buildingBlocks[2]
   const ensureAtomState = buildingBlocks[11]
-  const stack: AnyAtom[] = [atom]
-  while (stack.length) {
-    const a = stack.pop()!
-    const aState = ensureAtomState(store, a)
+  const atomStack: AnyAtom[] = []
+  const stateStack: AtomState[] = []
+  for (const atom of atoms) {
+    atomStack.push(atom)
+    stateStack.push(ensureAtomState(store, atom))
+  }
+  while (atomStack.length) {
+    const a = atomStack.pop()!
+    const aState = stateStack.pop()!
     for (const d of getMountedOrPendingDependents(a, aState, mountedMap)) {
       const dState = ensureAtomState(store, d)
       if (invalidatedAtoms.get(d) !== dState.n) {
         invalidatedAtoms.set(d, dState.n)
-        stack.push(d)
+        atomStack.push(d)
+        stateStack.push(dState)
       }
     }
   }
@@ -752,7 +758,7 @@ const BUILDING_BLOCK_writeAtomState: WriteAtomState = (
         if (prevEpochNumber !== aState.n) {
           ++storeEpochHolder[0]
           changedAtoms.add(a)
-          invalidateDependents(store, a)
+          invalidateDependents(store, [a])
           storeHooks.c?.(a)
         }
         return undefined as R
@@ -785,6 +791,7 @@ const BUILDING_BLOCK_mountDependencies: MountDependencies = (store, atom) => {
   const atomState = ensureAtomState(store, atom)
   const mounted = mountedMap.get(atom)
   if (mounted) {
+    const staleDeps = new Set<AnyAtom>()
     for (const [a, n] of atomState.d) {
       if (!mounted.d.has(a)) {
         const aState = ensureAtomState(store, a)
@@ -793,11 +800,16 @@ const BUILDING_BLOCK_mountDependencies: MountDependencies = (store, atom) => {
         mounted.d.add(a)
         if (n !== aState.n) {
           changedAtoms.add(a)
-          invalidateDependents(store, a)
-          storeHooks.c?.(a)
+          staleDeps.add(a)
         }
       }
     }
+    if (staleDeps.size) {
+      invalidateDependents(store, staleDeps)
+      for (const a of staleDeps) {
+        storeHooks.c?.(a)
+      }
+    }
     for (const a of mounted.d) {
       if (!atomState.d.has(a)) {
         mounted.d.delete(a)
diff --git a/tests/vanilla/internals.test.tsx b/tests/vanilla/internals.test.tsx
index acc2df4634..6777d8ee56 100644
--- a/tests/vanilla/internals.test.tsx
+++ b/tests/vanilla/internals.test.tsx
@@ -210,7 +210,7 @@ describe('internals', () => {
 
     const unsub = store.sub(leafAtom, () => {})
     const invalidateDependents = INTERNAL_getBuildingBlocks(store)[15]
-    expect(() => invalidateDependents(store, baseAtom)).not.toThrow()
+    expect(() => invalidateDependents(store, [baseAtom])).not.toThrow()
     unsub()
   })
 })