From 8847084aa651851d50bff1ae548ef5e5fc2ad6b8 Mon Sep 17 00:00:00 2001 From: David Maskasky Date: Wed, 3 Apr 2024 19:01:26 -0700 Subject: [PATCH 1/2] [BREAKING] selectAtom does not resolve promises internally (#2435) * test: selectAtom should not return async value when the base atom curr and prev values are synchronous * selectAtom replaces promise with fulfilled value when promise resolves * replace refAtom with weakMap implementation * partition state by store * BREAKING: selectAtom does not internally resolve promises --------- Co-authored-by: David Maskasky Co-authored-by: Daishi Kato --- src/vanilla/utils/selectAtom.ts | 17 ++- tests/react/vanilla-utils/selectAtom.test.tsx | 106 +----------------- tests/vanilla/utils/types.test.tsx | 10 -- 3 files changed, 9 insertions(+), 124 deletions(-) diff --git a/src/vanilla/utils/selectAtom.ts b/src/vanilla/utils/selectAtom.ts index 1e705839d2..e759f5879f 100644 --- a/src/vanilla/utils/selectAtom.ts +++ b/src/vanilla/utils/selectAtom.ts @@ -17,20 +17,20 @@ const memo3 = ( export function selectAtom( anAtom: Atom, - selector: (v: Awaited, prevSlice?: Slice) => Slice, + selector: (v: Value, prevSlice?: Slice) => Slice, equalityFn?: (a: Slice, b: Slice) => boolean, -): Atom ? Promise : Slice> +): Atom export function selectAtom( anAtom: Atom, - selector: (v: Awaited, prevSlice?: Slice) => Slice, - equalityFn: (a: Slice, b: Slice) => boolean = Object.is, + selector: (v: Value, prevSlice?: Slice) => Slice, + equalityFn: (prevSlice: Slice, slice: Slice) => boolean = Object.is, ) { return memo3( () => { const EMPTY = Symbol() const selectValue = ([value, prevSlice]: readonly [ - Awaited, + Value, Slice | typeof EMPTY, ]) => { if (prevSlice === EMPTY) { @@ -39,15 +39,12 @@ export function selectAtom( const slice = selector(value, prevSlice) return equalityFn(prevSlice, slice) ? prevSlice : slice } - const derivedAtom: Atom | typeof EMPTY> & { + const derivedAtom: Atom & { init?: typeof EMPTY } = atom((get) => { const prev = get(derivedAtom) const value = get(anAtom) - if (value instanceof Promise || prev instanceof Promise) { - return Promise.all([value, prev] as const).then(selectValue) - } - return selectValue([value as Awaited, prev] as const) + return selectValue([value, prev] as const) }) // HACK to read derived atom before initialization derivedAtom.init = EMPTY diff --git a/tests/react/vanilla-utils/selectAtom.test.tsx b/tests/react/vanilla-utils/selectAtom.test.tsx index 9dcf10c498..cbd26cde15 100644 --- a/tests/react/vanilla-utils/selectAtom.test.tsx +++ b/tests/react/vanilla-utils/selectAtom.test.tsx @@ -1,7 +1,7 @@ -import { StrictMode, Suspense, useEffect, useRef } from 'react' +import { StrictMode, useEffect, useRef } from 'react' import { fireEvent, render } from '@testing-library/react' import { it } from 'vitest' -import { useAtom, useAtomValue, useSetAtom } from 'jotai/react' +import { useAtomValue, useSetAtom } from 'jotai/react' import { atom } from 'jotai/vanilla' import { selectAtom } from 'jotai/vanilla/utils' @@ -58,54 +58,6 @@ it('selectAtom works as expected', async () => { await findByText('a: 3') }) -it('selectAtom works with async atom', async () => { - const bigAtom = atom({ a: 0, b: 'othervalue' }) - const bigAtomAsync = atom((get) => Promise.resolve(get(bigAtom))) - const littleAtom = selectAtom(bigAtomAsync, (v) => v.a) - - const Parent = () => { - const setValue = useSetAtom(bigAtom) - return ( - <> - - - ) - } - - const Selector = () => { - const a = useAtomValue(littleAtom) - return ( - <> -
a: {a}
- - ) - } - - const { findByText, getByText } = render( - - - - - - , - ) - - await findByText('a: 0') - - fireEvent.click(getByText('increment')) - await findByText('a: 1') - fireEvent.click(getByText('increment')) - await findByText('a: 2') - fireEvent.click(getByText('increment')) - await findByText('a: 3') -}) - it('do not update unless equality function says value has changed', async () => { const bigAtom = atom({ a: 0 }) const littleAtom = selectAtom( @@ -177,57 +129,3 @@ it('do not update unless equality function says value has changed', async () => await findByText('value: {"a":3}') await findByText('commits: 4') }) - -it('equality function works even if suspend', async () => { - const bigAtom = atom({ a: 0 }) - const bigAtomAsync = atom((get) => Promise.resolve(get(bigAtom))) - const littleAtom = selectAtom( - bigAtomAsync, - (value) => value, - (left, right) => left.a === right.a, - ) - - const Controls = () => { - const [value, setValue] = useAtom(bigAtom) - return ( - <> -
bigValue: {JSON.stringify(value)}
- - - - ) - } - - const Selector = () => { - const value = useAtomValue(littleAtom) - return
littleValue: {JSON.stringify(value)}
- } - - const { findByText, getByText } = render( - - - - - - , - ) - - await findByText('bigValue: {"a":0}') - await findByText('littleValue: {"a":0}') - - fireEvent.click(getByText('increment')) - await findByText('bigValue: {"a":1}') - await findByText('littleValue: {"a":1}') - - fireEvent.click(getByText('other')) - await findByText('bigValue: {"a":1,"b":2}') - await findByText('littleValue: {"a":1}') -}) diff --git a/tests/vanilla/utils/types.test.tsx b/tests/vanilla/utils/types.test.tsx index 76db62a09a..8a2a07dd8b 100644 --- a/tests/vanilla/utils/types.test.tsx +++ b/tests/vanilla/utils/types.test.tsx @@ -10,16 +10,6 @@ it('selectAtom() should return the correct types', () => { const syncAtom = atom(0) const syncSelectedAtom = selectAtom(syncAtom, doubleCount) expectType, typeof syncSelectedAtom>>(true) - - const asyncAtom = atom(Promise.resolve(0)) - const asyncSelectedAtom = selectAtom(asyncAtom, doubleCount) - expectType>, typeof asyncSelectedAtom>>(true) - - const maybeAsyncAtom = atom(Promise.resolve(0) as number | Promise) - const maybeAsyncSelectedAtom = selectAtom(maybeAsyncAtom, doubleCount) - expectType< - TypeEqual>, typeof maybeAsyncSelectedAtom> - >(true) }) it('unwrap() should return the correct types', () => { From bf075ef8439daae9cc0dc3b5f504902cbcdc5f3e Mon Sep 17 00:00:00 2001 From: Daishi Kato Date: Thu, 4 Apr 2024 11:13:36 +0900 Subject: [PATCH 2/2] New store implementation as `store2.ts` (#2463) * wip * wip * wip * wip * wip * wip * wip * wip * wip * cherry-pick #2462 * wip * wip * wip * wip * wip * wip * wip * wip * wip * wip * follow #2471 * expose store2 so that we can experiment it * follow #2472 * separate experimental export * Update src/vanilla/store2.ts Co-authored-by: Iwo Plaza * refactor --------- Co-authored-by: Iwo Plaza --- package.json | 3 +- rollup.config.js | 4 + src/experimental.ts | 11 + src/types.d.ts | 1 + src/vanilla.ts | 14 +- src/vanilla/store2.ts | 724 +++++++++++++++++++++++++++++++++++ tests/vanilla/store.test.tsx | 6 +- 7 files changed, 760 insertions(+), 3 deletions(-) create mode 100644 src/experimental.ts create mode 100644 src/vanilla/store2.ts diff --git a/package.json b/package.json index cb793935d3..9005b55916 100644 --- a/package.json +++ b/package.json @@ -71,6 +71,7 @@ "build:vanilla:utils": "rollup -c --config-vanilla_utils", "build:react": "rollup -c --config-react --client-only", "build:react:utils": "rollup -c --config-react_utils --client-only", + "build:experimental": "rollup -c --config-experimental", "postbuild": "yarn patch-d-ts && yarn copy && yarn patch-ts3.8 && yarn patch-old-ts && yarn patch-esm-ts && yarn patch-readme", "prettier": "prettier '*.{js,json,md}' '{src,tests,benchmarks,docs}/**/*.{ts,tsx,md,mdx}' --write", "prettier:ci": "prettier '*.{js,json,md}' '{src,tests,benchmarks,docs}/**/*.{ts,tsx,md,mdx}' --list-different", @@ -78,7 +79,7 @@ "eslint:ci": "eslint --no-eslintrc --c .eslintrc.json '*.{js,json,ts}' '{src,tests,benchmarks}/**/*.{ts,tsx}'", "pretest": "tsc", "test": "vitest --ui --coverage", - "test:ci": "vitest", + "test:ci": "vitest && USE_STORE2=true vitest", "patch-d-ts": "node -e \"var {entries}=require('./rollup.config.js');require('shelljs').find('dist/**/*.d.ts').forEach(f=>{entries.forEach(({find,replacement})=>require('shelljs').sed('-i',new RegExp(' from \\''+find.source.slice(0,-1)+'\\';$'),' from \\''+replacement+'\\';',f));require('shelljs').sed('-i',/ from '(\\.[^']+)\\.ts';$/,' from \\'\\$1\\';',f)})\"", "copy": "shx cp -r dist/src/* dist/esm && shx cp -r dist/src/* dist && shx rm -rf dist/src && shx rm -rf dist/{src,tests} && downlevel-dts dist dist/ts3.8 --to=3.8 && shx cp package.json readme.md LICENSE dist && json -I -f dist/package.json -e \"this.private=false; this.devDependencies=undefined; this.optionalDependencies=undefined; this.scripts=undefined; this.prettier=undefined;\"", "patch-ts3.8": "node -e \"require('shelljs').find('dist/ts3.8/**/*.d.ts').forEach(f=>require('fs').appendFileSync(f,'declare type Awaited = T extends Promise ? V : T;'))\"", diff --git a/rollup.config.js b/rollup.config.js index afd185f7a6..f7e9368155 100644 --- a/rollup.config.js +++ b/rollup.config.js @@ -76,6 +76,7 @@ function createESMConfig(input, output, clientOnly) { 'import.meta.env?.MODE': '(import.meta.env ? import.meta.env.MODE : undefined)', }), + 'import.meta.env?.USE_STORE2': 'false', delimiters: ['\\b', '\\b(?!(\\.|/))'], preventAssignment: true, }), @@ -95,6 +96,7 @@ function createCommonJSConfig(input, output, clientOnly) { resolve({ extensions }), replace({ 'import.meta.env?.MODE': 'process.env.NODE_ENV', + 'import.meta.env?.USE_STORE2': 'false', delimiters: ['\\b', '\\b(?!(\\.|/))'], preventAssignment: true, }), @@ -132,6 +134,7 @@ function createUMDConfig(input, output, env, clientOnly) { resolve({ extensions }), replace({ 'import.meta.env?.MODE': JSON.stringify(env), + 'import.meta.env?.USE_STORE2': 'false', delimiters: ['\\b', '\\b(?!(\\.|/))'], preventAssignment: true, }), @@ -155,6 +158,7 @@ function createSystemConfig(input, output, env, clientOnly) { resolve({ extensions }), replace({ 'import.meta.env?.MODE': JSON.stringify(env), + 'import.meta.env?.USE_STORE2': 'false', delimiters: ['\\b', '\\b(?!(\\.|/))'], preventAssignment: true, }), diff --git a/src/experimental.ts b/src/experimental.ts new file mode 100644 index 0000000000..3b17070f81 --- /dev/null +++ b/src/experimental.ts @@ -0,0 +1,11 @@ +export { atom } from './vanilla/atom.ts' +export type { Atom, WritableAtom, PrimitiveAtom } from './vanilla/atom.ts' +export { createStore, getDefaultStore } from './vanilla/store2.ts' +export type { + Getter, + Setter, + ExtractAtomValue, + ExtractAtomArgs, + ExtractAtomResult, + SetStateAction, +} from './vanilla/typeUtils.ts' diff --git a/src/types.d.ts b/src/types.d.ts index f10fd4e38b..e099c12ad6 100644 --- a/src/types.d.ts +++ b/src/types.d.ts @@ -1,5 +1,6 @@ declare interface ImportMeta { env?: { MODE: string + USE_STORE2?: string } } diff --git a/src/vanilla.ts b/src/vanilla.ts index 3918c93f6e..5169a4c30d 100644 --- a/src/vanilla.ts +++ b/src/vanilla.ts @@ -1,6 +1,18 @@ export { atom } from './vanilla/atom.ts' export type { Atom, WritableAtom, PrimitiveAtom } from './vanilla/atom.ts' -export { createStore, getDefaultStore } from './vanilla/store.ts' + +// export { createStore, getDefaultStore } from './vanilla/store.ts' +import * as store from './vanilla/store.ts' +import * as store2 from './vanilla/store2.ts' +type CreateStore = typeof store.createStore +type GetDefaultStore = typeof store.getDefaultStore +export const createStore: CreateStore = import.meta.env?.USE_STORE2 + ? store2.createStore + : store.createStore +export const getDefaultStore: GetDefaultStore = import.meta.env?.USE_STORE2 + ? store2.getDefaultStore + : store.getDefaultStore + export type { Getter, Setter, diff --git a/src/vanilla/store2.ts b/src/vanilla/store2.ts new file mode 100644 index 0000000000..d3084248e1 --- /dev/null +++ b/src/vanilla/store2.ts @@ -0,0 +1,724 @@ +import type { Atom, WritableAtom } from './atom.ts' + +type AnyValue = unknown +type AnyError = unknown +type AnyAtom = Atom +type AnyWritableAtom = WritableAtom +type OnUnmount = () => void +type Getter = Parameters[0] +type Setter = Parameters[1] + +const isSelfAtom = (atom: AnyAtom, a: AnyAtom): boolean => + atom.unstable_is ? atom.unstable_is(a) : a === atom + +const hasInitialValue = >( + atom: T, +): atom is T & (T extends Atom ? { init: Value } : never) => + 'init' in atom + +const isActuallyWritableAtom = (atom: AnyAtom): atom is AnyWritableAtom => + !!(atom as AnyWritableAtom).write + +// +// Pending Set +// + +type PendingPair = [ + // TODO We should probably separate queues to notify and (un)mount + pendingForSync: Set void)> | undefined, + pendingForAsync: Set void)>, +] + +const createPendingPair = (): PendingPair => [new Set(), new Set()] + +const addPending = ( + pendingPair: PendingPair, + pending: readonly [AnyAtom, AtomState] | (() => void), +) => { + ;(pendingPair[0] || pendingPair[1]).add(pending) +} + +const flushPending = (pendingPair: PendingPair, isAsync?: true) => { + let pendingSet: Set void)> + if (isAsync) { + if (pendingPair[0]) { + // sync flush hasn't been called yet + return + } + pendingSet = pendingPair[1] + } else { + if (!pendingPair[0]) { + throw new Error('[Bug] cannot sync flush twice') + } + pendingSet = pendingPair[0] + } + const flushed = new Set() + while (pendingSet.size) { + const copy = new Set(pendingSet) + pendingSet.clear() + copy.forEach((pending) => { + if (typeof pending === 'function') { + pending() + } else { + const [atom, atomState] = pending + if (!flushed.has(atom) && atomState.m) { + atomState.m.l.forEach((listener) => listener()) + flushed.add(atom) + } + } + }) + } + pendingPair[0] = undefined + return flushed +} + +// +// Continuable Promise +// + +const CONTINUE_PROMISE = Symbol( + import.meta.env?.MODE !== 'production' ? 'CONTINUE_PROMISE' : '', +) + +const PENDING = 'pending' +const FULFILLED = 'fulfilled' +const REJECTED = 'rejected' + +type ContinuePromise = ( + nextPromise: PromiseLike | undefined, + nextAbort: () => void, +) => void + +type ContinuablePromise = Promise & + ( + | { status: typeof PENDING } + | { status: typeof FULFILLED; value?: T } + | { status: typeof REJECTED; reason?: AnyError } + ) & { + [CONTINUE_PROMISE]: ContinuePromise + } + +const isContinuablePromise = ( + promise: unknown, +): promise is ContinuablePromise => + typeof promise === 'object' && promise !== null && CONTINUE_PROMISE in promise + +const continuablePromiseMap: WeakMap< + PromiseLike, + ContinuablePromise +> = new WeakMap() + +/** + * Create a continuable promise from a regular promise. + */ +const createContinuablePromise = ( + promise: PromiseLike, + abort: () => void, + complete: () => void, +): ContinuablePromise => { + if (!continuablePromiseMap.has(promise)) { + let continuePromise: ContinuePromise + const p: any = new Promise((resolve, reject) => { + let curr = promise + const onFullfilled = (me: PromiseLike) => (v: T) => { + if (curr === me) { + p.status = FULFILLED + p.value = v + resolve(v) + complete() + } + } + const onRejected = (me: PromiseLike) => (e: AnyError) => { + if (curr === me) { + p.status = REJECTED + p.reason = e + reject(e) + complete() + } + } + promise.then(onFullfilled(promise), onRejected(promise)) + continuePromise = (nextPromise, nextAbort) => { + if (nextPromise) { + continuablePromiseMap.set(nextPromise, p) + curr = nextPromise + nextPromise.then(onFullfilled(nextPromise), onRejected(nextPromise)) + } + abort() + abort = nextAbort + } + }) + p.status = PENDING + p[CONTINUE_PROMISE] = continuePromise! + continuablePromiseMap.set(promise, p) + } + return continuablePromiseMap.get(promise) as ContinuablePromise +} + +const isPromiseLike = (x: unknown): x is PromiseLike => + typeof (x as any)?.then === 'function' + +const getPendingContinuablePromise = (atomState: AtomState) => { + const value: unknown = (atomState as any).s?.v + if (isContinuablePromise(value) && value.status === PENDING) { + return value + } + return null +} + +/** + * State tracked for mounted atoms. An atom is considered "mounted" if it has a + * subscriber, or is a transitive dependency of another atom that has a + * subscriber. + * + * The mounted state of an atom is freed once it is no longer mounted. + */ +type Mounted = { + /** Set of listeners to notify when the atom value changes. */ + readonly l: Set<() => void> + /** Set of mounted atoms that the atom depends on. */ + readonly d: Set + /** Function to run when the atom is unmounted. */ + u?: OnUnmount +} + +/** + * Mutable atom state, + * tracked for both mounted and unmounted atoms in a store. + */ +type AtomState = { + /** + * Map of atoms that the atom depends on. + * The map value is value/error of the dependency. + */ + readonly d: Map + /** Set of atoms that depends on the atom. */ + readonly t: Set + /** Object to store mounted state of the atom. */ + m?: Mounted // only available if the atom is mounted + /** Atom value, atom error or empty. */ + s?: { readonly v: Value } | { readonly e: AnyError } +} + +type WithS = T & { s: NonNullable } + +const returnAtomValue = (atomState: WithS>): Value => { + if ('e' in atomState.s) { + throw atomState.s.e + } + return atomState.s.v +} + +const setAtomStateValueOrPromise = ( + atomState: AtomState, + valueOrPromise: unknown, + abortPromise = () => {}, + completePromise = () => {}, +) => { + const pendingPromise = getPendingContinuablePromise(atomState) + if (isPromiseLike(valueOrPromise)) { + if (pendingPromise) { + if (pendingPromise !== valueOrPromise) { + pendingPromise[CONTINUE_PROMISE](valueOrPromise, abortPromise) + } + } else { + const continuablePromise = createContinuablePromise( + valueOrPromise, + abortPromise, + completePromise, + ) + atomState.s = { v: continuablePromise } + } + } else { + if (pendingPromise) { + pendingPromise[CONTINUE_PROMISE]( + Promise.resolve(valueOrPromise), + abortPromise, + ) + } + atomState.s = { v: valueOrPromise } + } +} + +// for debugging purpose only +type StoreListenerRev2 = ( + action: + | { type: 'write'; flushed: Set } + | { type: 'async-write'; flushed: Set } + | { type: 'sub'; flushed: Set } + | { type: 'unsub' } + | { type: 'restore'; flushed: Set }, +) => void + +type OldAtomState = { d: Map } & ( + | { e: AnyError } + | { v: AnyValue } +) +type OldMounted = { l: Set<() => void>; t: Set; u?: OnUnmount } + +export type Store = { + get: (atom: Atom) => Value + set: ( + atom: WritableAtom, + ...args: Args + ) => Result + sub: (atom: AnyAtom, listener: () => void) => () => void + dev_subscribe_store?: (l: StoreListenerRev2, rev: 2) => () => void + dev_get_mounted_atoms?: () => IterableIterator + dev_get_atom_state?: (a: AnyAtom) => OldAtomState | undefined + dev_get_mounted?: (a: AnyAtom) => OldMounted | undefined + dev_restore_atoms?: (values: Iterable) => void +} + +export const createStore = (): Store => { + const atomStateMap = new WeakMap() + + const getAtomState = (atom: Atom) => { + let atomState = atomStateMap.get(atom) as AtomState | undefined + if (!atomState) { + atomState = { d: new Map(), t: new Set() } + atomStateMap.set(atom, atomState) + } + return atomState + } + + let storeListenersRev2: Set + let mountedAtoms: Set + if (import.meta.env?.MODE !== 'production') { + storeListenersRev2 = new Set() + mountedAtoms = new Set() + } + + const clearDependencies = (atom: Atom) => { + const atomState = getAtomState(atom) + for (const a of atomState.d.keys()) { + getAtomState(a).t.delete(atom) + } + atomState.d.clear() + } + + const addDependency = ( + atom: Atom, + a: AnyAtom, + aState: WithS, + isSync: boolean, + ) => { + 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.s) + aState.t.add(atom) + if (!isSync && atomState.m) { + const pendingPair = createPendingPair() + mountDependencies(pendingPair, atomState) + flushPending(pendingPair) + } + } + + const readAtomState = ( + atom: Atom, + force?: true, + ): WithS> => { + // See if we can skip recomputing this atom. + const atomState = getAtomState(atom) + if (!force && 's' in atomState) { + // If the atom is mounted, we can use the cache. + // because it should have been updated by dependencies. + if (atomState.m) { + return atomState as WithS + } + // 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, s]) => { + // Recursively, read the atom state of the dependency, and + const aState = readAtomState(a) + // Check if the atom value is unchanged + return 'v' in s && 'v' in aState.s && Object.is(s.v, aState.s.v) + }) + ) { + return atomState as WithS + } + } + // Compute a new state for this atom. + clearDependencies(atom) + let isSync = true + const getter: Getter = (a: Atom) => { + if (isSelfAtom(atom, a)) { + const aState = getAtomState(a) + if (!aState.s) { + if (hasInitialValue(a)) { + setAtomStateValueOrPromise(aState, a.init) + } else { + // NOTE invalid derived atoms can reach here + throw new Error('no atom init') + } + } + return returnAtomValue(aState as WithS) + } + // a !== atom + const aState = readAtomState(a) + addDependency(atom, a, aState, isSync) + 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) + setAtomStateValueOrPromise( + atomState, + valueOrPromise, + () => controller?.abort(), + () => { + if (atomState.m) { + const pendingPair = createPendingPair() + mountDependencies(pendingPair, atomState) + flushPending(pendingPair) + } + }, + ) + return atomState as WithS + } catch (error) { + atomState.s = { e: error } + return atomState as WithS + } finally { + isSync = false + } + } + + const readAtom = (atom: Atom): Value => + returnAtomValue(readAtomState(atom)) + + const recomputeDependents = (pendingPair: PendingPair, atom: AnyAtom) => { + // 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 getAtomState(n).t) { + // we shouldn't use isSelfAtom here. + 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) + } + // 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]) + for (let i = topsortedAtoms.length - 1; i >= 0; --i) { + const a = topsortedAtoms[i]! + const aState = getAtomState(a) + const prev = aState.s + let hasChangedDeps = false + for (const dep of aState.d.keys()) { + if (dep !== a && changedAtoms.has(dep)) { + hasChangedDeps = true + break + } + } + if (hasChangedDeps) { + // only recompute if it is mounted or it has a pending promise + if (aState.m || getPendingContinuablePromise(aState)) { + readAtomState(a, true) + mountDependencies(pendingPair, aState) + if ( + !prev || + !('v' in prev) || + !('v' in aState.s!) || + !Object.is(prev.v, aState.s.v) + ) { + addPending(pendingPair, [a, aState]) + changedAtoms.add(a) + } + } + } + } + } + + const writeAtomState = ( + pendingPair: PendingPair, + atom: WritableAtom, + ...args: Args + ): Result => { + const getter: Getter = (a: Atom) => returnAtomValue(readAtomState(a)) + const setter: Setter = ( + a: WritableAtom, + ...args: As + ) => { + let r: R | undefined + if (isSelfAtom(atom, a)) { + if (!hasInitialValue(a)) { + // NOTE technically possible but restricted as it may cause bugs + throw new Error('atom not writable') + } + const aState = getAtomState(a) + const prev = aState.s + const v = args[0] as V + setAtomStateValueOrPromise(aState, v) + mountDependencies(pendingPair, aState) + const curr = (aState as WithS).s + if ( + !prev || + !('v' in prev) || + !('v' in curr) || + !Object.is(prev.v, curr.v) + ) { + addPending(pendingPair, [a, aState]) + recomputeDependents(pendingPair, a) + } + } else { + r = writeAtomState(pendingPair, a as AnyWritableAtom, ...args) as R + } + const flushed = flushPending(pendingPair, true) + if (import.meta.env?.MODE !== 'production' && flushed) { + storeListenersRev2.forEach((l) => l({ type: 'async-write', flushed })) + } + return r as R + } + const result = atom.write(getter, setter, ...args) + return result + } + + const writeAtom = ( + atom: WritableAtom, + ...args: Args + ): Result => { + const pendingPair = createPendingPair() + const result = writeAtomState(pendingPair, atom, ...args) + const flushed = flushPending(pendingPair) + if (import.meta.env?.MODE !== 'production') { + storeListenersRev2.forEach((l) => l({ type: 'write', flushed: flushed! })) + } + return result + } + + const mountDependencies = ( + pendingPair: PendingPair, + atomState: AtomState, + ) => { + if (atomState.m && !getPendingContinuablePromise(atomState)) { + for (const a of atomState.d.keys()) { + if (!atomState.m.d.has(a)) { + mountAtom(pendingPair, a) + atomState.m.d.add(a) + } + } + for (const a of atomState.m.d || []) { + if (!atomState.d.has(a)) { + unmountAtom(pendingPair, a) + atomState.m.d.delete(a) + } + } + } + } + + const mountAtom = (pendingPair: PendingPair, atom: AnyAtom): Mounted => { + const atomState = getAtomState(atom) + if (!atomState.m) { + // recompute atom state + readAtomState(atom) + // mount dependents first + for (const a of atomState.d.keys()) { + mountAtom(pendingPair, a) + } + // mount self + atomState.m = { l: new Set(), d: new Set(atomState.d.keys()) } + if (import.meta.env?.MODE !== 'production') { + mountedAtoms.add(atom) + } + if (isActuallyWritableAtom(atom) && atom.onMount) { + const mounted = atomState.m + const { onMount } = atom + addPending(pendingPair, () => { + const onUnmount = onMount((...args) => + writeAtomState(pendingPair, atom, ...args), + ) + if (onUnmount) { + mounted.u = onUnmount + } + }) + } + } + return atomState.m + } + + const unmountAtom = (pendingPair: PendingPair, atom: AnyAtom) => { + const atomState = getAtomState(atom) + if ( + atomState.m && + !atomState.m.l.size && + !Array.from(atomState.t).some((a) => getAtomState(a).m) + ) { + // unmount self + const onUnmount = atomState.m.u + if (onUnmount) { + addPending(pendingPair, onUnmount) + } + delete atomState.m + if (import.meta.env?.MODE !== 'production') { + mountedAtoms.delete(atom) + } + // unmount dependencies + for (const a of atomState.d.keys()) { + unmountAtom(pendingPair, a) + } + // abort pending promise + const pendingPromise = getPendingContinuablePromise(atomState) + if (pendingPromise) { + // FIXME using `undefined` is kind of a hack. + pendingPromise[CONTINUE_PROMISE](undefined, () => {}) + } + } + } + + const subscribeAtom = (atom: AnyAtom, listener: () => void) => { + let prevMounted: Mounted | undefined + if (import.meta.env?.MODE !== 'production') { + prevMounted = atomStateMap.get(atom)?.m + } + const pendingPair = createPendingPair() + const mounted = mountAtom(pendingPair, atom) + const flushed = flushPending(pendingPair) + const listeners = mounted.l + listeners.add(listener) + if (import.meta.env?.MODE !== 'production') { + if (!prevMounted) { + flushed!.add(atom) // HACK to include self + } + storeListenersRev2.forEach((l) => l({ type: 'sub', flushed: flushed! })) + } + return () => { + listeners.delete(listener) + const pendingPair = createPendingPair() + unmountAtom(pendingPair, atom) + flushPending(pendingPair) + if (import.meta.env?.MODE !== 'production') { + // devtools uses this to detect if it _can_ unmount or not + storeListenersRev2.forEach((l) => l({ type: 'unsub' })) + } + } + } + + if (import.meta.env?.MODE !== 'production') { + return { + get: readAtom, + set: writeAtom, + sub: subscribeAtom, + // store dev methods (these are tentative and subject to change without notice) + dev_subscribe_store: (l: StoreListenerRev2, rev: 2) => { + if (rev !== 2) { + throw new Error('The current StoreListener revision is 2.') + } + storeListenersRev2.add(l as StoreListenerRev2) + return () => { + storeListenersRev2.delete(l as StoreListenerRev2) + } + }, + dev_get_mounted_atoms: () => mountedAtoms.values(), + dev_get_atom_state: (a: AnyAtom) => { + const getOldAtomState = (a: AnyAtom): OldAtomState | undefined => { + const aState = atomStateMap.get(a) + return ( + aState && + aState.s && { + d: new Map( + Array.from(aState.d.keys()).flatMap((a) => { + const s = getOldAtomState(a) + return s ? [[a, s]] : [] + }), + ), + ...aState.s, + } + ) + } + return getOldAtomState(a) + }, + dev_get_mounted: (a: AnyAtom) => { + const aState = atomStateMap.get(a) + return ( + aState && + aState.m && { + l: aState.m.l, + t: new Set([...aState.t, a]), // HACK to include self + ...(aState.m.u ? { u: aState.m.u } : {}), + } + ) + }, + dev_restore_atoms: (values: Iterable) => { + const pendingPair = createPendingPair() + for (const [atom, value] of values) { + setAtomStateValueOrPromise(getAtomState(atom), value) + recomputeDependents(pendingPair, atom) + } + const flushed = flushPending(pendingPair) + storeListenersRev2.forEach((l) => + l({ type: 'restore', flushed: flushed! }), + ) + }, + } + } + return { + get: readAtom, + set: writeAtom, + sub: subscribeAtom, + } +} + +let defaultStore: Store | undefined + +export const getDefaultStore = (): Store => { + if (!defaultStore) { + defaultStore = createStore() + if (import.meta.env?.MODE !== 'production') { + ;(globalThis as any).__JOTAI_DEFAULT_STORE__ ||= defaultStore + if ((globalThis as any).__JOTAI_DEFAULT_STORE__ !== defaultStore) { + console.warn( + 'Detected multiple Jotai instances. It may cause unexpected behavior with the default store. https://github.com/pmndrs/jotai/discussions/2044', + ) + } + } + } + return defaultStore +} diff --git a/tests/vanilla/store.test.tsx b/tests/vanilla/store.test.tsx index 8c6d6f7873..a38ee11cd0 100644 --- a/tests/vanilla/store.test.tsx +++ b/tests/vanilla/store.test.tsx @@ -45,7 +45,11 @@ describe('[DEV-ONLY] dev-only methods', () => { store.set(countAtom, 1) const result = store.dev_get_mounted_atoms?.() || [] - expect(Array.from(result)).toStrictEqual([ + expect( + Array.from(result).sort( + (a, b) => Object.keys(a).length - Object.keys(b).length, + ), + ).toStrictEqual([ { toString: expect.any(Function), read: expect.any(Function) }, { toString: expect.any(Function),