use iterative approach in recompute dependents

src/vanilla/store.ts

@@ -439,6 +439,107 @@ const buildStore = (
     return dependents
   }
 
+  class Queue<T> {
+    private head = 0
+    private tail = 0
+    private items: { [key: number]: T } = {}
+    enqueue(item: T) {
+      this.items[this.tail++] = item
+    }
+    dequeue(): T | undefined {
+      if (this.isEmpty()) {
+        return undefined
+      }
+      const item = this.items[this.head]
+      delete this.items[this.head++]
+      return item
+    }
+    peek(): T | undefined {
+      return this.items[this.head]
+    }
+    isEmpty(): boolean {
+      return this.head === this.tail
+    }
+  }
+
+  // BFS implementation ~350ms for stack size limit nodes
+  function getSortedDependentsBFS(
+    pending: Pending,
+    startAtom: AnyAtom,
+    startAtomState: AtomState,
+  ): [[AnyAtom, AtomState, number][], Set<AnyAtom>] {
+    const sorted: [AnyAtom, AtomState, number][] = []
+    const visited = new Set<AnyAtom>()
+    const visiting = new Set<AnyAtom>()
+    const queue = new Queue<[AnyAtom, AtomState]>()
+    queue.enqueue([startAtom, startAtomState])
+    while (!queue.isEmpty()) {
+      const [currentAtom, currentAtomState] = queue.peek()!
+      if (visited.has(currentAtom)) {
+        queue.dequeue()
+        continue
+      }
+      if (visiting.has(currentAtom)) {
+        visiting.delete(currentAtom)
+        visited.add(currentAtom)
+        sorted.push([currentAtom, currentAtomState, currentAtomState.n])
+        continue
+      }
+      visiting.add(currentAtom)
+      for (const [dependentAtom, dependentAtomState] of getDependents(
+        pending,
+        currentAtom,
+        currentAtomState,
+      )) {
+        if (!visited.has(dependentAtom)) {
+          queue.enqueue([dependentAtom, dependentAtomState])
+        }
+      }
+    }
+    return [sorted, visited]
+  }
+
+  // DFS implementation ~280ms for stack size limit nodes
+  function getSortedDependents(
+    pending: Pending,
+    rootAtom: AnyAtom,
+    rootAtomState: AtomState,
+  ): [[AnyAtom, AtomState, number][], Set<AnyAtom>] {
+    const sorted: [atom: AnyAtom, atomState: AtomState, epochNumber: number][] =
+      []
+    const visiting = new Set<AnyAtom>()
+    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][] = [[rootAtom, rootAtomState]]
+    while (stack.length > 0) {
+      const [a, aState] = stack[stack.length - 1]!
+      if (visited.has(a)) {
+        // All dependents have been processed, now process this atom
+        stack.pop()
+        continue
+      }
+      if (visiting.has(a)) {
+        // The algorithm calls for pushing onto the front of the list. For
+        // performance, we will simply push onto the end, and then will iterate in
+        // reverse order later.
+        sorted.push([a, aState, aState.n])
+        // Atom has been visited but not yet processed
+        visited.add(a)
+        stack.pop()
+        continue
+      }
+      visiting.add(a)
+      // Push unvisited dependents onto the stack
+      for (const [d, s] of getDependents(pending, a, aState)) {
+        if (a !== d && !visiting.has(d)) {
+          stack.push([d, s])
+        }
+      }
+    }
+    return [sorted.reverse(), visiting]
+  }
+
   const recomputeDependents = <Value>(
     pending: Pending,
     atom: Atom<Value>,
@@ -450,35 +551,16 @@ const buildStore = (
 
     // Step 1: traverse the dependency graph to build the topsorted atom list
     // We don't bother to check for cycles, which simplifies the algorithm.
-    const topsortedAtoms: (readonly [
-      atom: AnyAtom,
-      atomState: AtomState,
-      epochNumber: number,
-    ])[] = []
-    const markedAtoms = new Set<AnyAtom>()
-    const visit = (a: AnyAtom, aState: AtomState) => {
-      if (markedAtoms.has(a)) {
-        return
-      }
-      markedAtoms.add(a)
-      for (const [d, s] of getDependents(pending, a, aState)) {
-        if (a !== d) {
-          visit(d, s)
-        }
-      }
-      // The algorithm calls for pushing onto the front of the list. For
-      // performance, we will simply push onto the end, and then will iterate in
-      // reverse order later.
-      topsortedAtoms.push([a, aState, aState.n])
-    }
-    // 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
-    visit(atom, atomState)
+    const [topsortedAtoms, markedAtoms] = getSortedDependents(
+      pending,
+      atom,
+      atomState,
+    )
+
     // Step 2: use the topsorted atom list to recompute all affected atoms
     // Track what's changed, so that we can short circuit when possible
     const changedAtoms = new Set<AnyAtom>([atom])
-    for (let i = topsortedAtoms.length - 1; i >= 0; --i) {
-      const [a, aState, prevEpochNumber] = topsortedAtoms[i]!
+    for (const [a, aState, prevEpochNumber] of topsortedAtoms) {
       let hasChangedDeps = false
       for (const dep of aState.d.keys()) {
         if (dep !== a && changedAtoms.has(dep)) {

tests/vanilla/store.test.tsx

@@ -1,7 +1,7 @@
 import { waitFor } from '@testing-library/react'
 import { assert, describe, expect, it, vi } from 'vitest'
 import { atom, createStore } from 'jotai/vanilla'
-import type { Getter } from 'jotai/vanilla'
+import type { Atom, Getter, PrimitiveAtom } from 'jotai/vanilla'
 
 it('should not fire on subscribe', async () => {
   const store = createStore()
@@ -910,3 +910,33 @@ it('should use the correct pending on unmount', () => {
   expect(store.get(a)).toBe(1)
   expect(aListener).toHaveBeenCalledTimes(1)
 })
+
+it('processes deep atom a graph beyond maxDepth', () => {
+  function getMaxDepth() {
+    let depth = 0
+    function d(): number {
+      ++depth
+      try {
+        return d()
+      } catch (error) {
+        return depth
+      }
+    }
+    return d()
+  }
+  const maxDepth = getMaxDepth()
+  const store = createStore()
+  const baseAtom = atom(0)
+  const atoms: [PrimitiveAtom<number>, ...rest: Atom<number>[]] = [baseAtom]
+  Array.from({ length: maxDepth }, (_, i) => {
+    const prevAtom = atoms[i]!
+    const a = atom((get) => get(prevAtom))
+    atoms.push(a)
+  })
+  const lastAtom = atoms[maxDepth]!
+  // store.get(lastAtom) // FIXME: This is causing a stack overflow
+  expect(() => store.sub(lastAtom, () => {})).not.toThrow()
+  // store.get(lastAtom) // FIXME: This is causing a stack overflow
+  expect(() => store.set(baseAtom, 1)).not.toThrow()
+  // store.set(lastAtom) // FIXME: This is causing a stack overflow
+})