From ece5b0bb34a8b00013e7fad1f6cc10de004d47c5 Mon Sep 17 00:00:00 2001 From: David Maskasky Date: Fri, 10 Jan 2025 22:56:31 -0800 Subject: [PATCH] attempt 2: refactor with store hooks --- .eslintrc.json | 26 +- .prettierrc | 2 +- __tests__/atomEffect.test.tsx | 90 ++-- __tests__/store.ts | 761 ++++++++++++++++++++++++++++++++++ package.json | 2 +- pnpm-lock.yaml | 10 +- src/atomEffect.ts | 224 +++++----- 7 files changed, 937 insertions(+), 178 deletions(-) create mode 100644 __tests__/store.ts diff --git a/.eslintrc.json b/.eslintrc.json index 03cbeec..951ceb4 100644 --- a/.eslintrc.json +++ b/.eslintrc.json @@ -14,14 +14,7 @@ "plugin:import/errors", "plugin:import/warnings" ], - "plugins": [ - "@typescript-eslint", - "react", - "prettier", - "react-hooks", - "import", - "jest" - ], + "plugins": ["@typescript-eslint", "react", "prettier", "react-hooks", "import", "jest"], "parser": "@typescript-eslint/parser", "parserOptions": { "ecmaVersion": 2018, @@ -39,6 +32,7 @@ "import/no-unresolved": ["error", { "commonjs": true, "amd": true }], "import/export": "error", "import/no-duplicates": ["error"], + "prettier/prettier": ["error", { "printWidth": 100 }], "@typescript-eslint/explicit-module-boundary-types": "off", "@typescript-eslint/no-unused-vars": [ "warn", @@ -47,23 +41,13 @@ "@typescript-eslint/no-use-before-define": "off", "@typescript-eslint/no-empty-function": "off", "@typescript-eslint/no-explicit-any": "off", - "jest/consistent-test-it": [ - "error", - { "fn": "it", "withinDescribe": "it" } - ], + "@typescript-eslint/no-unused-expressions": "off", + "jest/consistent-test-it": ["error", { "fn": "it", "withinDescribe": "it" }], "import/order": [ "error", { "alphabetize": { "order": "asc", "caseInsensitive": true }, - "groups": [ - "builtin", - "external", - "internal", - "parent", - "sibling", - "index", - "object" - ], + "groups": ["builtin", "external", "internal", "parent", "sibling", "index", "object"], "newlines-between": "never", "pathGroups": [ { diff --git a/.prettierrc b/.prettierrc index 77f2042..864f8d4 100644 --- a/.prettierrc +++ b/.prettierrc @@ -4,5 +4,5 @@ "singleQuote": true, "bracketSameLine": true, "tabWidth": 2, - "printWidth": 80 + "printWidth": 100 } diff --git a/__tests__/atomEffect.test.tsx b/__tests__/atomEffect.test.tsx index 7cf2b7f..2192c8c 100644 --- a/__tests__/atomEffect.test.tsx +++ b/__tests__/atomEffect.test.tsx @@ -1,27 +1,23 @@ import React, { createElement, useEffect } from 'react' import { act, render, renderHook, waitFor } from '@testing-library/react' import { Provider, useAtom, useAtomValue, useSetAtom } from 'jotai/react' -import { atom, createStore, getDefaultStore } from 'jotai/vanilla' +import { atom } from 'jotai/vanilla' import { atomEffect } from '../src/atomEffect' -import { - ErrorBoundary, - assert, - delay, - increment, - incrementLetter, -} from './test-utils' +import { createStore } from './store' +import { ErrorBoundary, assert, delay, increment, incrementLetter } from './test-utils' it('should run the effect on vanilla store', () => { - const store = createStore().unstable_derive( - (getAtomState, setAtomState, ...rest) => [ - getAtomState, - (atom, atomState) => - Object.assign(setAtomState(atom, atomState), { + const store = createStore().unstable_derive((getAtomState, setAtomState, ...rest) => [ + getAtomState, + (atom, atomState) => + setAtomState( + atom, + Object.assign(atomState, { label: atom.debugLabel, - }), - ...rest, - ] - ) + }) + ), + ...rest, + ]) const countAtom = atom(0) countAtom.debugLabel = 'count' const effectAtom = atomEffect((_, set) => { @@ -85,6 +81,7 @@ it('should run the effect on mount and cleanup on unmount and whenever countAtom }) let didMount = false + const store = createStore() function useTest() { const [count, setCount] = useAtom(countAtom) useAtomValue(effectAtom) @@ -93,7 +90,10 @@ it('should run the effect on mount and cleanup on unmount and whenever countAtom }, [count]) return setCount } - const { result, rerender, unmount } = renderHook(useTest) + const wrapper = ({ children }: { children: React.ReactNode }) => ( + {children} + ) + const { result, rerender, unmount } = renderHook(useTest, { wrapper }) function incrementCount() { const setCount = result.current setCount(increment) @@ -142,7 +142,7 @@ it('should not cause infinite loops when effect updates the watched atom', () => runCount++ set(watchedAtom, increment) }) - const store = getDefaultStore() + const store = createStore() store.sub(effectAtom, () => void 0) const incrementWatched = () => store.set(watchedAtom, increment) @@ -164,7 +164,7 @@ it('should not cause infinite loops when effect updates the watched atom asynchr set(watchedAtom, increment) }, 0) }) - const store = getDefaultStore() + const store = createStore() store.sub(effectAtom, () => void 0) // changing the value should run the effect again one time store.set(watchedAtom, increment) @@ -183,7 +183,7 @@ it('should allow synchronous recursion with set.recurse for first run', () => { } recurse(watchedAtom, increment) }) - const store = getDefaultStore() + const store = createStore() store.sub(effectAtom, () => void 0) expect({ runCount, watched: store.get(watchedAtom) }).toEqual({ runCount: 4, // 2 @@ -206,7 +206,7 @@ it('should allow synchronous recursion with set.recurse', () => { } recurse(watchedAtom, increment) }) - const store = getDefaultStore() + const store = createStore() store.sub(effectAtom, () => void 0) store.set(watchedAtom, increment) expect(store.get(watchedAtom)).toBe(5) @@ -229,7 +229,7 @@ it('should allow multiple synchronous recursion with set.recurse', () => { recurse(watchedAtom, increment) recurse(watchedAtom, increment) }) - const store = getDefaultStore() + const store = createStore() store.sub(effectAtom, () => void 0) store.set(watchedAtom, increment) expect({ runCount, value: store.get(watchedAtom) }).toEqual({ @@ -266,7 +266,7 @@ it('should batch updates during synchronous recursion with set.recurse', () => { ]) set.recurse(updateAtom) }) - const store = getDefaultStore() + const store = createStore() store.sub(effectAtom, () => void 0) store.set(watchedAtom, increment) expect(store.get(lettersAndNumbersAtom)).toEqual(['a0', 'b1']) @@ -289,7 +289,7 @@ it('should allow asynchronous recursion with task delay with set.recurse', async recurse(watchedAtom, increment) }) }) - const store = getDefaultStore() + const store = createStore() store.sub(effectAtom, () => void 0) await waitFor(() => assert(done)) expect(store.get(watchedAtom)).toBe(3) @@ -311,7 +311,7 @@ it('should allow asynchronous recursion with microtask delay with set.recurse', recurse(watchedAtom, increment) }) }) - const store = getDefaultStore() + const store = createStore() store.sub(effectAtom, () => void 0) await waitFor(() => assert(store.get(watchedAtom) >= 3)) expect(store.get(watchedAtom)).toBe(3) @@ -321,23 +321,26 @@ it('should allow asynchronous recursion with microtask delay with set.recurse', it('should work with both set.recurse and set', () => { expect.assertions(3) let runCount = 0 - const watchedAtom = atom(0) + const valueAtom = atom(0) const countAtom = atom(0) const effectAtom = atomEffect((get, set) => { - const value = get(watchedAtom) + const value = get(valueAtom) + if (value >= 5) { + throw new Error() + } get(countAtom) runCount++ if (value === 0 || value % 3) { - set.recurse(watchedAtom, increment) + set.recurse(valueAtom, increment) set(countAtom, increment) return } - set(watchedAtom, increment) + set(valueAtom, increment) }) - const store = getDefaultStore() + const store = createStore() store.sub(effectAtom, () => void 0) expect(store.get(countAtom)).toBe(3) - expect(store.get(watchedAtom)).toBe(4) + expect(store.get(valueAtom)).toBe(4) expect(runCount).toBe(4) }) @@ -354,7 +357,7 @@ it('should disallow synchronous set.recurse in cleanup', () => { }) return cleanup }) - const store = getDefaultStore() + const store = createStore() store.sub(effectAtom, () => void 0) store.set(anotherAtom, increment) expect(() => store.set(anotherAtom, increment)).toThrowError( @@ -380,7 +383,7 @@ it('should return value from set.recurse', () => { return } }) - const store = getDefaultStore() + const store = createStore() store.sub(effectAtom, () => void 0) expect(results).toEqual([1, 2, 3, 4, 5]) }) @@ -654,7 +657,7 @@ it('should batch synchronous updates as a single transaction', () => { letters + String(numbers), ]) }) - const store = getDefaultStore() + const store = createStore() store.sub(effectAtom, () => void 0) expect(runCount).toBe(1) @@ -722,11 +725,9 @@ it('should abort the previous promise', async () => { const completedRuns: number[] = [] const resolves: (() => void)[] = [] const countAtom = atom(0) - const abortControllerAtom = atom<{ abortController: AbortController | null }>( - { - abortController: null, - } - ) + const abortControllerAtom = atom<{ abortController: AbortController | null }>({ + abortController: null, + }) const effectAtom = atomEffect((get) => { const currentRun = runCount++ get(countAtom) @@ -818,7 +819,7 @@ it('should not infinite loop with nested atomEffects', async () => { get(readOnlyAtom) }) - const store = getDefaultStore() + const store = createStore() store.sub(effect2Atom, () => void 0) await waitFor(() => assert(delayedIncrement)) @@ -851,7 +852,7 @@ it('should not rerun with get.peek', () => { get.peek(countAtom) runCount++ }) - const store = getDefaultStore() + const store = createStore() store.sub(effectAtom, () => void 0) store.set(countAtom, increment) expect(runCount).toBe(1) @@ -869,10 +870,7 @@ it('should trigger the error boundary when an error is thrown', async () => { let didThrow = false function wrapper() { return ( - (didThrow = true)} - children={} - /> + (didThrow = true)} children={} /> ) } const originalConsoleError = console.error diff --git a/__tests__/store.ts b/__tests__/store.ts new file mode 100644 index 0000000..3468eee --- /dev/null +++ b/__tests__/store.ts @@ -0,0 +1,761 @@ +import type { Atom, WritableAtom } from 'jotai/vanilla' + +type AnyValue = unknown +type AnyError = unknown +type AnyAtom = Atom +type AnyWritableAtom = WritableAtom +type OnUnmount = () => void +type Getter = Parameters[0] +type Setter = Parameters[1] +type EpochNumber = number + +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 + +// +// Cancelable Promise +// + +type CancelHandler = (nextValue: unknown) => void +type PromiseState = [cancelHandlers: Set, settled: boolean] + +const cancelablePromiseMap = new WeakMap, PromiseState>() + +const isPendingPromise = (value: unknown): value is PromiseLike => + isPromiseLike(value) && !cancelablePromiseMap.get(value)?.[1] + +const cancelPromise = (promise: PromiseLike, nextValue: unknown) => { + const promiseState = cancelablePromiseMap.get(promise) + if (promiseState) { + promiseState[1] = true + promiseState[0].forEach((fn) => fn(nextValue)) + } else if (process.env?.MODE !== 'production') { + throw new Error('[Bug] cancelable promise not found') + } +} + +const patchPromiseForCancelability = (promise: PromiseLike) => { + if (cancelablePromiseMap.has(promise)) { + // already patched + return + } + const promiseState: PromiseState = [new Set(), false] + cancelablePromiseMap.set(promise, promiseState) + const settle = () => { + promiseState[1] = true + } + promise.then(settle, settle) + ;(promise as { onCancel?: (fn: CancelHandler) => void }).onCancel = (fn) => { + promiseState[0].add(fn) + } +} + +const isPromiseLike = ( + p: unknown +): p is PromiseLike & { onCancel?: (fn: CancelHandler) => void } => + typeof (p as any)?.then === 'function' + +/** + * 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 + /** Set of mounted atoms that depends on the atom. */ + readonly t: Set + /** Function to run when the atom is unmounted. */ + u?: () => void +} + +/** + * Mutable atom state, + * tracked for both mounted and unmounted atoms in a store. + * + * This should be garbage collectable. + * We can mutate it during atom read. (except for fields with TODO) + */ +type AtomState = { + /** + * Map of atoms that the atom depends on. + * The map value is the epoch number of the dependency. + */ + readonly d: Map + /** + * Set of atoms with pending promise that depend on the atom. + * + * This may cause memory leaks, but it's for the capability to continue promises + * TODO(daishi): revisit how to handle this + */ + readonly p: Set + /** The epoch number of the atom. */ + n: EpochNumber + /** + * Object to store mounted state of the atom. + * TODO(daishi): move this out of AtomState + */ + m?: Mounted // only available if the atom is mounted + /** + * Listener to notify when the atom value is updated. + * This is an experimental API and will be changed in the next minor. + * TODO(daishi): move this store hooks + */ + u?: () => void + /** + * Listener to notify when the atom is mounted or unmounted. + * This is an experimental API and will be changed in the next minor. + * TODO(daishi): move this store hooks + */ + h?: () => void + /** Atom value */ + v?: Value + /** Atom error */ + e?: AnyError +} + +const isAtomStateInitialized = (atomState: AtomState) => + 'v' in atomState || 'e' in atomState + +const returnAtomValue = (atomState: AtomState): Value => { + if ('e' in atomState) { + throw atomState.e + } + if (process.env?.MODE !== 'production' && !('v' in atomState)) { + throw new Error('[Bug] atom state is not initialized') + } + return atomState.v! +} + +const addPendingPromiseToDependency = ( + atom: AnyAtom, + promise: PromiseLike, + dependencyAtomState: AtomState +) => { + if (!dependencyAtomState.p.has(atom)) { + dependencyAtomState.p.add(atom) + promise.then( + () => { + dependencyAtomState.p.delete(atom) + }, + () => { + dependencyAtomState.p.delete(atom) + } + ) + } +} + +const addDependency = ( + atom: Atom, + atomState: AtomState, + a: AnyAtom, + aState: AtomState +) => { + if (process.env?.MODE !== 'production' && a === atom) { + throw new Error('[Bug] atom cannot depend on itself') + } + atomState.d.set(a, aState.n) + if (isPendingPromise(atomState.v)) { + addPendingPromiseToDependency(atom, atomState.v, aState) + } + aState.m?.t.add(atom) +} + +// internal & unstable type +type StoreArgs = readonly [ + getAtomState: (atom: Atom) => AtomState | undefined, + setAtomState: (atom: Atom, atomState: AtomState) => void, + atomRead: (atom: Atom, ...params: Parameters['read']>) => Value, + atomWrite: ( + atom: WritableAtom, + ...params: Parameters['write']> + ) => Result, + atomOnInit: (atom: Atom, store: Store) => void, + atomOnMount: ( + atom: WritableAtom, + setAtom: (...args: Args) => Result + ) => OnUnmount | void, +] + +// for debugging purpose only +type DevStoreRev4 = { + dev4_get_internal_weak_map: () => { + get: (atom: AnyAtom) => AtomState | undefined + } + dev4_get_mounted_atoms: () => Set + dev4_restore_atoms: (values: Iterable) => void +} + +type Store = { + get: (atom: Atom) => Value + set: ( + atom: WritableAtom, + ...args: Args + ) => Result + sub: (atom: AnyAtom, listener: () => void) => () => void + unstable_derive: (fn: (...args: StoreArgs) => StoreArgs) => Store +} + +export type INTERNAL_DevStoreRev4 = DevStoreRev4 +export type INTERNAL_PrdStore = Store + +/** + * This is an experimental API and will be changed in the next minor. + */ +const INTERNAL_flushStoreHook = Symbol.for('JOTAI.EXPERIMENTAL.FLUSHSTOREHOOK') + +const buildStore = (...storeArgs: StoreArgs): Store => { + const [getAtomState, setAtomState, atomRead, atomWrite, atomOnInit, atomOnMount] = storeArgs + const ensureAtomState = (atom: Atom) => { + if (process.env?.MODE !== 'production' && !atom) { + throw new Error('Atom is undefined or null') + } + let atomState = getAtomState(atom) + if (!atomState) { + atomState = { d: new Map(), p: new Set(), n: 0 } + setAtomState(atom, atomState) + atomOnInit?.(atom, store) + } + return atomState + } + + // These are store state. + // As they are not garbage collectable, they shouldn't be mutated during atom read. + const invalidatedAtoms = new WeakMap() + const changedAtoms = new Map() + const unmountCallbacks = new Set<() => void>() + const mountCallbacks = new Set<() => void>() + + const flushCallbacks = () => { + const errors: unknown[] = [] + const call = (fn: () => void) => { + try { + fn() + } catch (e) { + errors.push(e) + } + } + do { + ;(store as any)[INTERNAL_flushStoreHook]?.() + changedAtoms.forEach((atomState) => atomState.m?.l.forEach(call)) + changedAtoms.clear() + unmountCallbacks.forEach(call) + unmountCallbacks.clear() + mountCallbacks.forEach(call) + mountCallbacks.clear() + recomputeInvalidatedAtoms() + } while (changedAtoms.size) + if (errors.length) { + throw errors[0] + } + } + + const setAtomStateValueOrPromise = ( + atom: AnyAtom, + atomState: AtomState, + valueOrPromise: unknown + ) => { + const hasPrevValue = 'v' in atomState + const prevValue = atomState.v + const pendingPromise = isPendingPromise(atomState.v) ? atomState.v : null + if (isPromiseLike(valueOrPromise)) { + patchPromiseForCancelability(valueOrPromise) + for (const a of atomState.d.keys()) { + addPendingPromiseToDependency(atom, valueOrPromise, ensureAtomState(a)) + } + atomState.v = valueOrPromise + } else { + atomState.v = valueOrPromise + } + delete atomState.e + if (!hasPrevValue || !Object.is(prevValue, atomState.v)) { + ++atomState.n + if (pendingPromise) { + cancelPromise(pendingPromise, valueOrPromise) + } + } + } + + const readAtomState = (atom: Atom): AtomState => { + const atomState = ensureAtomState(atom) + // See if we can skip recomputing this atom. + if (isAtomStateInitialized(atomState)) { + // If the atom is mounted, we can use cached atom state. + // because it should have been updated by dependencies. + // We can't use the cache if the atom is invalidated. + if (atomState.m && invalidatedAtoms.get(atom) !== atomState.n) { + return atomState + } + // 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(a).n === n + ) + ) { + return atomState + } + } + // Compute a new state for this atom. + atomState.d.clear() + let isSync = true + const getter: Getter = (a: Atom) => { + if (isSelfAtom(atom, a)) { + const aState = ensureAtomState(a) + if (!isAtomStateInitialized(aState)) { + if (hasInitialValue(a)) { + setAtomStateValueOrPromise(a, aState, a.init) + } else { + // NOTE invalid derived atoms can reach here + throw new Error('no atom init') + } + } + return returnAtomValue(aState) + } + // a !== atom + const aState = readAtomState(a) + try { + return returnAtomValue(aState) + } finally { + addDependency(atom, atomState, a, aState) + if (!isSync) { + mountDependencies(atom, atomState) + flushCallbacks() + } + } + } + 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 (process.env?.MODE !== 'production' && !isActuallyWritableAtom(atom)) { + console.warn('setSelf function cannot be used with read-only atom') + } + if (!setSelf && isActuallyWritableAtom(atom)) { + setSelf = (...args) => { + if (process.env?.MODE !== 'production' && isSync) { + console.warn('setSelf function cannot be called in sync') + } + if (!isSync) { + return writeAtom(atom, ...args) + } + } + } + return setSelf + }, + } + try { + const valueOrPromise = atomRead(atom, getter, options as never) + setAtomStateValueOrPromise(atom, atomState, valueOrPromise) + if (isPromiseLike(valueOrPromise)) { + valueOrPromise.onCancel?.(() => controller?.abort()) + const complete = () => { + if (atomState.m) { + mountDependencies(atom, atomState) + flushCallbacks() + } + } + valueOrPromise.then(complete, complete) + } + return atomState + } catch (error) { + delete atomState.v + atomState.e = error + ++atomState.n + return atomState + } finally { + isSync = false + } + } + + const readAtom = (atom: Atom): Value => returnAtomValue(readAtomState(atom)) + + const getMountedOrPendingDependents = ( + atomState: AtomState + ): Map => { + const dependents = new Map() + for (const a of atomState.m?.t || []) { + const aState = ensureAtomState(a) + if (aState.m) { + dependents.set(a, aState) + } + } + for (const atomWithPendingPromise of atomState.p) { + dependents.set(atomWithPendingPromise, ensureAtomState(atomWithPendingPromise)) + } + return dependents + } + + const invalidateDependents = (atomState: AtomState) => { + const visited = new WeakSet() + const stack: AtomState[] = [atomState] + while (stack.length) { + const aState = stack.pop()! + if (!visited.has(aState)) { + visited.add(aState) + for (const [d, s] of getMountedOrPendingDependents(aState)) { + invalidatedAtoms.set(d, s.n) + stack.push(s) + } + } + } + } + + const recomputeInvalidatedAtoms = () => { + // Step 1: traverse the dependency graph to build the topsorted atom list + // We don't bother to check for cycles, which simplifies the algorithm. + // This is a topological sort via depth-first search, slightly modified from + // what's described here for simplicity and performance reasons: + // https://en.wikipedia.org/wiki/Topological_sorting#Depth-first_search + const topSortedReversed: [atom: AnyAtom, atomState: AtomState, epochNumber: EpochNumber][] = [] + const visiting = new WeakSet() + const visited = new WeakSet() + // Visit the root atom. This is the only atom in the dependency graph + // without incoming edges, which is one reason we can simplify the algorithm + const stack: [a: AnyAtom, aState: AtomState][] = Array.from(changedAtoms) + while (stack.length) { + const [a, aState] = stack[stack.length - 1]! + if (visited.has(a)) { + // All dependents have been processed, now process this atom + stack.pop() + continue + } + if (visiting.has(a)) { + // 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. + if (invalidatedAtoms.get(a) === aState.n) { + topSortedReversed.push([a, aState, aState.n]) + } else { + invalidatedAtoms.delete(a) + changedAtoms.set(a, aState) + } + // Atom has been visited but not yet processed + visited.add(a) + stack.pop() + continue + } + visiting.add(a) + // Push unvisited dependents onto the stack + for (const [d, s] of getMountedOrPendingDependents(aState)) { + if (!visiting.has(d)) { + stack.push([d, s]) + } + } + } + + // Step 2: use the topSortedReversed atom list to recompute all affected atoms + // Track what's changed, so that we can short circuit when possible + for (let i = topSortedReversed.length - 1; i >= 0; --i) { + const [a, aState, prevEpochNumber] = topSortedReversed[i]! + let hasChangedDeps = false + for (const dep of aState.d.keys()) { + if (dep !== a && changedAtoms.has(dep)) { + hasChangedDeps = true + break + } + } + if (hasChangedDeps) { + readAtomState(a) + mountDependencies(a, aState) + if (prevEpochNumber !== aState.n) { + changedAtoms.set(a, aState) + aState.u?.() + } + } + invalidatedAtoms.delete(a) + } + } + + const writeAtomState = ( + atom: WritableAtom, + ...args: Args + ): Result => { + let isSync = true + const getter: Getter = (a: Atom) => returnAtomValue(readAtomState(a)) + const setter: Setter = (a: WritableAtom, ...args: As) => { + const aState = ensureAtomState(a) + try { + if (isSelfAtom(atom, a)) { + if (!hasInitialValue(a)) { + // NOTE technically possible but restricted as it may cause bugs + throw new Error('atom not writable') + } + const prevEpochNumber = aState.n + const v = args[0] as V + setAtomStateValueOrPromise(a, aState, v) + mountDependencies(a, aState) + if (prevEpochNumber !== aState.n) { + changedAtoms.set(a, aState) + aState.u?.() + invalidateDependents(aState) + } + return undefined as R + } else { + return writeAtomState(a, ...args) + } + } finally { + if (!isSync) { + recomputeInvalidatedAtoms() + flushCallbacks() + } + } + } + try { + return atomWrite(atom, getter, setter, ...args) + } finally { + isSync = false + } + } + + const writeAtom = ( + atom: WritableAtom, + ...args: Args + ): Result => { + try { + return writeAtomState(atom, ...args) + } finally { + recomputeInvalidatedAtoms() + flushCallbacks() + } + } + + const mountDependencies = (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(a, ensureAtomState(a)) + aMounted.t.add(atom) + atomState.m.d.add(a) + } + } + for (const a of atomState.m.d || []) { + if (!atomState.d.has(a)) { + atomState.m.d.delete(a) + const aMounted = unmountAtom(a, ensureAtomState(a)) + aMounted?.t.delete(atom) + } + } + } + } + + const mountAtom = (atom: Atom, atomState: AtomState): Mounted => { + if (!atomState.m) { + // recompute atom state + readAtomState(atom) + // mount dependencies first + for (const a of atomState.d.keys()) { + const aMounted = mountAtom(a, ensureAtomState(a)) + aMounted.t.add(atom) + } + // mount self + atomState.m = { + l: new Set(), + d: new Set(atomState.d.keys()), + t: new Set(), + } + atomState.h?.() + if (isActuallyWritableAtom(atom)) { + const mounted = atomState.m + let setAtom: (...args: unknown[]) => unknown + const createInvocationContext = (fn: () => T) => { + let isSync = true + setAtom = (...args: unknown[]) => { + try { + return writeAtomState(atom, ...args) + } finally { + if (!isSync) { + recomputeInvalidatedAtoms() + flushCallbacks() + } + } + } + try { + return fn() + } finally { + isSync = false + } + } + const processOnMount = () => { + const onUnmount = createInvocationContext(() => + atomOnMount(atom, (...args) => setAtom(...args)) + ) + if (onUnmount) { + mounted.u = () => createInvocationContext(onUnmount) + } + } + mountCallbacks.add(processOnMount) + } + } + return atomState.m + } + + const unmountAtom = ( + atom: Atom, + atomState: AtomState + ): Mounted | undefined => { + if ( + atomState.m && + !atomState.m.l.size && + !Array.from(atomState.m.t).some((a) => ensureAtomState(a).m?.d.has(atom)) + ) { + // unmount self + const onUnmount = atomState.m.u + if (onUnmount) { + unmountCallbacks.add(onUnmount) + } + delete atomState.m + atomState.h?.() + // unmount dependencies + for (const a of atomState.d.keys()) { + const aMounted = unmountAtom(a, ensureAtomState(a)) + aMounted?.t.delete(atom) + } + return undefined + } + return atomState.m + } + + const subscribeAtom = (atom: AnyAtom, listener: () => void) => { + const atomState = ensureAtomState(atom) + const mounted = mountAtom(atom, atomState) + const listeners = mounted.l + listeners.add(listener) + flushCallbacks() + return () => { + listeners.delete(listener) + unmountAtom(atom, atomState) + flushCallbacks() + } + } + + const unstable_derive: Store['unstable_derive'] = (fn) => buildStore(...fn(...storeArgs)) + + const store: Store = { + get: readAtom, + set: writeAtom, + sub: subscribeAtom, + unstable_derive, + } + return store +} + +const deriveDevStoreRev4 = (store: Store): Store & DevStoreRev4 => { + const debugMountedAtoms = new Set() + let savedGetAtomState: StoreArgs[0] + let inRestoreAtom = 0 + const derivedStore = store.unstable_derive((...storeArgs: [...StoreArgs]) => { + const [getAtomState, setAtomState, , atomWrite] = storeArgs + savedGetAtomState = getAtomState + storeArgs[1] = function devSetAtomState(atom, atomState) { + setAtomState(atom, atomState) + const originalMounted = atomState.h + atomState.h = () => { + originalMounted?.() + if (atomState.m) { + debugMountedAtoms.add(atom) + } else { + debugMountedAtoms.delete(atom) + } + } + } + storeArgs[3] = function devAtomWrite(atom, getter, setter, ...args) { + if (inRestoreAtom) { + return setter(atom, ...args) + } + return atomWrite(atom, getter, setter, ...args) + } + return storeArgs + }) + const savedStoreSet = derivedStore.set + const devStore: DevStoreRev4 = { + // store dev methods (these are tentative and subject to change without notice) + dev4_get_internal_weak_map: () => ({ + get: (atom) => { + const atomState = savedGetAtomState(atom) + if (!atomState || atomState.n === 0) { + // for backward compatibility + return undefined + } + return atomState + }, + }), + dev4_get_mounted_atoms: () => debugMountedAtoms, + dev4_restore_atoms: (values) => { + const restoreAtom: WritableAtom = { + read: () => null, + write: (_get, set) => { + ++inRestoreAtom + try { + for (const [atom, value] of values) { + if (hasInitialValue(atom)) { + set(atom as never, value) + } + } + } finally { + --inRestoreAtom + } + }, + } + savedStoreSet(restoreAtom) + }, + } + return Object.assign(derivedStore, devStore) +} + +type PrdOrDevStore = Store | (Store & DevStoreRev4) + +export const createStore = (): PrdOrDevStore => { + const atomStateMap = new WeakMap() + const store = buildStore( + (atom) => atomStateMap.get(atom), + (atom, atomState) => atomStateMap.set(atom, atomState).get(atom), + (atom, ...params) => atom.read(...params), + (atom, ...params) => atom.write(...params), + // eslint-disable-next-line @typescript-eslint/ban-ts-comment + // @ts-ignore + (atom, ...params) => atom.unstable_onInit?.(...params), + (atom, ...params) => atom.onMount?.(...params) + ) + if (process.env?.MODE !== 'production') { + return deriveDevStoreRev4(store) + } + return store +} + +let defaultStore: PrdOrDevStore | undefined + +export const getDefaultStore = (): PrdOrDevStore => { + if (!defaultStore) { + defaultStore = createStore() + if (process.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/package.json b/package.json index 76c0028..b3df44b 100644 --- a/package.json +++ b/package.json @@ -76,7 +76,7 @@ "html-webpack-plugin": "^5.5.3", "jest": "^29.7.0", "jest-environment-jsdom": "^29.7.0", - "jotai": "https://pkg.csb.dev/pmndrs/jotai/commit/3b435d79/jotai", + "jotai": "https://pkg.csb.dev/pmndrs/jotai/commit/2cf0a88f/jotai", "microbundle": "^0.15.1", "npm-run-all": "^4.1.5", "prettier": "^3.0.3", diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index fc21d74..dfe95fb 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -63,8 +63,8 @@ importers: specifier: ^29.7.0 version: 29.7.0 jotai: - specifier: https://pkg.csb.dev/pmndrs/jotai/commit/3b435d79/jotai - version: https://pkg.csb.dev/pmndrs/jotai/commit/3b435d79/jotai(@types/react@18.3.18)(react@18.3.1) + specifier: https://pkg.csb.dev/pmndrs/jotai/commit/2cf0a88f/jotai + version: https://pkg.csb.dev/pmndrs/jotai/commit/2cf0a88f/jotai(@types/react@18.3.18)(react@18.3.1) microbundle: specifier: ^0.15.1 version: 0.15.1(@types/babel__core@7.20.5) @@ -2948,8 +2948,8 @@ packages: node-notifier: optional: true - jotai@https://pkg.csb.dev/pmndrs/jotai/commit/3b435d79/jotai: - resolution: {tarball: https://pkg.csb.dev/pmndrs/jotai/commit/3b435d79/jotai} + jotai@https://pkg.csb.dev/pmndrs/jotai/commit/2cf0a88f/jotai: + resolution: {tarball: https://pkg.csb.dev/pmndrs/jotai/commit/2cf0a88f/jotai} version: 2.11.0 engines: {node: '>=12.20.0'} peerDependencies: @@ -8374,7 +8374,7 @@ snapshots: - supports-color - ts-node - jotai@https://pkg.csb.dev/pmndrs/jotai/commit/3b435d79/jotai(@types/react@18.3.18)(react@18.3.1): + jotai@https://pkg.csb.dev/pmndrs/jotai/commit/2cf0a88f/jotai(@types/react@18.3.18)(react@18.3.1): optionalDependencies: '@types/react': 18.3.18 react: 18.3.1 diff --git a/src/atomEffect.ts b/src/atomEffect.ts index 109e2b5..0dc1dca 100644 --- a/src/atomEffect.ts +++ b/src/atomEffect.ts @@ -1,14 +1,19 @@ -import type { Atom, Getter, Setter, createStore } from 'jotai/vanilla' +import type { Atom, Getter, Setter } from 'jotai/vanilla' import { atom } from 'jotai/vanilla' -type Store = ReturnType +const INTERNAL_flushStoreHook = Symbol.for('JOTAI.EXPERIMENTAL.FLUSHSTOREHOOK') +const INTERNAL_syncEffectChannel = Symbol.for('JOTAI-EFFECT.EXPERIMENTAL.SYNCEFFECTCHANNEL') + +type Store = Parameters>[0] +type StoreWithHooks = Store & { + [INTERNAL_flushStoreHook]: () => void + [INTERNAL_syncEffectChannel]?: Set<() => void> +} type GetAtomState = Parameters[0]>[0] type AtomState = NonNullable> -type Batch = Parameters>[0] - type AnyAtom = Atom type GetterWithPeek = Getter & { peek: Getter } @@ -21,102 +26,116 @@ export type Effect = (get: GetterWithPeek, set: SetterWithRecurse) => void | Cle type Ref = { /** epoch */ - x: number + epoch: number /** in progress */ - i: number - /** recursing */ - rc: number - /** refreshing */ - rf?: boolean - /** mounted */ - m?: boolean - /** from cleanup */ - fc?: boolean - /** getter */ - g?: Getter - /** cleanup */ - c?: Cleanup | void + inProgress: number /** pending error */ - e?: unknown + error?: unknown + /** getter */ + get?: Getter } export function atomEffect(effect: Effect) { const refreshAtom = atom(0) - const refAtom = atom(() => ({ i: 0, x: 0, rc: 0 } as Ref)) + const refAtom = atom(() => ({ inProgress: 0, epoch: 0 }) as Ref) - const internalAtom = atom( - (get) => { - get(refreshAtom) - const ref = get(refAtom) - throwPendingError(ref) - ref.g = get - return ++ref.x + const internalAtom = atom((get) => { + get(refreshAtom) + const ref = get(refAtom) + throwPendingError(ref) + if (ref.inProgress) { + return ref.epoch } - ) + ref.get = get + return ++ref.epoch + }) internalAtom.unstable_onInit = (store) => { const ref = store.get(refAtom) + let runCleanup: Cleanup | void + let isMounted = false + let isRecursing = false + let isRefreshing = false + let fromCleanup = false + function runEffect() { - if ( - !ref.m || - ref.rc || - ref.i && !ref.rf - ) { + if (!isMounted || (ref.inProgress && !isRefreshing) || isRecursing) { return } const deps = new Map() - const getter = ((a) => { - const value = ref.g!(a) + const getter: GetterWithPeek = ((a) => { + const value = ref.get!(a) deps.set(a, value) return value }) as GetterWithPeek + getter.peek = store.get - const setter = ((a, ...args) => { + const setter: SetterWithRecurse = ((a, ...args) => { try { - ++ref.i + ++ref.inProgress return store.set(a, ...args) } finally { - deps.keys().forEach(ref.g!) // TODO - do we still need this? - --ref.i + --ref.inProgress } }) as SetterWithRecurse setter.recurse = (a, ...args) => { - if (ref.fc) { + if (fromCleanup) { if (process.env.NODE_ENV !== 'production') { throw new Error('set.recurse is not allowed in cleanup') } - return void 0 as any + return undefined as any } try { - ++ref.rc + isRecursing = true return store.set(a, ...args) } finally { - try { - const depsChanged = Array.from(deps).some(areDifferent) - if (depsChanged) { - refresh() - } - } finally { - --ref.rc + isRecursing = false + const depsChanged = Array.from(deps).some(areDifferent) + if (depsChanged) { + refresh() } } } try { - ++ref.i - cleanup() - ref.c = effectAtom.effect(getter, setter) + ++ref.inProgress + runCleanup?.() + const cleanup = effectAtom.effect(getter, setter) + if (typeof cleanup === 'function') { + runCleanup = () => { + try { + fromCleanup = true + return cleanup() + } catch (error) { + ref.error = error + refresh() + } finally { + fromCleanup = false + runCleanup = undefined + } + } + } } catch (e) { - ref.e = e + ref.error = e refresh() } finally { - --ref.i + Array.from(deps.keys(), ref.get!) + --ref.inProgress + } + + function refresh() { + try { + isRefreshing = true + store.set(refreshAtom, (v) => v + 1) + } finally { + isRefreshing = false + } } function areDifferent([a, v]: [Atom, unknown]) { @@ -125,54 +144,26 @@ export function atomEffect(effect: Effect) { } const atomState = getAtomState(store, internalAtom) + const syncEffectChannel = ensureSyncEffectChannel(store) - const originalMountHook = atomState.h - atomState.h = (batch) => { - originalMountHook?.(batch) + hookInto(atomState, 'h', function atomOnMount() { + if (ref.inProgress) { + return + } if (atomState.m) { - ref.m = true - scheduleListener(batch, runEffect) + isMounted = true + syncEffectChannel.add(runEffect) } else { - ref.m = false - scheduleListener(batch, cleanup) - } - } - - const originalUpdateHook = atomState.u - atomState.u = (batch) => { - originalUpdateHook?.(batch) - batch[0].add(runEffect) - } - - function scheduleListener(batch: Batch, listener: () => void) { - batch[0].add(listener) - } - - function refresh() { - try { - ref.rf = true - store.set(refreshAtom, (v) => v + 1) - } finally { - ref.rf = false + isMounted = false + if (runCleanup) { + syncEffectChannel.add(runCleanup) + } } - } + }) - function cleanup() { - if (typeof ref.c !== 'function') { - return - } - try { - ref.fc = true - ref.c() - } - catch(e) { - ref.e = e - refresh() - } finally { - ref.fc = false - delete ref.c - } - } + hookInto(atomState, 'u', function atomOnUpdate() { + syncEffectChannel.add(runEffect) + }) } if (process.env.NODE_ENV !== 'production') { @@ -188,27 +179,40 @@ export function atomEffect(effect: Effect) { } const effectAtom = Object.assign( - atom((get) => get(internalAtom)), + atom((get) => void get(internalAtom)), { effect } ) return effectAtom - + function throwPendingError(ref: Ref) { if ('e' in ref) { - const error = ref.e - delete ref.e + const error = ref.error + delete ref.error throw error } } } +function ensureSyncEffectChannel(store: Store) { + const storeWithHooks = store as StoreWithHooks + let syncEffectChannel = storeWithHooks[INTERNAL_syncEffectChannel] + if (!syncEffectChannel) { + storeWithHooks[INTERNAL_syncEffectChannel] = syncEffectChannel = new Set<() => void>() + hookInto(storeWithHooks, INTERNAL_flushStoreHook, () => { + syncEffectChannel!.forEach((fn: () => void) => fn()) + syncEffectChannel!.clear() + }) + } + return syncEffectChannel +} + const getAtomStateMap = new WeakMap() /** - * HACK: steal atomState to synchronously determine if - * the atom is mounted + * HACK: Steals atomState to synchronously determine if + * the atom is mounted. * We return null to cause the buildStore(...args) to throw - * to abort creating a derived store + * to abort creating a derived store. */ function getAtomState(store: Store, atom: AnyAtom): AtomState { let getAtomStateFn = getAtomStateMap.get(store) @@ -225,3 +229,15 @@ function getAtomState(store: Store, atom: AnyAtom): AtomState { } return getAtomStateFn!(atom)! } + +function hookInto void }>( + obj: T, + methodName: M, + newMethod: NonNullable +) { + const originalMethod = obj[methodName] + obj[methodName] = ((...args: any[]) => { + originalMethod?.(...args) + newMethod(...args) + }) as T[typeof methodName] +}