diff --git a/src/vanilla/store.ts b/src/vanilla/store.ts index 0a7d7e368a..395ff345c2 100644 --- a/src/vanilla/store.ts +++ b/src/vanilla/store.ts @@ -271,7 +271,7 @@ type DevStoreRev4 = { dev4_restore_atoms: (values: Iterable) => void } -type PrdStore = { +type Store = { get: (atom: Atom) => Value set: ( atom: WritableAtom, @@ -281,21 +281,12 @@ type PrdStore = { unstable_derive: (fn: (...args: StoreArgs) => StoreArgs) => Store } -type Store = PrdStore | (PrdStore & DevStoreRev4) - export type INTERNAL_DevStoreRev4 = DevStoreRev4 -export type INTERNAL_PrdStore = PrdStore +export type INTERNAL_PrdStore = Store const buildStore = ( ...[getAtomState, atomRead, atomWrite, atomOnMount]: StoreArgs ): Store => { - // for debugging purpose only - let debugMountedAtoms: Set - - if (import.meta.env?.MODE !== 'production') { - debugMountedAtoms = new Set() - } - const setAtomStateValueOrPromise = ( atom: AnyAtom, atomState: AtomState, @@ -634,9 +625,6 @@ const buildStore = ( d: new Set(atomState.d.keys()), t: new Set(), } - if (import.meta.env?.MODE !== 'production') { - debugMountedAtoms.add(atom) - } if (isActuallyWritableAtom(atom)) { const mounted = atomState.m let setAtom: (...args: unknown[]) => unknown @@ -686,9 +674,6 @@ const buildStore = ( addBatchFunc(batch, 'L', () => onUnmount(batch)) } delete atomState.m - if (import.meta.env?.MODE !== 'production') { - debugMountedAtoms.delete(atom) - } // unmount dependencies for (const a of atomState.d.keys()) { const aMounted = unmountAtom(batch, a, getAtomState(a)) @@ -723,43 +708,90 @@ const buildStore = ( sub: subscribeAtom, unstable_derive, } - 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: () => ({ - get: (atom) => { - const atomState = getAtomState(atom) - if (atomState.n === 0) { - // for backward compatibility - return undefined + return store +} + +const deriveDevStoreRev4 = (store: Store): Store & DevStoreRev4 => { + const proxyAtomStateMap = new WeakMap() + const debugMountedAtoms = new Set() + let savedGetAtomState: StoreArgs[0] + let inRestoreAtom = 0 + const derivedStore = store.unstable_derive( + (getAtomState, atomRead, atomWrite, atomOnMount) => { + savedGetAtomState = getAtomState + return [ + (atom) => { + let proxyAtomState = proxyAtomStateMap.get(atom) + if (!proxyAtomState) { + const atomState = getAtomState(atom) + proxyAtomState = new Proxy(atomState, { + set(target, prop, value) { + if (prop === 'm') { + debugMountedAtoms.add(atom) + } + return Reflect.set(target, prop, value) + }, + deleteProperty(target, prop) { + if (prop === 'm') { + debugMountedAtoms.delete(atom) + } + return Reflect.deleteProperty(target, prop) + }, + }) + proxyAtomStateMap.set(atom, proxyAtomState) } - return atomState + return proxyAtomState }, - }), - dev4_get_mounted_atoms: () => debugMountedAtoms, - dev4_restore_atoms: (values) => { - const batch = createBatch() - for (const [atom, value] of values) { - if (hasInitialValue(atom)) { - const atomState = getAtomState(atom) - const prevEpochNumber = atomState.n - setAtomStateValueOrPromise(atom, atomState, value) - mountDependencies(batch, atom, atomState) - if (prevEpochNumber !== atomState.n) { - registerBatchAtom(batch, atom, atomState) - recomputeDependents(batch, atom, atomState) - } + atomRead, + (atom, getter, setter, ...args) => { + if (inRestoreAtom) { + return setter(atom, ...args) } + return atomWrite(atom, getter, setter, ...args) + }, + atomOnMount, + ] + }, + ) + 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.n === 0) { + // for backward compatibility + return undefined } - flushBatch(batch) + return atomState }, - } - Object.assign(store, devStore) + }), + 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 store + return Object.assign(derivedStore, devStore) } -export const createStore = (): Store => { +type PrdOrDevStore = Store | (Store & DevStoreRev4) + +export const createStore = (): PrdOrDevStore => { const atomStateMap = new WeakMap() const getAtomState = (atom: Atom) => { if (import.meta.env?.MODE !== 'production' && !atom) { @@ -772,17 +804,21 @@ export const createStore = (): Store => { } return atomState } - return buildStore( + const store = buildStore( getAtomState, (atom, ...params) => atom.read(...params), (atom, ...params) => atom.write(...params), (atom, ...params) => atom.onMount?.(...params), ) + if (import.meta.env?.MODE !== 'production') { + return deriveDevStoreRev4(store) + } + return store } -let defaultStore: Store | undefined +let defaultStore: PrdOrDevStore | undefined -export const getDefaultStore = (): Store => { +export const getDefaultStore = (): PrdOrDevStore => { if (!defaultStore) { defaultStore = createStore() if (import.meta.env?.MODE !== 'production') {