diff --git a/src/vanilla/store.ts b/src/vanilla/store.ts index b126cfbf2a..accc08d727 100644 --- a/src/vanilla/store.ts +++ b/src/vanilla/store.ts @@ -165,6 +165,12 @@ type Pending = readonly [ dependents: Map>, atomStates: Map, functions: Set<() => void>, + // atoms read by readAtomState + // atoms written by writeAtomState + // 1. an atom is written to by writeAtomState + // 2. mark the atom as "dirty" + // 3. store the dependents of the dirty atom + // 4. ] const createPending = (): Pending => [new Map(), new Map(), new Set()] @@ -439,6 +445,28 @@ const buildStore = ( return dependents } + function getAllDependents( + pending: Pending, + atom: AnyAtom, + atomState: AtomState, + ): Set { + const visited = new Set() + const stack: [AnyAtom, AtomState][] = [[atom, atomState]] + while (stack.length > 0) { + const [a, aState] = stack.pop()! + if (visited.has(a)) { + continue + } + visited.add(a) + for (const [d, dState] of getDependents(pending, a, aState)) { + if (!visited.has(d)) { + stack.push([d, dState]) + } + } + } + return visited + } + // 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 @@ -525,8 +553,9 @@ const buildStore = ( ...args: Args ): Result => { let isSync = true - const getter: Getter = (a: Atom) => - returnAtomValue(readAtomState(pending, a)) + const getter: Getter = (a: Atom) => { + return returnAtomValue(readAtomState(pending, a)) + } const setter: Setter = ( a: WritableAtom, ...args: As