diff --git a/src/types.ts b/src/types.ts index a3a6b99e..58244681 100644 --- a/src/types.ts +++ b/src/types.ts @@ -3,6 +3,7 @@ import type { Atom, WritableAtom, createStore } from 'jotai/vanilla'; import type { INTERNAL_AtomState, INTERNAL_Store, + INTERNAL_getBuildingBlocksRev1, } from 'jotai/vanilla/internals'; export type DevStore = { @@ -17,6 +18,10 @@ export type StoreWithoutDevMethods = ReturnType; export type StoreWithDevMethods = INTERNAL_Store & DevStore; export type Store = StoreWithoutDevMethods | StoreWithDevMethods; +type Mutable = { -readonly [P in keyof T]: T[P] }; +export type BuildingBlocks = Mutable< + ReturnType +>; export type Options = Parameters[0]; diff --git a/src/utils/internals/compose-with-devtools.ts b/src/utils/internals/compose-with-devtools.ts index 05e9e818..87663695 100644 --- a/src/utils/internals/compose-with-devtools.ts +++ b/src/utils/internals/compose-with-devtools.ts @@ -1,13 +1,15 @@ import { Atom, WritableAtom } from 'jotai'; -import { INTERNAL_overrideCreateStore } from 'jotai/vanilla'; +import { INTERNAL_overrideCreateStore, createStore } from 'jotai/vanilla'; import { INTERNAL_buildStoreRev1 as INTERNAL_buildStore, + INTERNAL_getBuildingBlocksRev1, INTERNAL_initializeStoreHooks, } from 'jotai/vanilla/internals'; import { AnyAtom, AnyAtomError, AnyAtomValue, + BuildingBlocks, DevStore, Store, StoreWithDevMethods, @@ -156,27 +158,26 @@ const __composeWithDevTools = ( return store as typeof store & DevToolsStoreMethods; }; -const createDevStore = (): StoreWithDevMethods => { +const createDevStore = ( + prevCreateStore: (() => Store) | undefined, +): StoreWithDevMethods => { let inRestoreAtom = 0; - const storeHooks = INTERNAL_initializeStoreHooks({}); - const atomStateMap = new WeakMap(); - const mountedAtoms = new WeakMap(); - const store = INTERNAL_buildStore( - atomStateMap, - mountedAtoms, - undefined, - undefined, - undefined, - undefined, - storeHooks, - undefined, - (atom, get, set, ...args) => { - if (inRestoreAtom) { - return set(atom, ...(args as any)); - } - return atom.write(get, set, ...(args as any)); - }, - ); + const prevStore = prevCreateStore?.() ?? INTERNAL_buildStore(); + const buildingBlocks = [ + ...INTERNAL_getBuildingBlocksRev1(prevStore), + ] as BuildingBlocks; + const storeHooks = INTERNAL_initializeStoreHooks(buildingBlocks[6]); + const atomWrite = buildingBlocks[8]; + buildingBlocks[8] = (atom, get, set, ...args) => { + if (inRestoreAtom) { + return set(atom, ...args); + } + return atomWrite + ? atomWrite(atom, get, set, ...args) + : atom.write(get, set, ...args); + }; + const store = INTERNAL_buildStore(...buildingBlocks); + const [atomStateMap, mountedAtoms] = buildingBlocks; const debugMountedAtoms = new Set>(); storeHooks.m.add(undefined, (atom) => { debugMountedAtoms.add(atom); @@ -214,10 +215,7 @@ const createDevStore = (): StoreWithDevMethods => { }, }; - return { - ...store, - ...devStore, - }; + return Object.assign(store, devStore); }; const isDevStore = (store: Store): store is StoreWithDevMethods => { @@ -225,7 +223,7 @@ const isDevStore = (store: Store): store is StoreWithDevMethods => { }; INTERNAL_overrideCreateStore((prev) => { - return createDevStore; + return () => createDevStore(prev); }); export const composeWithDevTools = (