diff --git a/src/vanilla/store.ts b/src/vanilla/store.ts index a1e10d804c..b126cfbf2a 100644 --- a/src/vanilla/store.ts +++ b/src/vanilla/store.ts @@ -439,41 +439,62 @@ const buildStore = ( return dependents } + // This is a topological sort via depth-first search, slightly modified from + // what's described here for simplicity and performance reasons: + // https://en.wikipedia.org/wiki/Topological_sorting#Depth-first_search + function getSortedDependents( + pending: Pending, + rootAtom: AnyAtom, + rootAtomState: AtomState, + ): [[AnyAtom, AtomState, number][], Set] { + const sorted: [atom: AnyAtom, atomState: AtomState, epochNumber: number][] = + [] + const visiting = new Set() + const visited = new Set() + // 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, visited] + } + const recomputeDependents = ( pending: Pending, atom: Atom, atomState: AtomState, ) => { - // This is a topological sort via depth-first search, slightly modified from - // what's described here for simplicity and performance reasons: - // https://en.wikipedia.org/wiki/Topological_sorting#Depth-first_search - // 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() - 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([atom]) diff --git a/tests/vanilla/store.test.tsx b/tests/vanilla/store.test.tsx index 9c190fc3c1..2d978ad545 100644 --- a/tests/vanilla/store.test.tsx +++ b/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() @@ -933,3 +933,33 @@ it('should call subscribers after setAtom updates atom value on mount but not on expect(unmount).toHaveBeenCalledTimes(1) expect(listener).toHaveBeenCalledTimes(0) }) + +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, ...Atom[]] = [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 +})