diff --git a/src/vanilla/store.ts b/src/vanilla/store.ts index 5cd5bb719d..73eb9fb06c 100644 --- a/src/vanilla/store.ts +++ b/src/vanilla/store.ts @@ -104,6 +104,8 @@ type AtomState = { v?: Value /** Atom error */ e?: AnyError + /** Indicates that the atom value has been changed */ + x?: true } const isAtomStateInitialized = (atomState: AtomState) => @@ -164,6 +166,8 @@ const addDependency = ( type Batch = Readonly<{ /** Atom dependents map */ D: Map> + /** High priority functions */ + H: Set<() => void> /** Medium priority functions */ M: Set<() => void> /** Low priority functions */ @@ -172,10 +176,15 @@ type Batch = Readonly<{ const createBatch = (): Batch => ({ D: new Map(), + H: new Set(), M: new Set(), L: new Set(), }) +const addBatchFuncHigh = (batch: Batch, fn: () => void) => { + batch.H.add(fn) +} + const addBatchFuncMedium = (batch: Batch, fn: () => void) => { batch.M.add(fn) } @@ -232,6 +241,7 @@ const flushBatch = (batch: Batch) => { } while (batch.M.size || batch.L.size) { batch.D.clear() + copySetAndClear(batch.H).forEach(call) copySetAndClear(batch.M).forEach(call) copySetAndClear(batch.L).forEach(call) } @@ -305,11 +315,11 @@ const buildStore = ( addPendingPromiseToDependency(atom, valueOrPromise, getAtomState(a)) } atomState.v = valueOrPromise - delete atomState.e } else { atomState.v = valueOrPromise - delete atomState.e } + delete atomState.e + delete atomState.x if (!hasPrevValue || !Object.is(prevValue, atomState.v)) { ++atomState.n if (pendingPromise) { @@ -321,7 +331,6 @@ const buildStore = ( const readAtomState = ( batch: Batch | undefined, atom: Atom, - dirtyAtoms?: Set, ): AtomState => { const atomState = getAtomState(atom) // See if we can skip recomputing this atom. @@ -329,7 +338,7 @@ const buildStore = ( // If the atom is mounted, we can use cached atom state. // because it should have been updated by dependencies. // We can't use the cache if the atom is dirty. - if (atomState.m && !dirtyAtoms?.has(atom)) { + if (atomState.m && !atomState.x) { return atomState } // Otherwise, check if the dependencies have changed. @@ -339,7 +348,7 @@ const buildStore = ( ([a, n]) => // Recursively, read the atom state of the dependency, and // check if the atom epoch number is unchanged - readAtomState(batch, a, dirtyAtoms).n === n, + readAtomState(batch, a).n === n, ) ) { return atomState @@ -362,7 +371,7 @@ const buildStore = ( return returnAtomValue(aState) } // a !== atom - const aState = readAtomState(batch, a, dirtyAtoms) + const aState = readAtomState(batch, a) try { return returnAtomValue(aState) } finally { @@ -423,6 +432,7 @@ const buildStore = ( } catch (error) { delete atomState.v atomState.e = error + delete atomState.x ++atomState.n return atomState } finally { @@ -457,21 +467,26 @@ 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( + const recomputeDependents = ( batch: Batch, - rootAtom: AnyAtom, - rootAtomState: AtomState, - ): [[AnyAtom, AtomState, number][], Set] { - const sorted: [atom: AnyAtom, atomState: AtomState, epochNumber: number][] = - [] + atom: Atom, + atomState: AtomState, + ) => { + // 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 + // what's described here for simplicity and performance reasons: + // https://en.wikipedia.org/wiki/Topological_sorting#Depth-first_search + const topSortedReversed: [ + 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]] + const stack: [a: AnyAtom, aState: AtomState][] = [[atom, atomState]] while (stack.length > 0) { const [a, aState] = stack[stack.length - 1]! if (visited.has(a)) { @@ -483,9 +498,11 @@ const buildStore = ( // 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]) + 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 } @@ -497,44 +514,31 @@ const buildStore = ( } } } - return [sorted, visited] - } - - const recomputeDependents = ( - batch: Batch, - atom: Atom, - atomState: AtomState, - ) => { - // 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, markedAtoms] = getSortedDependents( - batch, - atom, - atomState, - ) - // Step 2: use the topsorted atom list to recompute all affected atoms + // 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([atom]) - for (let i = topsortedAtoms.length - 1; i >= 0; --i) { - const [a, aState, prevEpochNumber] = topsortedAtoms[i]! - let hasChangedDeps = false - for (const dep of aState.d.keys()) { - if (dep !== a && changedAtoms.has(dep)) { - hasChangedDeps = true - break + addBatchFuncHigh(batch, () => { + const changedAtoms = new Set([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 + } } - } - if (hasChangedDeps) { - readAtomState(batch, a, markedAtoms) - 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 } - markedAtoms.delete(a) - } + }) } const writeAtomState = ( diff --git a/tests/vanilla/dependency.test.tsx b/tests/vanilla/dependency.test.tsx index a530a67a4d..90f0c114e7 100644 --- a/tests/vanilla/dependency.test.tsx +++ b/tests/vanilla/dependency.test.tsx @@ -405,3 +405,23 @@ it('can cache reading an atom in write function (with mounting)', () => { store.set(w) expect(aReadCount).toBe(1) }) + +it('batches sync writes', () => { + const a = atom(0) + const b = atom((get) => get(a)) + const fetch = vi.fn() + const c = atom((get) => fetch(get(a))) + const w = atom(null, (get, set) => { + set(a, 1) + expect(get(b)).toBe(1) + expect(fetch).toHaveBeenCalledTimes(0) + }) + const store = createStore() + store.sub(b, () => {}) + store.sub(c, () => {}) + fetch.mockClear() + store.set(w) + expect(fetch).toHaveBeenCalledOnce() + expect(fetch).toBeCalledWith(1) + expect(store.get(a)).toBe(1) +}) diff --git a/tests/vanilla/store.test.tsx b/tests/vanilla/store.test.tsx index 0677aa092a..ea6a97e019 100644 --- a/tests/vanilla/store.test.tsx +++ b/tests/vanilla/store.test.tsx @@ -963,3 +963,23 @@ it('processes deep atom a graph beyond maxDepth', () => { expect(() => store.set(baseAtom, 1)).not.toThrow() // store.set(lastAtom) // FIXME: This is causing a stack overflow }) + +it('mounted atom should be recomputed eagerly', () => { + const result: string[] = [] + const a = atom(0) + const b = atom((get) => { + result.push('bRead') + return get(a) + }) + const store = createStore() + store.sub(a, () => { + result.push('aCallback') + }) + store.sub(b, () => { + result.push('bCallback') + }) + expect(result).toEqual(['bRead']) + result.splice(0) + store.set(a, 1) + expect(result).toEqual(['bRead', 'aCallback', 'bCallback']) +})