feat: dev store with unstable_derive (#2852)

* wip: dev store with unstable_derive

* hack lint rule

* fix mounted map

* implement restore atoms

* save derivedStore.set

* fix it

* refactor

* rename

src/vanilla/store.ts

@@ -271,7 +271,7 @@ type DevStoreRev4 = {
   dev4_restore_atoms: (values: Iterable<readonly [AnyAtom, AnyValue]>) => void
 }
 
-type PrdStore = {
+type Store = {
   get: <Value>(atom: Atom<Value>) => Value
   set: <Value, Args extends unknown[], Result>(
     atom: WritableAtom<Value, Args, Result>,
@@ -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<AnyAtom>
-
-  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<AnyAtom>()
+  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<null, [], void> = {
+        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 = <Value>(atom: Atom<Value>) => {
     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') {