diff --git a/src/vanilla/store.ts b/src/vanilla/store.ts index d5010a8ccf..5cd5bb719d 100644 --- a/src/vanilla/store.ts +++ b/src/vanilla/store.ts @@ -77,7 +77,7 @@ type Mounted = { /** Set of mounted atoms that depends on the atom. */ readonly t: Set /** Function to run when the atom is unmounted. */ - u?: (pending: Pending) => void + u?: (batch: Batch) => void } /** @@ -138,7 +138,7 @@ const addPendingPromiseToDependency = ( } const addDependency = ( - pending: Pending | undefined, + batch: Batch | undefined, atom: Atom, atomState: AtomState, a: AnyAtom, @@ -152,53 +152,72 @@ const addDependency = ( addPendingPromiseToDependency(atom, atomState.v, aState) } aState.m?.t.add(atom) - if (pending) { - addPendingDependent(pending, a, atom) + if (batch) { + addBatchAtomDependent(batch, a, atom) } } // -// Pending +// Batch // -type Pending = readonly [ - dependents: Map>, - atomStates: Map, - functions: Set<() => void>, -] +type Batch = Readonly<{ + /** Atom dependents map */ + D: Map> + /** Medium priority functions */ + M: Set<() => void> + /** Low priority functions */ + L: Set<() => void> +}> + +const createBatch = (): Batch => ({ + D: new Map(), + M: new Set(), + L: new Set(), +}) + +const addBatchFuncMedium = (batch: Batch, fn: () => void) => { + batch.M.add(fn) +} -const createPending = (): Pending => [new Map(), new Map(), new Set()] +const addBatchFuncLow = (batch: Batch, fn: () => void) => { + batch.L.add(fn) +} -const addPendingAtom = ( - pending: Pending, +const registerBatchAtom = ( + batch: Batch, atom: AnyAtom, atomState: AtomState, ) => { - if (!pending[0].has(atom)) { - pending[0].set(atom, new Set()) + if (!batch.D.has(atom)) { + batch.D.set(atom, new Set()) + addBatchFuncMedium(batch, () => { + atomState.m?.l.forEach((listener) => listener()) + }) } - pending[1].set(atom, atomState) } -const addPendingDependent = ( - pending: Pending, +const addBatchAtomDependent = ( + batch: Batch, atom: AnyAtom, dependent: AnyAtom, ) => { - const dependents = pending[0].get(atom) + const dependents = batch.D.get(atom) if (dependents) { dependents.add(dependent) } } -const getPendingDependents = (pending: Pending, atom: AnyAtom) => - pending[0].get(atom) +const getBatchAtomDependents = (batch: Batch, atom: AnyAtom) => + batch.D.get(atom) -const addPendingFunction = (pending: Pending, fn: () => void) => { - pending[2].add(fn) +const copySetAndClear = (origSet: Set): Set => { + const newSet = new Set(origSet) + origSet.clear() + return newSet } -const flushPending = (pending: Pending) => { +const flushBatch = (batch: Batch) => { let error: AnyError let hasError = false const call = (fn: () => void) => { @@ -211,14 +230,10 @@ const flushPending = (pending: Pending) => { } } } - while (pending[1].size || pending[2].size) { - pending[0].clear() - const atomStates = new Set(pending[1].values()) - pending[1].clear() - const functions = new Set(pending[2]) - pending[2].clear() - atomStates.forEach((atomState) => atomState.m?.l.forEach(call)) - functions.forEach(call) + while (batch.M.size || batch.L.size) { + batch.D.clear() + copySetAndClear(batch.M).forEach(call) + copySetAndClear(batch.L).forEach(call) } if (hasError) { throw error @@ -304,7 +319,7 @@ const buildStore = ( } const readAtomState = ( - pending: Pending | undefined, + batch: Batch | undefined, atom: Atom, dirtyAtoms?: Set, ): AtomState => { @@ -324,7 +339,7 @@ const buildStore = ( ([a, n]) => // Recursively, read the atom state of the dependency, and // check if the atom epoch number is unchanged - readAtomState(pending, a, dirtyAtoms).n === n, + readAtomState(batch, a, dirtyAtoms).n === n, ) ) { return atomState @@ -347,17 +362,17 @@ const buildStore = ( return returnAtomValue(aState) } // a !== atom - const aState = readAtomState(pending, a, dirtyAtoms) + const aState = readAtomState(batch, a, dirtyAtoms) try { return returnAtomValue(aState) } finally { if (isSync) { - addDependency(pending, atom, atomState, a, aState) + addDependency(batch, atom, atomState, a, aState) } else { - const pending = createPending() - addDependency(pending, atom, atomState, a, aState) - mountDependencies(pending, atom, atomState) - flushPending(pending) + const batch = createBatch() + addDependency(batch, atom, atomState, a, aState) + mountDependencies(batch, atom, atomState) + flushBatch(batch) } } } @@ -397,9 +412,9 @@ const buildStore = ( valueOrPromise.onCancel?.(() => controller?.abort()) const complete = () => { if (atomState.m) { - const pending = createPending() - mountDependencies(pending, atom, atomState) - flushPending(pending) + const batch = createBatch() + mountDependencies(batch, atom, atomState) + flushBatch(batch) } } valueOrPromise.then(complete, complete) @@ -418,8 +433,8 @@ const buildStore = ( const readAtom = (atom: Atom): Value => returnAtomValue(readAtomState(undefined, atom)) - const getMountedOrPendingDependents = ( - pending: Pending, + const getMountedOrBatchDependents = ( + batch: Batch, atom: Atom, atomState: AtomState, ): Map => { @@ -436,7 +451,7 @@ const buildStore = ( getAtomState(atomWithPendingPromise), ) } - getPendingDependents(pending, atom)?.forEach((dependent) => { + getBatchAtomDependents(batch, atom)?.forEach((dependent) => { dependents.set(dependent, getAtomState(dependent)) }) return dependents @@ -446,7 +461,7 @@ const buildStore = ( // what's described here for simplicity and performance reasons: // https://en.wikipedia.org/wiki/Topological_sorting#Depth-first_search function getSortedDependents( - pending: Pending, + batch: Batch, rootAtom: AnyAtom, rootAtomState: AtomState, ): [[AnyAtom, AtomState, number][], Set] { @@ -476,7 +491,7 @@ const buildStore = ( } visiting.add(a) // Push unvisited dependents onto the stack - for (const [d, s] of getMountedOrPendingDependents(pending, a, aState)) { + for (const [d, s] of getMountedOrBatchDependents(batch, a, aState)) { if (a !== d && !visiting.has(d)) { stack.push([d, s]) } @@ -486,14 +501,14 @@ const buildStore = ( } const recomputeDependents = ( - pending: Pending, + 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( - pending, + batch, atom, atomState, ) @@ -511,10 +526,10 @@ const buildStore = ( } } if (hasChangedDeps) { - readAtomState(pending, a, markedAtoms) - mountDependencies(pending, a, aState) + readAtomState(batch, a, markedAtoms) + mountDependencies(batch, a, aState) if (prevEpochNumber !== aState.n) { - addPendingAtom(pending, a, aState) + registerBatchAtom(batch, a, aState) changedAtoms.add(a) } } @@ -523,13 +538,13 @@ const buildStore = ( } const writeAtomState = ( - pending: Pending, + batch: Batch, atom: WritableAtom, ...args: Args ): Result => { let isSync = true const getter: Getter = (a: Atom) => - returnAtomValue(readAtomState(pending, a)) + returnAtomValue(readAtomState(batch, a)) const setter: Setter = ( a: WritableAtom, ...args: As @@ -544,18 +559,18 @@ const buildStore = ( const prevEpochNumber = aState.n const v = args[0] as V setAtomStateValueOrPromise(a, aState, v) - mountDependencies(pending, a, aState) + mountDependencies(batch, a, aState) if (prevEpochNumber !== aState.n) { - addPendingAtom(pending, a, aState) - recomputeDependents(pending, a, aState) + registerBatchAtom(batch, a, aState) + recomputeDependents(batch, a, aState) } return undefined as R } else { - return writeAtomState(pending, a, ...args) + return writeAtomState(batch, a, ...args) } } finally { if (!isSync) { - flushPending(pending) + flushBatch(batch) } } } @@ -570,23 +585,23 @@ const buildStore = ( atom: WritableAtom, ...args: Args ): Result => { - const pending = createPending() + const batch = createBatch() try { - return writeAtomState(pending, atom, ...args) + return writeAtomState(batch, atom, ...args) } finally { - flushPending(pending) + flushBatch(batch) } } const mountDependencies = ( - pending: Pending, + batch: Batch, atom: AnyAtom, atomState: AtomState, ) => { if (atomState.m && !isPendingPromise(atomState.v)) { for (const a of atomState.d.keys()) { if (!atomState.m.d.has(a)) { - const aMounted = mountAtom(pending, a, getAtomState(a)) + const aMounted = mountAtom(batch, a, getAtomState(a)) aMounted.t.add(atom) atomState.m.d.add(a) } @@ -594,7 +609,7 @@ const buildStore = ( for (const a of atomState.m.d || []) { if (!atomState.d.has(a)) { atomState.m.d.delete(a) - const aMounted = unmountAtom(pending, a, getAtomState(a)) + const aMounted = unmountAtom(batch, a, getAtomState(a)) aMounted?.t.delete(atom) } } @@ -602,16 +617,16 @@ const buildStore = ( } const mountAtom = ( - pending: Pending, + batch: Batch, atom: Atom, atomState: AtomState, ): Mounted => { if (!atomState.m) { // recompute atom state - readAtomState(pending, atom) + readAtomState(batch, atom) // mount dependencies first for (const a of atomState.d.keys()) { - const aMounted = mountAtom(pending, a, getAtomState(a)) + const aMounted = mountAtom(batch, a, getAtomState(a)) aMounted.t.add(atom) } // mount self @@ -626,14 +641,14 @@ const buildStore = ( if (isActuallyWritableAtom(atom)) { const mounted = atomState.m let setAtom: (...args: unknown[]) => unknown - const createInvocationContext = (pending: Pending, fn: () => T) => { + const createInvocationContext = (batch: Batch, fn: () => T) => { let isSync = true setAtom = (...args: unknown[]) => { try { - return writeAtomState(pending, atom, ...args) + return writeAtomState(batch, atom, ...args) } finally { if (!isSync) { - flushPending(pending) + flushBatch(batch) } } } @@ -643,12 +658,12 @@ const buildStore = ( isSync = false } } - addPendingFunction(pending, () => { - const onUnmount = createInvocationContext(pending, () => + addBatchFuncLow(batch, () => { + const onUnmount = createInvocationContext(batch, () => atomOnMount(atom, (...args) => setAtom(...args)), ) if (onUnmount) { - mounted.u = (pending) => createInvocationContext(pending, onUnmount) + mounted.u = (batch) => createInvocationContext(batch, onUnmount) } }) } @@ -657,7 +672,7 @@ const buildStore = ( } const unmountAtom = ( - pending: Pending, + batch: Batch, atom: Atom, atomState: AtomState, ): Mounted | undefined => { @@ -669,7 +684,7 @@ const buildStore = ( // unmount self const onUnmount = atomState.m.u if (onUnmount) { - addPendingFunction(pending, () => onUnmount(pending)) + addBatchFuncLow(batch, () => onUnmount(batch)) } delete atomState.m if (import.meta.env?.MODE !== 'production') { @@ -677,7 +692,7 @@ const buildStore = ( } // unmount dependencies for (const a of atomState.d.keys()) { - const aMounted = unmountAtom(pending, a, getAtomState(a)) + const aMounted = unmountAtom(batch, a, getAtomState(a)) aMounted?.t.delete(atom) } return undefined @@ -686,17 +701,17 @@ const buildStore = ( } const subscribeAtom = (atom: AnyAtom, listener: () => void) => { - const pending = createPending() + const batch = createBatch() const atomState = getAtomState(atom) - const mounted = mountAtom(pending, atom, atomState) + const mounted = mountAtom(batch, atom, atomState) const listeners = mounted.l listeners.add(listener) - flushPending(pending) + flushBatch(batch) return () => { listeners.delete(listener) - const pending = createPending() - unmountAtom(pending, atom, atomState) - flushPending(pending) + const batch = createBatch() + unmountAtom(batch, atom, atomState) + flushBatch(batch) } } @@ -724,20 +739,20 @@ const buildStore = ( }), dev4_get_mounted_atoms: () => debugMountedAtoms, dev4_restore_atoms: (values) => { - const pending = createPending() + const batch = createBatch() for (const [atom, value] of values) { if (hasInitialValue(atom)) { const atomState = getAtomState(atom) const prevEpochNumber = atomState.n setAtomStateValueOrPromise(atom, atomState, value) - mountDependencies(pending, atom, atomState) + mountDependencies(batch, atom, atomState) if (prevEpochNumber !== atomState.n) { - addPendingAtom(pending, atom, atomState) - recomputeDependents(pending, atom, atomState) + registerBatchAtom(batch, atom, atomState) + recomputeDependents(batch, atom, atomState) } } } - flushPending(pending) + flushBatch(batch) }, } Object.assign(store, devStore) diff --git a/tests/vanilla/store.test.tsx b/tests/vanilla/store.test.tsx index 2d978ad545..0677aa092a 100644 --- a/tests/vanilla/store.test.tsx +++ b/tests/vanilla/store.test.tsx @@ -628,7 +628,7 @@ describe('should invoke flushPending only after all atoms are updated (#2804)', store.sub(a, () => { mountResult.push('a value changed - ' + store.get(a)) }) - const unsub = store.sub(m, () => {}) + store.sub(m, () => {}) mountResult.push('after store.sub') expect(mountResult).not.toEqual([ 'before store.sub', @@ -736,7 +736,7 @@ describe('should mount and trigger listeners even when an error is thrown', () = set(a, 1) get(e) }) - const w = atom(null, async (get, set) => { + const w = atom(null, async (_get, set) => { setTimeout(() => { try { set(b)