diff --git a/src/vanilla/store.ts b/src/vanilla/store.ts index af064d367b..fc31b0d55f 100644 --- a/src/vanilla/store.ts +++ b/src/vanilla/store.ts @@ -148,6 +148,7 @@ type PrdStore = { ) => Result sub: (atom: AnyAtom, listener: () => void) => () => void unstable_resolve?: >(atom: A) => A + unstable_derive: () => Store } type Store = PrdStore & Partial @@ -181,668 +182,677 @@ export const createStore = (): Store => { devListenersRev2 = new Set() mountedAtoms = new Set() } + const buildStore = (): Store => { + const resolveAtom = >(atom: A): A => { + return store.unstable_resolve?.(atom) ?? atom + } - const resolveAtom = >(atom: A): A => { - return store.unstable_resolve?.(atom) ?? atom - } - - const getAtomState = (atom: Atom) => - atomStateMap.get(atom) as AtomState | undefined + const getAtomState = (atom: Atom) => + atomStateMap.get(atom) as AtomState | undefined - const addPendingDependent = (atom: AnyAtom, atomState: AtomState) => { - atomState.d.forEach((_, a) => { - if (!pendingMap.has(a)) { - const aState = getAtomState(a) - pendingMap.set(a, [aState, new Set()]) - if (aState) { - addPendingDependent(a, aState) + const addPendingDependent = (atom: AnyAtom, atomState: AtomState) => { + atomState.d.forEach((_, a) => { + if (!pendingMap.has(a)) { + const aState = getAtomState(a) + pendingMap.set(a, [aState, new Set()]) + if (aState) { + addPendingDependent(a, aState) + } } - } - pendingMap.get(a)![1].add(atom) - }) - } - - const setAtomState = ( - atom: Atom, - atomState: AtomState, - ): void => { - if (import.meta.env?.MODE !== 'production') { - Object.freeze(atomState) - } - const prevAtomState = getAtomState(atom) - atomStateMap.set(atom, atomState) - pendingStack[pendingStack.length - 1]?.add(atom) - if (!pendingMap.has(atom)) { - pendingMap.set(atom, [prevAtomState, new Set()]) - addPendingDependent(atom, atomState) - } - if (hasPromiseAtomValue(prevAtomState)) { - const next = - 'v' in atomState - ? atomState.v instanceof Promise - ? atomState.v - : Promise.resolve(atomState.v) - : Promise.reject(atomState.e) - if (prevAtomState.v !== next) { - cancelPromise(prevAtomState.v, next) - } + pendingMap.get(a)![1].add(atom) + }) } - } - const updateDependencies = ( - atom: Atom, - nextAtomState: AtomState, - nextDependencies: NextDependencies, - keepPreviousDependencies?: boolean, - ): void => { - const dependencies: Dependencies = new Map( - keepPreviousDependencies ? nextAtomState.d : null, - ) - let changed = false - nextDependencies.forEach((aState, a) => { - if (!aState && a === atom) { - aState = nextAtomState - } - if (aState) { - dependencies.set(a, aState) - if (nextAtomState.d.get(a) !== aState) { - changed = true + const setAtomState = ( + atom: Atom, + atomState: AtomState, + ): void => { + if (import.meta.env?.MODE !== 'production') { + Object.freeze(atomState) + } + const prevAtomState = getAtomState(atom) + atomStateMap.set(atom, atomState) + pendingStack[pendingStack.length - 1]?.add(atom) + if (!pendingMap.has(atom)) { + pendingMap.set(atom, [prevAtomState, new Set()]) + addPendingDependent(atom, atomState) + } + if (hasPromiseAtomValue(prevAtomState)) { + const next = + 'v' in atomState + ? atomState.v instanceof Promise + ? atomState.v + : Promise.resolve(atomState.v) + : Promise.reject(atomState.e) + if (prevAtomState.v !== next) { + cancelPromise(prevAtomState.v, next) } - } else if (import.meta.env?.MODE !== 'production') { - console.warn('[Bug] atom state not found') } - }) - if (changed || nextAtomState.d.size !== dependencies.size) { - nextAtomState.d = dependencies } - } - const setAtomValue = ( - atom: Atom, - value: Value, - nextDependencies?: NextDependencies, - keepPreviousDependencies?: boolean, - ): AtomState => { - const prevAtomState = getAtomState(atom) - const nextAtomState: AtomState = { - d: prevAtomState?.d || new Map(), - v: value, - } - if (nextDependencies) { - updateDependencies( - atom, - nextAtomState, - nextDependencies, - keepPreviousDependencies, + const updateDependencies = ( + atom: Atom, + nextAtomState: AtomState, + nextDependencies: NextDependencies, + keepPreviousDependencies?: boolean, + ): void => { + const dependencies: Dependencies = new Map( + keepPreviousDependencies ? nextAtomState.d : null, ) - } - if ( - isEqualAtomValue(prevAtomState, nextAtomState) && - prevAtomState.d === nextAtomState.d - ) { - // bail out - return prevAtomState - } - if ( - hasPromiseAtomValue(prevAtomState) && - hasPromiseAtomValue(nextAtomState) && - isEqualPromiseAtomValue(prevAtomState, nextAtomState) - ) { - if (prevAtomState.d === nextAtomState.d) { - // bail out - return prevAtomState - } else { - // restore the wrapped promise - nextAtomState.v = prevAtomState.v + let changed = false + nextDependencies.forEach((aState, a) => { + if (!aState && a === atom) { + aState = nextAtomState + } + if (aState) { + dependencies.set(a, aState) + if (nextAtomState.d.get(a) !== aState) { + changed = true + } + } else if (import.meta.env?.MODE !== 'production') { + console.warn('[Bug] atom state not found') + } + }) + if (changed || nextAtomState.d.size !== dependencies.size) { + nextAtomState.d = dependencies } } - setAtomState(atom, nextAtomState) - return nextAtomState - } - const setAtomValueOrPromise = ( - atom: Atom, - valueOrPromise: Value, - nextDependencies?: NextDependencies, - abortPromise?: () => void, - ): AtomState => { - if (isPromiseLike(valueOrPromise)) { - let continuePromise: (next: Promise>) => void - const updatePromiseDependencies = () => { - const prevAtomState = getAtomState(atom) - if ( - !hasPromiseAtomValue(prevAtomState) || - prevAtomState.v !== promise - ) { - // not the latest promise - return - } - // update dependencies, that could have changed - const nextAtomState = setAtomValue( + const setAtomValue = ( + atom: Atom, + value: Value, + nextDependencies?: NextDependencies, + keepPreviousDependencies?: boolean, + ): AtomState => { + const prevAtomState = getAtomState(atom) + const nextAtomState: AtomState = { + d: prevAtomState?.d || new Map(), + v: value, + } + if (nextDependencies) { + updateDependencies( atom, - promise as Value, + nextAtomState, nextDependencies, + keepPreviousDependencies, ) - if (mountedMap.has(atom) && prevAtomState.d !== nextAtomState.d) { - mountDependencies(atom, nextAtomState, prevAtomState.d) + } + if ( + isEqualAtomValue(prevAtomState, nextAtomState) && + prevAtomState.d === nextAtomState.d + ) { + // bail out + return prevAtomState + } + if ( + hasPromiseAtomValue(prevAtomState) && + hasPromiseAtomValue(nextAtomState) && + isEqualPromiseAtomValue(prevAtomState, nextAtomState) + ) { + if (prevAtomState.d === nextAtomState.d) { + // bail out + return prevAtomState + } else { + // restore the wrapped promise + nextAtomState.v = prevAtomState.v } } - const promise: Promise> & PromiseMeta> = - new Promise((resolve, reject) => { - let settled = false - valueOrPromise.then( - (v) => { - if (!settled) { - settled = true - resolvePromise(promise, v) - resolve(v as Awaited) - updatePromiseDependencies() - } - }, - (e) => { + setAtomState(atom, nextAtomState) + return nextAtomState + } + + const setAtomValueOrPromise = ( + atom: Atom, + valueOrPromise: Value, + nextDependencies?: NextDependencies, + abortPromise?: () => void, + ): AtomState => { + if (isPromiseLike(valueOrPromise)) { + let continuePromise: (next: Promise>) => void + const updatePromiseDependencies = () => { + const prevAtomState = getAtomState(atom) + if ( + !hasPromiseAtomValue(prevAtomState) || + prevAtomState.v !== promise + ) { + // not the latest promise + return + } + // update dependencies, that could have changed + const nextAtomState = setAtomValue( + atom, + promise as Value, + nextDependencies, + ) + if (mountedMap.has(atom) && prevAtomState.d !== nextAtomState.d) { + mountDependencies(atom, nextAtomState, prevAtomState.d) + } + } + const promise: Promise> & PromiseMeta> = + new Promise((resolve, reject) => { + let settled = false + valueOrPromise.then( + (v) => { + if (!settled) { + settled = true + resolvePromise(promise, v) + resolve(v as Awaited) + updatePromiseDependencies() + } + }, + (e) => { + if (!settled) { + settled = true + rejectPromise(promise, e) + reject(e) + updatePromiseDependencies() + } + }, + ) + continuePromise = (next) => { if (!settled) { settled = true - rejectPromise(promise, e) - reject(e) - updatePromiseDependencies() + next.then( + (v) => resolvePromise(promise, v), + (e) => rejectPromise(promise, e), + ) + resolve(next) } - }, - ) - continuePromise = (next) => { - if (!settled) { - settled = true - next.then( - (v) => resolvePromise(promise, v), - (e) => rejectPromise(promise, e), - ) - resolve(next) } + }) + promise.orig = valueOrPromise as PromiseLike> + promise.status = 'pending' + registerCancelPromise(promise, (next) => { + if (next) { + continuePromise(next as Promise>) } + abortPromise?.() }) - promise.orig = valueOrPromise as PromiseLike> - promise.status = 'pending' - registerCancelPromise(promise, (next) => { - if (next) { - continuePromise(next as Promise>) - } - abortPromise?.() - }) - return setAtomValue(atom, promise as Value, nextDependencies, true) - } - return setAtomValue(atom, valueOrPromise, nextDependencies) - } - - const setAtomError = ( - atom: Atom, - error: AnyError, - nextDependencies?: NextDependencies, - ): AtomState => { - const prevAtomState = getAtomState(atom) - const nextAtomState: AtomState = { - d: prevAtomState?.d || new Map(), - e: error, - } - if (nextDependencies) { - updateDependencies(atom, nextAtomState, nextDependencies) - } - if ( - isEqualAtomError(prevAtomState, nextAtomState) && - prevAtomState.d === nextAtomState.d - ) { - // bail out - return prevAtomState - } - setAtomState(atom, nextAtomState) - return nextAtomState - } - - const readAtomState = ( - atom: Atom, - force?: (a: AnyAtom) => boolean, - ): AtomState => { - // See if we can skip recomputing this atom. - const atomState = getAtomState(atom) - if (!force?.(atom) && atomState) { - // If the atom is mounted, we can use the cache. - // because it should have been updated by dependencies. - if (mountedMap.has(atom)) { - return atomState - } - // Otherwise, check if the dependencies have changed. - // If all dependencies haven't changed, we can use the cache. + return setAtomValue(atom, promise as Value, nextDependencies, true) + } + return setAtomValue(atom, valueOrPromise, nextDependencies) + } + + const setAtomError = ( + atom: Atom, + error: AnyError, + nextDependencies?: NextDependencies, + ): AtomState => { + const prevAtomState = getAtomState(atom) + const nextAtomState: AtomState = { + d: prevAtomState?.d || new Map(), + e: error, + } + if (nextDependencies) { + updateDependencies(atom, nextAtomState, nextDependencies) + } if ( - Array.from(atomState.d).every(([a, s]) => { - if (a === atom) { - return true - } - const aState = readAtomState(a, force) - // Check if the atom state is unchanged, or - // check the atom value in case only dependencies are changed - return aState === s || isEqualAtomValue(aState, s) - }) + isEqualAtomError(prevAtomState, nextAtomState) && + prevAtomState.d === nextAtomState.d ) { - return atomState - } - } - // Compute a new state for this atom. - const nextDependencies: NextDependencies = new Map() - let isSync = true - const getter: Getter = (a: Atom) => { - const resolvedA = resolveAtom(a) - if (resolvedA === (atom as AnyAtom)) { - const aState = getAtomState(resolvedA) - if (aState) { - nextDependencies.set(resolvedA, aState) - return returnAtomValue(aState) - } - if (hasInitialValue(resolvedA)) { - nextDependencies.set(resolvedA, undefined) - return resolvedA.init - } - // NOTE invalid derived atoms can reach here - throw new Error('no atom init') + // bail out + return prevAtomState } - // resolvedA !== atom - const aState = readAtomState(resolvedA, force) - nextDependencies.set(resolvedA, aState) - return returnAtomValue(aState) + setAtomState(atom, nextAtomState) + return nextAtomState } - let controller: AbortController | undefined - let setSelf: ((...args: unknown[]) => unknown) | undefined - const options = { - get signal() { - if (!controller) { - controller = new AbortController() + + const readAtomState = ( + atom: Atom, + force?: (a: AnyAtom) => boolean, + ): AtomState => { + // See if we can skip recomputing this atom. + const atomState = getAtomState(atom) + if (!force?.(atom) && atomState) { + // If the atom is mounted, we can use the cache. + // because it should have been updated by dependencies. + if (mountedMap.has(atom)) { + return atomState } - return controller.signal - }, - get setSelf() { + // Otherwise, check if the dependencies have changed. + // If all dependencies haven't changed, we can use the cache. if ( - import.meta.env?.MODE !== 'production' && - !isActuallyWritableAtom(atom) + Array.from(atomState.d).every(([a, s]) => { + if (a === atom) { + return true + } + const aState = readAtomState(a, force) + // Check if the atom state is unchanged, or + // check the atom value in case only dependencies are changed + return aState === s || isEqualAtomValue(aState, s) + }) ) { - console.warn('setSelf function cannot be used with read-only atom') + return atomState } - if (!setSelf && isActuallyWritableAtom(atom)) { - setSelf = (...args) => { - if (import.meta.env?.MODE !== 'production' && isSync) { - console.warn('setSelf function cannot be called in sync') - } - if (!isSync) { - return writeAtom(atom, ...args) - } + } + // Compute a new state for this atom. + const nextDependencies: NextDependencies = new Map() + let isSync = true + const getter: Getter = (a: Atom) => { + const resolvedA = resolveAtom(a) + if (resolvedA === (atom as AnyAtom)) { + const aState = getAtomState(resolvedA) + if (aState) { + nextDependencies.set(resolvedA, aState) + return returnAtomValue(aState) + } + if (hasInitialValue(resolvedA)) { + nextDependencies.set(resolvedA, undefined) + return resolvedA.init } + // NOTE invalid derived atoms can reach here + throw new Error('no atom init') } - return setSelf - }, - } - try { - const valueOrPromise = atom.read(getter, options as never) - return setAtomValueOrPromise(atom, valueOrPromise, nextDependencies, () => - controller?.abort(), - ) - } catch (error) { - return setAtomError(atom, error, nextDependencies) - } finally { - isSync = false + // resolvedA !== atom + const aState = readAtomState(resolvedA, force) + nextDependencies.set(resolvedA, aState) + return returnAtomValue(aState) + } + let controller: AbortController | undefined + let setSelf: ((...args: unknown[]) => unknown) | undefined + const options = { + get signal() { + if (!controller) { + controller = new AbortController() + } + return controller.signal + }, + get setSelf() { + if ( + import.meta.env?.MODE !== 'production' && + !isActuallyWritableAtom(atom) + ) { + console.warn('setSelf function cannot be used with read-only atom') + } + if (!setSelf && isActuallyWritableAtom(atom)) { + setSelf = (...args) => { + if (import.meta.env?.MODE !== 'production' && isSync) { + console.warn('setSelf function cannot be called in sync') + } + if (!isSync) { + return writeAtom(atom, ...args) + } + } + } + return setSelf + }, + } + try { + const valueOrPromise = atom.read(getter, options as never) + return setAtomValueOrPromise( + atom, + valueOrPromise, + nextDependencies, + () => controller?.abort(), + ) + } catch (error) { + return setAtomError(atom, error, nextDependencies) + } finally { + isSync = false + } } - } - const readAtom = (atom: Atom): Value => - returnAtomValue(readAtomState(resolveAtom(atom))) + const readAtom = (atom: Atom): Value => + returnAtomValue(readAtomState(resolveAtom(atom))) - const recomputeDependents = (atom: AnyAtom): void => { - const getDependents = (a: AnyAtom): Dependents => { - const dependents = new Set(mountedMap.get(a)?.t) - pendingMap.get(a)?.[1].forEach((dependent) => { - dependents.add(dependent) - }) - return dependents - } + const recomputeDependents = (atom: AnyAtom): void => { + const getDependents = (a: AnyAtom): Dependents => { + const dependents = new Set(mountedMap.get(a)?.t) + pendingMap.get(a)?.[1].forEach((dependent) => { + dependents.add(dependent) + }) + 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 + // 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 = new Array() - const markedAtoms = new Set() - const visit = (n: AnyAtom) => { - if (markedAtoms.has(n)) { - return - } - markedAtoms.add(n) - for (const m of getDependents(n)) { - if (n !== m) { - visit(m) + // 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 = new Array() + const markedAtoms = new Set() + const visit = (n: AnyAtom) => { + if (markedAtoms.has(n)) { + return } + markedAtoms.add(n) + for (const m of getDependents(n)) { + if (n !== m) { + visit(m) + } + } + // 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(n) } - // 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(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) - - // 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]) - const isMarked = (a: AnyAtom) => markedAtoms.has(a) - for (let i = topsortedAtoms.length - 1; i >= 0; --i) { - const a = topsortedAtoms[i]! - const prevAtomState = getAtomState(a) - if (!prevAtomState) { - continue - } - let hasChangedDeps = false - for (const dep of prevAtomState.d.keys()) { - if (dep !== a && changedAtoms.has(dep)) { - hasChangedDeps = true - break + // 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) + + // 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]) + const isMarked = (a: AnyAtom) => markedAtoms.has(a) + for (let i = topsortedAtoms.length - 1; i >= 0; --i) { + const a = topsortedAtoms[i]! + const prevAtomState = getAtomState(a) + if (!prevAtomState) { + continue } - } - if (hasChangedDeps) { - const nextAtomState = readAtomState(a, isMarked) - addPendingDependent(a, nextAtomState) - if (!isEqualAtomValue(prevAtomState, nextAtomState)) { - changedAtoms.add(a) + let hasChangedDeps = false + for (const dep of prevAtomState.d.keys()) { + if (dep !== a && changedAtoms.has(dep)) { + hasChangedDeps = true + break + } + } + if (hasChangedDeps) { + const nextAtomState = readAtomState(a, isMarked) + addPendingDependent(a, nextAtomState) + if (!isEqualAtomValue(prevAtomState, nextAtomState)) { + changedAtoms.add(a) + } } + markedAtoms.delete(a) } - markedAtoms.delete(a) } - } - const writeAtomState = ( - atom: WritableAtom, - ...args: Args - ): Result => { - const getter: Getter = (a: Atom) => - returnAtomValue(readAtomState(resolveAtom(a))) - const setter: Setter = ( - a: WritableAtom, - ...args: As - ) => { - const resolvedA = resolveAtom(a) - const isSync = pendingStack.length > 0 - if (!isSync) { - pendingStack.push(new Set([resolvedA])) - } - let r: R | undefined - if (resolvedA === (atom as AnyAtom)) { - if (!hasInitialValue(resolvedA)) { - // NOTE technically possible but restricted as it may cause bugs - throw new Error('atom not writable') + const writeAtomState = ( + atom: WritableAtom, + ...args: Args + ): Result => { + const getter: Getter = (a: Atom) => + returnAtomValue(readAtomState(resolveAtom(a))) + const setter: Setter = ( + a: WritableAtom, + ...args: As + ) => { + const resolvedA = resolveAtom(a) + const isSync = pendingStack.length > 0 + if (!isSync) { + pendingStack.push(new Set([resolvedA])) } - const prevAtomState = getAtomState(resolvedA) - const nextAtomState = setAtomValueOrPromise(resolvedA, args[0] as V) - if (!isEqualAtomValue(prevAtomState, nextAtomState)) { - recomputeDependents(resolvedA) + let r: R | undefined + if (resolvedA === (atom as AnyAtom)) { + if (!hasInitialValue(resolvedA)) { + // NOTE technically possible but restricted as it may cause bugs + throw new Error('atom not writable') + } + const prevAtomState = getAtomState(resolvedA) + const nextAtomState = setAtomValueOrPromise(resolvedA, args[0] as V) + if (!isEqualAtomValue(prevAtomState, nextAtomState)) { + recomputeDependents(resolvedA) + } + } else { + r = writeAtomState(resolvedA as AnyWritableAtom, ...args) as R } - } else { - r = writeAtomState(resolvedA as AnyWritableAtom, ...args) as R - } - if (!isSync) { - const flushed = flushPending(pendingStack.pop()!) - if (import.meta.env?.MODE !== 'production') { - devListenersRev2.forEach((l) => - l({ type: 'async-write', flushed: flushed! }), - ) + if (!isSync) { + const flushed = flushPending(pendingStack.pop()!) + if (import.meta.env?.MODE !== 'production') { + devListenersRev2.forEach((l) => + l({ type: 'async-write', flushed: flushed! }), + ) + } } + return r as R } - return r as R + const result = atom.write(getter, setter, ...args) + return result } - const result = atom.write(getter, setter, ...args) - return result - } - const writeAtom = ( - atom: WritableAtom, - ...args: Args - ): Result => { - const resolvedAtom = resolveAtom(atom) - pendingStack.push(new Set([resolvedAtom])) - const result = writeAtomState(resolvedAtom, ...args) + const writeAtom = ( + atom: WritableAtom, + ...args: Args + ): Result => { + const resolvedAtom = resolveAtom(atom) + pendingStack.push(new Set([resolvedAtom])) + const result = writeAtomState(resolvedAtom, ...args) - const flushed = flushPending(pendingStack.pop()!) - if (import.meta.env?.MODE !== 'production') { - devListenersRev2.forEach((l) => l({ type: 'write', flushed: flushed! })) - } - return result - } - - const mountAtom = ( - atom: Atom, - initialDependent?: AnyAtom, - onMountQueue?: (() => void)[], - ): Mounted => { - const existingMount = mountedMap.get(atom) - if (existingMount) { - if (initialDependent) { - existingMount.t.add(initialDependent) - } - return existingMount - } + const flushed = flushPending(pendingStack.pop()!) + if (import.meta.env?.MODE !== 'production') { + devListenersRev2.forEach((l) => l({ type: 'write', flushed: flushed! })) + } + return result + } + + const mountAtom = ( + atom: Atom, + initialDependent?: AnyAtom, + onMountQueue?: (() => void)[], + ): Mounted => { + const existingMount = mountedMap.get(atom) + if (existingMount) { + if (initialDependent) { + existingMount.t.add(initialDependent) + } + return existingMount + } - const queue = onMountQueue || [] - // mount dependencies before mounting self - getAtomState(atom)?.d.forEach((_, a) => { - if (a !== atom) { - mountAtom(a, atom, queue) - } - }) - // recompute atom state - readAtomState(atom) - // mount self - const mounted: Mounted = { - t: new Set(initialDependent && [initialDependent]), - l: new Set(), - } - mountedMap.set(atom, mounted) - if (import.meta.env?.MODE !== 'production') { - mountedAtoms.add(atom) - } - // onMount - if (isActuallyWritableAtom(atom) && atom.onMount) { - const { onMount } = atom - queue.push(() => { - const onUnmount = onMount((...args) => writeAtom(atom, ...args)) - if (onUnmount) { - mounted.u = onUnmount + const queue = onMountQueue || [] + // mount dependencies before mounting self + getAtomState(atom)?.d.forEach((_, a) => { + if (a !== atom) { + mountAtom(a, atom, queue) } }) + // recompute atom state + readAtomState(atom) + // mount self + const mounted: Mounted = { + t: new Set(initialDependent && [initialDependent]), + l: new Set(), + } + mountedMap.set(atom, mounted) + if (import.meta.env?.MODE !== 'production') { + mountedAtoms.add(atom) + } + // onMount + if (isActuallyWritableAtom(atom) && atom.onMount) { + const { onMount } = atom + queue.push(() => { + const onUnmount = onMount((...args) => writeAtom(atom, ...args)) + if (onUnmount) { + mounted.u = onUnmount + } + }) + } + if (!onMountQueue) { + queue.forEach((f) => f()) + } + return mounted } - if (!onMountQueue) { - queue.forEach((f) => f()) - } - return mounted - } - // FIXME doesn't work with mutually dependent atoms - const canUnmountAtom = (atom: AnyAtom, mounted: Mounted) => - !mounted.l.size && - (!mounted.t.size || (mounted.t.size === 1 && mounted.t.has(atom))) + // FIXME doesn't work with mutually dependent atoms + const canUnmountAtom = (atom: AnyAtom, mounted: Mounted) => + !mounted.l.size && + (!mounted.t.size || (mounted.t.size === 1 && mounted.t.has(atom))) - const tryUnmountAtom = (atom: Atom, mounted: Mounted): void => { - if (!canUnmountAtom(atom, mounted)) { - return - } - // unmount self - const onUnmount = mounted.u - if (onUnmount) { - onUnmount() - } - mountedMap.delete(atom) - if (import.meta.env?.MODE !== 'production') { - mountedAtoms.delete(atom) - } - // unmount dependencies afterward - const atomState = getAtomState(atom) - if (atomState) { - // cancel promise - if (hasPromiseAtomValue(atomState)) { - cancelPromise(atomState.v) + const tryUnmountAtom = ( + atom: Atom, + mounted: Mounted, + ): void => { + if (!canUnmountAtom(atom, mounted)) { + return } - atomState.d.forEach((_, a) => { - if (a !== atom) { - const mountedDep = mountedMap.get(a) - if (mountedDep) { - mountedDep.t.delete(atom) - tryUnmountAtom(a, mountedDep) + // unmount self + const onUnmount = mounted.u + if (onUnmount) { + onUnmount() + } + mountedMap.delete(atom) + if (import.meta.env?.MODE !== 'production') { + mountedAtoms.delete(atom) + } + // unmount dependencies afterward + const atomState = getAtomState(atom) + if (atomState) { + // cancel promise + if (hasPromiseAtomValue(atomState)) { + cancelPromise(atomState.v) + } + atomState.d.forEach((_, a) => { + if (a !== atom) { + const mountedDep = mountedMap.get(a) + if (mountedDep) { + mountedDep.t.delete(atom) + tryUnmountAtom(a, mountedDep) + } } + }) + } else if (import.meta.env?.MODE !== 'production') { + console.warn('[Bug] could not find atom state to unmount', atom) + } + } + + const mountDependencies = ( + atom: Atom, + atomState: AtomState, + prevDependencies?: Dependencies, + ): void => { + const depSet = new Set(atomState.d.keys()) + const maybeUnmountAtomSet = new Set() + prevDependencies?.forEach((_, a) => { + if (depSet.has(a)) { + // not changed + depSet.delete(a) + return + } + maybeUnmountAtomSet.add(a) + const mounted = mountedMap.get(a) + if (mounted) { + mounted.t.delete(atom) // delete from dependents + } + }) + depSet.forEach((a) => { + mountAtom(a, atom) + }) + maybeUnmountAtomSet.forEach((a) => { + const mounted = mountedMap.get(a) + if (mounted) { + tryUnmountAtom(a, mounted) } }) - } else if (import.meta.env?.MODE !== 'production') { - console.warn('[Bug] could not find atom state to unmount', atom) } - } - - const mountDependencies = ( - atom: Atom, - atomState: AtomState, - prevDependencies?: Dependencies, - ): void => { - const depSet = new Set(atomState.d.keys()) - const maybeUnmountAtomSet = new Set() - prevDependencies?.forEach((_, a) => { - if (depSet.has(a)) { - // not changed - depSet.delete(a) - return - } - maybeUnmountAtomSet.add(a) - const mounted = mountedMap.get(a) - if (mounted) { - mounted.t.delete(atom) // delete from dependents - } - }) - depSet.forEach((a) => { - mountAtom(a, atom) - }) - maybeUnmountAtomSet.forEach((a) => { - const mounted = mountedMap.get(a) - if (mounted) { - tryUnmountAtom(a, mounted) - } - }) - } - const flushPending = ( - pendingAtoms: AnyAtom[] | Set, - ): void | Set => { - let flushed: Set - if (import.meta.env?.MODE !== 'production') { - flushed = new Set() - } - const pending: [AnyAtom, AtomState | undefined][] = [] - const collectPending = (pendingAtom: AnyAtom) => { - if (!pendingMap.has(pendingAtom)) { - return + const flushPending = ( + pendingAtoms: AnyAtom[] | Set, + ): void | Set => { + let flushed: Set + if (import.meta.env?.MODE !== 'production') { + flushed = new Set() } - const [prevAtomState, dependents] = pendingMap.get(pendingAtom)! - pendingMap.delete(pendingAtom) - pending.push([pendingAtom, prevAtomState]) - dependents.forEach(collectPending) - // FIXME might be better if we can avoid collecting from dependencies - getAtomState(pendingAtom)?.d.forEach((_, a) => collectPending(a)) - } - pendingAtoms.forEach(collectPending) - pending.forEach(([atom, prevAtomState]) => { - const atomState = getAtomState(atom) - if (!atomState) { - if (import.meta.env?.MODE !== 'production') { - console.warn('[Bug] no atom state to flush') + const pending: [AnyAtom, AtomState | undefined][] = [] + const collectPending = (pendingAtom: AnyAtom) => { + if (!pendingMap.has(pendingAtom)) { + return } - return + const [prevAtomState, dependents] = pendingMap.get(pendingAtom)! + pendingMap.delete(pendingAtom) + pending.push([pendingAtom, prevAtomState]) + dependents.forEach(collectPending) + // FIXME might be better if we can avoid collecting from dependencies + getAtomState(pendingAtom)?.d.forEach((_, a) => collectPending(a)) } - if (atomState !== prevAtomState) { - const mounted = mountedMap.get(atom) - if (mounted && atomState.d !== prevAtomState?.d) { - mountDependencies(atom, atomState, prevAtomState?.d) + pendingAtoms.forEach(collectPending) + pending.forEach(([atom, prevAtomState]) => { + const atomState = getAtomState(atom) + if (!atomState) { + if (import.meta.env?.MODE !== 'production') { + console.warn('[Bug] no atom state to flush') + } + return } - if ( - mounted && - !( - // TODO This seems pretty hacky. Hope to fix it. - // Maybe we could `mountDependencies` in `setAtomState`? - ( - !hasPromiseAtomValue(prevAtomState) && - (isEqualAtomValue(prevAtomState, atomState) || - isEqualAtomError(prevAtomState, atomState)) + if (atomState !== prevAtomState) { + const mounted = mountedMap.get(atom) + if (mounted && atomState.d !== prevAtomState?.d) { + mountDependencies(atom, atomState, prevAtomState?.d) + } + if ( + mounted && + !( + // TODO This seems pretty hacky. Hope to fix it. + // Maybe we could `mountDependencies` in `setAtomState`? + ( + !hasPromiseAtomValue(prevAtomState) && + (isEqualAtomValue(prevAtomState, atomState) || + isEqualAtomError(prevAtomState, atomState)) + ) ) - ) - ) { - mounted.l.forEach((listener) => listener()) - if (import.meta.env?.MODE !== 'production') { - flushed.add(atom) + ) { + mounted.l.forEach((listener) => listener()) + if (import.meta.env?.MODE !== 'production') { + flushed.add(atom) + } } } + }) + if (import.meta.env?.MODE !== 'production') { + // @ts-expect-error Variable 'flushed' is used before being assigned. + return flushed } - }) - if (import.meta.env?.MODE !== 'production') { - // @ts-expect-error Variable 'flushed' is used before being assigned. - return flushed } - } - const subscribeAtom = (atom: AnyAtom, listener: () => void) => { - const resolvedAtom = resolveAtom(atom) - const mounted = mountAtom(resolvedAtom) - const flushed = flushPending([resolvedAtom]) - const listeners = mounted.l - listeners.add(listener) - if (import.meta.env?.MODE !== 'production') { - devListenersRev2.forEach((l) => - l({ type: 'sub', flushed: flushed as Set }), - ) - } - return () => { - listeners.delete(listener) - tryUnmountAtom(resolvedAtom, mounted) + const subscribeAtom = (atom: AnyAtom, listener: () => void) => { + const resolvedAtom = resolveAtom(atom) + const mounted = mountAtom(resolvedAtom) + const flushed = flushPending([resolvedAtom]) + const listeners = mounted.l + listeners.add(listener) if (import.meta.env?.MODE !== 'production') { - // devtools uses this to detect if it _can_ unmount or not - devListenersRev2.forEach((l) => l({ type: 'unsub' })) + devListenersRev2.forEach((l) => + l({ type: 'sub', flushed: flushed as Set }), + ) + } + return () => { + listeners.delete(listener) + tryUnmountAtom(resolvedAtom, mounted) + if (import.meta.env?.MODE !== 'production') { + // devtools uses this to detect if it _can_ unmount or not + devListenersRev2.forEach((l) => l({ type: 'unsub' })) + } } } - } - const store: Store = { - get: readAtom, - set: writeAtom, - sub: subscribeAtom, - } - if (import.meta.env?.MODE !== 'production') { - const devStore: DevStoreRev2 = { - // store dev methods (these are tentative and subject to change without notice) - dev_subscribe_store: (l) => { - devListenersRev2.add(l) - return () => { - devListenersRev2.delete(l) - } - }, - dev_get_mounted_atoms: () => mountedAtoms.values(), - dev_get_atom_state: getAtomState, - dev_get_mounted: (a) => mountedMap.get(a), - dev_restore_atoms: (values) => { - pendingStack.push(new Set()) - for (const [atom, valueOrPromise] of values) { - if (hasInitialValue(atom)) { - setAtomValueOrPromise(atom, valueOrPromise) - recomputeDependents(atom) + const store: Store = { + get: readAtom, + set: writeAtom, + sub: subscribeAtom, + unstable_derive: buildStore, + } + if (import.meta.env?.MODE !== 'production') { + const devStore: DevStoreRev2 = { + // store dev methods (these are tentative and subject to change without notice) + dev_subscribe_store: (l) => { + devListenersRev2.add(l) + return () => { + devListenersRev2.delete(l) } - } - const flushed = flushPending(pendingStack.pop()!) - devListenersRev2.forEach((l) => - l({ type: 'restore', flushed: flushed! }), - ) - }, + }, + dev_get_mounted_atoms: () => mountedAtoms.values(), + dev_get_atom_state: getAtomState, + dev_get_mounted: (a) => mountedMap.get(a), + dev_restore_atoms: (values) => { + pendingStack.push(new Set()) + for (const [atom, valueOrPromise] of values) { + if (hasInitialValue(atom)) { + setAtomValueOrPromise(atom, valueOrPromise) + recomputeDependents(atom) + } + } + const flushed = flushPending(pendingStack.pop()!) + devListenersRev2.forEach((l) => + l({ type: 'restore', flushed: flushed! }), + ) + }, + } + Object.assign(store, devStore) } - Object.assign(store, devStore) + return store } - return store as Store + return buildStore() } let defaultStore: Store | undefined diff --git a/src/vanilla/store2.ts b/src/vanilla/store2.ts index e6a687200d..1ba11eb449 100644 --- a/src/vanilla/store2.ts +++ b/src/vanilla/store2.ts @@ -252,6 +252,7 @@ type PrdStore = { ) => Result sub: (atom: AnyAtom, listener: () => void) => () => void unstable_resolve?: >(atom: A) => A + unstable_derive: () => Store } type Store = PrdStore | (PrdStore & DevStoreRev4) @@ -266,458 +267,465 @@ export const createStore = (): Store => { if (import.meta.env?.MODE !== 'production') { debugMountedAtoms = new Set() } - - const getAtomState = (atom: Atom) => { - let atomState = atomStateMap.get(atom) as AtomState | undefined - if (!atomState) { - atomState = { d: new Map(), p: new Set(), n: 0 } - atomStateMap.set(atom, atomState) + const buildStore = (): Store => { + const getAtomState = (atom: Atom) => { + let atomState = atomStateMap.get(atom) as AtomState | undefined + if (!atomState) { + atomState = { d: new Map(), p: new Set(), n: 0 } + atomStateMap.set(atom, atomState) + } + return atomState } - return atomState - } - const resolveAtom = >(atom: A): A => { - return store.unstable_resolve?.(atom) ?? atom - } + const resolveAtom = >(atom: A): A => { + return store.unstable_resolve?.(atom) ?? atom + } - const setAtomStateValueOrPromise = ( - atom: AnyAtom, - atomState: AtomState, - valueOrPromise: unknown, - abortPromise = () => {}, - completePromise = () => {}, - ) => { - const hasPrevValue = 'v' in atomState - const prevValue = atomState.v - const pendingPromise = getPendingContinuablePromise(atomState) - if (isPromiseLike(valueOrPromise)) { - if (pendingPromise) { - if (pendingPromise !== valueOrPromise) { - pendingPromise[CONTINUE_PROMISE](valueOrPromise, abortPromise) - ++atomState.n + const setAtomStateValueOrPromise = ( + atom: AnyAtom, + atomState: AtomState, + valueOrPromise: unknown, + abortPromise = () => {}, + completePromise = () => {}, + ) => { + const hasPrevValue = 'v' in atomState + const prevValue = atomState.v + const pendingPromise = getPendingContinuablePromise(atomState) + if (isPromiseLike(valueOrPromise)) { + if (pendingPromise) { + if (pendingPromise !== valueOrPromise) { + pendingPromise[CONTINUE_PROMISE](valueOrPromise, abortPromise) + ++atomState.n + } + } else { + const continuablePromise = createContinuablePromise( + valueOrPromise, + abortPromise, + completePromise, + ) + if (continuablePromise.status === PENDING) { + for (const a of atomState.d.keys()) { + const aState = getAtomState(a) + addPendingContinuablePromiseToDependency( + atom, + continuablePromise, + aState, + ) + } + } + atomState.v = continuablePromise + delete atomState.e } } else { - const continuablePromise = createContinuablePromise( - valueOrPromise, - abortPromise, - completePromise, - ) - if (continuablePromise.status === PENDING) { - for (const a of atomState.d.keys()) { - const aState = getAtomState(a) - addPendingContinuablePromiseToDependency( - atom, - continuablePromise, - aState, - ) - } + if (pendingPromise) { + pendingPromise[CONTINUE_PROMISE]( + Promise.resolve(valueOrPromise), + abortPromise, + ) } - atomState.v = continuablePromise + atomState.v = valueOrPromise delete atomState.e } - } else { - if (pendingPromise) { - pendingPromise[CONTINUE_PROMISE]( - Promise.resolve(valueOrPromise), - abortPromise, - ) + if (!hasPrevValue || !Object.is(prevValue, atomState.v)) { + ++atomState.n } - atomState.v = valueOrPromise - delete atomState.e } - if (!hasPrevValue || !Object.is(prevValue, atomState.v)) { - ++atomState.n - } - } - const addDependency = ( - pending: Pending | undefined, - atom: Atom, - a: AnyAtom, - aState: AtomState, - ) => { - if (import.meta.env?.MODE !== 'production' && a === atom) { - throw new Error('[Bug] atom cannot depend on itself') - } - const atomState = getAtomState(atom) - atomState.d.set(a, aState.n) - const continuablePromise = getPendingContinuablePromise(atomState) - if (continuablePromise) { - addPendingContinuablePromiseToDependency(atom, continuablePromise, aState) - } - aState.m?.t.add(atom) - if (pending) { - addPendingDependent(pending, a, atom) - } - } - - const readAtomState = ( - pending: Pending | undefined, - atom: Atom, - force?: (a: AnyAtom) => boolean, - ): AtomState => { - // See if we can skip recomputing this atom. - const atomState = getAtomState(atom) - if (!force?.(atom) && isAtomStateInitialized(atomState)) { - // If the atom is mounted, we can use the cache. - // because it should have been updated by dependencies. - if (atomState.m) { - return atomState + const addDependency = ( + pending: Pending | undefined, + atom: Atom, + a: AnyAtom, + aState: AtomState, + ) => { + if (import.meta.env?.MODE !== 'production' && a === atom) { + throw new Error('[Bug] atom cannot depend on itself') } - // Otherwise, check if the dependencies have changed. - // If all dependencies haven't changed, we can use the cache. - if ( - Array.from(atomState.d).every( - ([a, n]) => - // Recursively, read the atom state of the dependency, and - // check if the atom epoch number is unchanged - readAtomState(pending, a, force).n === n, + const atomState = getAtomState(atom) + atomState.d.set(a, aState.n) + const continuablePromise = getPendingContinuablePromise(atomState) + if (continuablePromise) { + addPendingContinuablePromiseToDependency( + atom, + continuablePromise, + aState, ) - ) { - return atomState } - } - // Compute a new state for this atom. - atomState.d.clear() - let isSync = true - const getter: Getter = (a: Atom) => { - const resolvedA = resolveAtom(a) - if (resolvedA === (atom as AnyAtom)) { - const aState = getAtomState(resolvedA) - if (!isAtomStateInitialized(aState)) { - if (hasInitialValue(resolvedA)) { - setAtomStateValueOrPromise(resolvedA, aState, resolvedA.init) - } else { - // NOTE invalid derived atoms can reach here - throw new Error('no atom init') - } - } - return returnAtomValue(aState) + aState.m?.t.add(atom) + if (pending) { + addPendingDependent(pending, a, atom) } - // resolvedA !== resolvedAtom - const aState = readAtomState(pending, resolvedA, force) - if (isSync) { - addDependency(pending, atom, resolvedA, aState) - } else { - const pending = createPending() - addDependency(pending, atom, resolvedA, aState) - mountDependencies(pending, atom, atomState) - flushPending(pending) - } - return returnAtomValue(aState) } - let controller: AbortController | undefined - let setSelf: ((...args: unknown[]) => unknown) | undefined - const options = { - get signal() { - if (!controller) { - controller = new AbortController() + + const readAtomState = ( + pending: Pending | undefined, + atom: Atom, + force?: (a: AnyAtom) => boolean, + ): AtomState => { + // See if we can skip recomputing this atom. + const atomState = getAtomState(atom) + if (!force?.(atom) && isAtomStateInitialized(atomState)) { + // If the atom is mounted, we can use the cache. + // because it should have been updated by dependencies. + if (atomState.m) { + return atomState } - return controller.signal - }, - get setSelf() { + // Otherwise, check if the dependencies have changed. + // If all dependencies haven't changed, we can use the cache. if ( - import.meta.env?.MODE !== 'production' && - !isActuallyWritableAtom(atom) + Array.from(atomState.d).every( + ([a, n]) => + // Recursively, read the atom state of the dependency, and + // check if the atom epoch number is unchanged + readAtomState(pending, a, force).n === n, + ) ) { - console.warn('setSelf function cannot be used with read-only atom') + return atomState } - if (!setSelf && isActuallyWritableAtom(atom)) { - setSelf = (...args) => { - if (import.meta.env?.MODE !== 'production' && isSync) { - console.warn('setSelf function cannot be called in sync') - } - if (!isSync) { - return writeAtom(atom, ...args) + } + // Compute a new state for this atom. + atomState.d.clear() + let isSync = true + const getter: Getter = (a: Atom) => { + const resolvedA = resolveAtom(a) + if (resolvedA === (atom as AnyAtom)) { + const aState = getAtomState(resolvedA) + if (!isAtomStateInitialized(aState)) { + if (hasInitialValue(resolvedA)) { + setAtomStateValueOrPromise(resolvedA, aState, resolvedA.init) + } else { + // NOTE invalid derived atoms can reach here + throw new Error('no atom init') } } + return returnAtomValue(aState) } - return setSelf - }, - } - try { - const valueOrPromise = atom.read(getter, options as never) - setAtomStateValueOrPromise( - atom, - atomState, - valueOrPromise, - () => controller?.abort(), - () => { - if (atomState.m) { - const pending = createPending() - mountDependencies(pending, atom, atomState) - flushPending(pending) + // resolvedA !== resolvedAtom + const aState = readAtomState(pending, resolvedA, force) + if (isSync) { + addDependency(pending, atom, resolvedA, aState) + } else { + const pending = createPending() + addDependency(pending, atom, resolvedA, aState) + mountDependencies(pending, atom, atomState) + flushPending(pending) + } + return returnAtomValue(aState) + } + let controller: AbortController | undefined + let setSelf: ((...args: unknown[]) => unknown) | undefined + const options = { + get signal() { + if (!controller) { + controller = new AbortController() } + return controller.signal }, - ) - return atomState - } catch (error) { - delete atomState.v - atomState.e = error - ++atomState.n - return atomState - } finally { - isSync = false + get setSelf() { + if ( + import.meta.env?.MODE !== 'production' && + !isActuallyWritableAtom(atom) + ) { + console.warn('setSelf function cannot be used with read-only atom') + } + if (!setSelf && isActuallyWritableAtom(atom)) { + setSelf = (...args) => { + if (import.meta.env?.MODE !== 'production' && isSync) { + console.warn('setSelf function cannot be called in sync') + } + if (!isSync) { + return writeAtom(atom, ...args) + } + } + } + return setSelf + }, + } + try { + const valueOrPromise = atom.read(getter, options as never) + setAtomStateValueOrPromise( + atom, + atomState, + valueOrPromise, + () => controller?.abort(), + () => { + if (atomState.m) { + const pending = createPending() + mountDependencies(pending, atom, atomState) + flushPending(pending) + } + }, + ) + return atomState + } catch (error) { + delete atomState.v + atomState.e = error + ++atomState.n + return atomState + } finally { + isSync = false + } } - } - const readAtom = (atom: Atom): Value => - returnAtomValue(readAtomState(undefined, resolveAtom(atom))) + const readAtom = (atom: Atom): Value => + returnAtomValue(readAtomState(undefined, resolveAtom(atom))) - const recomputeDependents = (pending: Pending, atom: AnyAtom) => { - const getDependents = (a: AnyAtom): Set => { - const aState = getAtomState(a) - const dependents = new Set(aState.m?.t) - for (const atomWithPendingContinuablePromise of aState.p) { - dependents.add(atomWithPendingContinuablePromise) + const recomputeDependents = (pending: Pending, atom: AnyAtom) => { + const getDependents = (a: AnyAtom): Set => { + const aState = getAtomState(a) + const dependents = new Set(aState.m?.t) + for (const atomWithPendingContinuablePromise of aState.p) { + dependents.add(atomWithPendingContinuablePromise) + } + getPendingDependents(pending, a)?.forEach((dependent) => { + dependents.add(dependent) + }) + return dependents } - getPendingDependents(pending, a)?.forEach((dependent) => { - dependents.add(dependent) - }) - 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 - - // 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: AnyAtom[] = [] - const markedAtoms = new Set() - const visit = (n: AnyAtom) => { - if (markedAtoms.has(n)) { - return - } - markedAtoms.add(n) - for (const m of getDependents(n)) { - if (n !== m) { - visit(m) + // 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: AnyAtom[] = [] + const markedAtoms = new Set() + const visit = (n: AnyAtom) => { + if (markedAtoms.has(n)) { + return + } + markedAtoms.add(n) + for (const m of getDependents(n)) { + if (n !== m) { + visit(m) + } } + // 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(n) } - // 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(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) - // 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]) - const isMarked = (a: AnyAtom) => markedAtoms.has(a) - for (let i = topsortedAtoms.length - 1; i >= 0; --i) { - const a = topsortedAtoms[i]! - const aState = getAtomState(a) - const prevEpochNumber = aState.n - let hasChangedDeps = false - for (const dep of aState.d.keys()) { - if (dep !== a && changedAtoms.has(dep)) { - hasChangedDeps = true - break + // 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) + // 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]) + const isMarked = (a: AnyAtom) => markedAtoms.has(a) + for (let i = topsortedAtoms.length - 1; i >= 0; --i) { + const a = topsortedAtoms[i]! + const aState = getAtomState(a) + const prevEpochNumber = aState.n + let hasChangedDeps = false + for (const dep of aState.d.keys()) { + if (dep !== a && changedAtoms.has(dep)) { + hasChangedDeps = true + break + } + } + if (hasChangedDeps) { + readAtomState(pending, a, isMarked) + mountDependencies(pending, a, aState) + if (prevEpochNumber !== aState.n) { + addPendingAtom(pending, a, aState) + changedAtoms.add(a) + } } + markedAtoms.delete(a) } - if (hasChangedDeps) { - readAtomState(pending, a, isMarked) - mountDependencies(pending, a, aState) - if (prevEpochNumber !== aState.n) { - addPendingAtom(pending, a, aState) - changedAtoms.add(a) + } + + const writeAtomState = ( + pending: Pending, + atom: WritableAtom, + ...args: Args + ): Result => { + const getter: Getter = (a: Atom) => + returnAtomValue(readAtomState(pending, resolveAtom(a))) + const setter: Setter = ( + a: WritableAtom, + ...args: As + ) => { + let r: R | undefined + const resolvedA = resolveAtom(a) + if (resolvedA === (atom as AnyAtom)) { + if (!hasInitialValue(resolvedA)) { + // NOTE technically possible but restricted as it may cause bugs + throw new Error('atom not writable') + } + const aState = getAtomState(resolvedA) + const hasPrevValue = 'v' in aState + const prevValue = aState.v + const v = args[0] as V + setAtomStateValueOrPromise(resolvedA, aState, v) + mountDependencies(pending, resolvedA, aState) + if (!hasPrevValue || !Object.is(prevValue, aState.v)) { + addPendingAtom(pending, resolvedA, aState) + recomputeDependents(pending, resolvedA) + } + } else { + r = writeAtomState(pending, resolvedA, ...args) } + flushPending(pending) + return r as R } - markedAtoms.delete(a) + const result = atom.write(getter, setter, ...args) + return result } - } - const writeAtomState = ( - pending: Pending, - atom: WritableAtom, - ...args: Args - ): Result => { - const getter: Getter = (a: Atom) => - returnAtomValue(readAtomState(pending, resolveAtom(a))) - const setter: Setter = ( - a: WritableAtom, - ...args: As + const writeAtom = ( + atom: WritableAtom, + ...args: Args + ): Result => { + const resolvedAtom = resolveAtom(atom) + const pending = createPending() + const result = writeAtomState(pending, resolvedAtom, ...args) + flushPending(pending) + return result + } + + const mountDependencies = ( + pending: Pending, + atom: AnyAtom, + atomState: AtomState, ) => { - let r: R | undefined - const resolvedA = resolveAtom(a) - if (resolvedA === (atom as AnyAtom)) { - if (!hasInitialValue(resolvedA)) { - // NOTE technically possible but restricted as it may cause bugs - throw new Error('atom not writable') + if (atomState.m && !getPendingContinuablePromise(atomState)) { + for (const a of atomState.d.keys()) { + if (!atomState.m.d.has(a)) { + const aMounted = mountAtom(pending, a) + aMounted.t.add(atom) + atomState.m.d.add(a) + } } - const aState = getAtomState(resolvedA) - const hasPrevValue = 'v' in aState - const prevValue = aState.v - const v = args[0] as V - setAtomStateValueOrPromise(resolvedA, aState, v) - mountDependencies(pending, resolvedA, aState) - if (!hasPrevValue || !Object.is(prevValue, aState.v)) { - addPendingAtom(pending, resolvedA, aState) - recomputeDependents(pending, resolvedA) + for (const a of atomState.m.d || []) { + if (!atomState.d.has(a)) { + const aMounted = unmountAtom(pending, a) + aMounted?.t.delete(atom) + atomState.m.d.delete(a) + } } - } else { - r = writeAtomState(pending, resolvedA, ...args) } - flushPending(pending) - return r as R } - const result = atom.write(getter, setter, ...args) - return result - } - const writeAtom = ( - atom: WritableAtom, - ...args: Args - ): Result => { - const resolvedAtom = resolveAtom(atom) - const pending = createPending() - const result = writeAtomState(pending, resolvedAtom, ...args) - flushPending(pending) - return result - } - - const mountDependencies = ( - pending: Pending, - atom: AnyAtom, - atomState: AtomState, - ) => { - if (atomState.m && !getPendingContinuablePromise(atomState)) { - for (const a of atomState.d.keys()) { - if (!atomState.m.d.has(a)) { + const mountAtom = (pending: Pending, atom: AnyAtom): Mounted => { + const atomState = getAtomState(atom) + if (!atomState.m) { + // recompute atom state + readAtomState(pending, atom) + // mount dependencies first + for (const a of atomState.d.keys()) { const aMounted = mountAtom(pending, a) aMounted.t.add(atom) - atomState.m.d.add(a) } - } - for (const a of atomState.m.d || []) { - if (!atomState.d.has(a)) { - const aMounted = unmountAtom(pending, a) - aMounted?.t.delete(atom) - atomState.m.d.delete(a) + // mount self + atomState.m = { + l: new Set(), + d: new Set(atomState.d.keys()), + t: new Set(), + } + if (import.meta.env?.MODE !== 'production') { + debugMountedAtoms.add(atom) + } + if (isActuallyWritableAtom(atom) && atom.onMount) { + const mounted = atomState.m + const { onMount } = atom + addPendingFunction(pending, () => { + const onUnmount = onMount((...args) => + writeAtomState(pending, atom, ...args), + ) + if (onUnmount) { + mounted.u = onUnmount + } + }) } } + return atomState.m } - } - - const mountAtom = (pending: Pending, atom: AnyAtom): Mounted => { - const atomState = getAtomState(atom) - if (!atomState.m) { - // recompute atom state - readAtomState(pending, atom) - // mount dependencies first - for (const a of atomState.d.keys()) { - const aMounted = mountAtom(pending, a) - aMounted.t.add(atom) - } - // mount self - atomState.m = { - l: new Set(), - d: new Set(atomState.d.keys()), - t: new Set(), - } - if (import.meta.env?.MODE !== 'production') { - debugMountedAtoms.add(atom) - } - if (isActuallyWritableAtom(atom) && atom.onMount) { - const mounted = atomState.m - const { onMount } = atom - addPendingFunction(pending, () => { - const onUnmount = onMount((...args) => - writeAtomState(pending, atom, ...args), - ) - if (onUnmount) { - mounted.u = onUnmount - } - }) - } - } - return atomState.m - } - const unmountAtom = ( - pending: Pending, - atom: AnyAtom, - ): Mounted | undefined => { - const atomState = getAtomState(atom) - if ( - atomState.m && - !atomState.m.l.size && - !Array.from(atomState.m.t).some((a) => getAtomState(a).m) - ) { - // unmount self - const onUnmount = atomState.m.u - if (onUnmount) { - addPendingFunction(pending, onUnmount) - } - delete atomState.m - if (import.meta.env?.MODE !== 'production') { - debugMountedAtoms.delete(atom) - } - // unmount dependencies - for (const a of atomState.d.keys()) { - const aMounted = unmountAtom(pending, a) - aMounted?.t.delete(atom) - } - // abort pending promise - const pendingPromise = getPendingContinuablePromise(atomState) - if (pendingPromise) { - // FIXME using `undefined` is kind of a hack. - pendingPromise[CONTINUE_PROMISE](undefined, () => {}) + const unmountAtom = ( + pending: Pending, + atom: AnyAtom, + ): Mounted | undefined => { + const atomState = getAtomState(atom) + if ( + atomState.m && + !atomState.m.l.size && + !Array.from(atomState.m.t).some((a) => getAtomState(a).m) + ) { + // unmount self + const onUnmount = atomState.m.u + if (onUnmount) { + addPendingFunction(pending, onUnmount) + } + delete atomState.m + if (import.meta.env?.MODE !== 'production') { + debugMountedAtoms.delete(atom) + } + // unmount dependencies + for (const a of atomState.d.keys()) { + const aMounted = unmountAtom(pending, a) + aMounted?.t.delete(atom) + } + // abort pending promise + const pendingPromise = getPendingContinuablePromise(atomState) + if (pendingPromise) { + // FIXME using `undefined` is kind of a hack. + pendingPromise[CONTINUE_PROMISE](undefined, () => {}) + } + return undefined } - return undefined + return atomState.m } - return atomState.m - } - const subscribeAtom = (atom: AnyAtom, listener: () => void) => { - const resolvedAtom = resolveAtom(atom) - const pending = createPending() - const mounted = mountAtom(pending, resolvedAtom) - flushPending(pending) - const listeners = mounted.l - listeners.add(listener) - return () => { - listeners.delete(listener) + const subscribeAtom = (atom: AnyAtom, listener: () => void) => { + const resolvedAtom = resolveAtom(atom) const pending = createPending() - unmountAtom(pending, resolvedAtom) + const mounted = mountAtom(pending, resolvedAtom) flushPending(pending) + const listeners = mounted.l + listeners.add(listener) + return () => { + listeners.delete(listener) + const pending = createPending() + unmountAtom(pending, resolvedAtom) + flushPending(pending) + } } - } - const store: Store = { - get: readAtom, - set: writeAtom, - sub: subscribeAtom, - } - if (import.meta.env?.MODE !== 'production') { - const devStore: DevStoreRev4 = { - // store dev methods (these are tentative and subject to change without notice) - dev4_get_internal_weak_map: () => atomStateMap, - dev4_get_mounted_atoms: () => debugMountedAtoms, - dev4_restore_atoms: (values) => { - const pending = createPending() - for (const [atom, value] of values) { - if (hasInitialValue(atom)) { - const aState = getAtomState(atom) - const hasPrevValue = 'v' in aState - const prevValue = aState.v - setAtomStateValueOrPromise(atom, aState, value) - mountDependencies(pending, atom, aState) - if (!hasPrevValue || !Object.is(prevValue, aState.v)) { - addPendingAtom(pending, atom, aState) - recomputeDependents(pending, atom) + const store: Store = { + get: readAtom, + set: writeAtom, + sub: subscribeAtom, + unstable_derive: buildStore, + } + if (import.meta.env?.MODE !== 'production') { + const devStore: DevStoreRev4 = { + // store dev methods (these are tentative and subject to change without notice) + dev4_get_internal_weak_map: () => atomStateMap, + dev4_get_mounted_atoms: () => debugMountedAtoms, + dev4_restore_atoms: (values) => { + const pending = createPending() + for (const [atom, value] of values) { + if (hasInitialValue(atom)) { + const aState = getAtomState(atom) + const hasPrevValue = 'v' in aState + const prevValue = aState.v + setAtomStateValueOrPromise(atom, aState, value) + mountDependencies(pending, atom, aState) + if (!hasPrevValue || !Object.is(prevValue, aState.v)) { + addPendingAtom(pending, atom, aState) + recomputeDependents(pending, atom) + } } } - } - flushPending(pending) - }, + flushPending(pending) + }, + } + Object.assign(store, devStore) } - Object.assign(store, devStore) + return store } - return store as Store + return buildStore() } let defaultStore: Store | undefined