refactor(core,utils,devtools): add dev methods in store (#717)

* refactor(core,utils,devtools): add dev methods in store

* fix: state listener should be invoked in flush

* simplify debug atom init

Author: dai-shi

src/core/Provider.ts

@@ -7,14 +7,14 @@ import {
 } from 'react'
 import type { PropsWithChildren } from 'react'
 import type { Atom, Scope } from './atom'
+import { createScopeContainer, getScopeContext } from './contexts'
+import type { ScopeContainer } from './contexts'
 import {
-  ScopeContainer,
-  createScopeContainer,
-  getScopeContext,
-  isDevScopeContainer,
-} from './contexts'
-import type { ScopeContainerForDevelopment } from './contexts'
-import { DEV_GET_ATOM_STATE, DEV_GET_MOUNTED } from './store'
+  DEV_GET_ATOM_STATE,
+  DEV_GET_MOUNTED,
+  DEV_GET_MOUNTED_ATOMS,
+  DEV_SUBSCRIBE_STATE,
+} from './store'
 import type { AtomState, Store } from './store'
 
 export const Provider = ({
@@ -34,8 +34,7 @@ export const Provider = ({
   if (
     typeof process === 'object' &&
     process.env.NODE_ENV !== 'production' &&
-    process.env.NODE_ENV !== 'test' &&
-    isDevScopeContainer(scopeContainerRef.current)
+    process.env.NODE_ENV !== 'test'
   ) {
     // eslint-disable-next-line react-hooks/rules-of-hooks
     useDebugState(scopeContainerRef.current)
@@ -77,15 +76,16 @@ const stateToPrintable = ([store, atoms]: [Store, Atom<unknown>[]]) =>
 
 // We keep a reference to the atoms in Provider's registeredAtoms in dev mode,
 // so atoms aren't garbage collected by the WeakMap of mounted atoms
-const useDebugState = (scopeContainer: ScopeContainerForDevelopment) => {
-  const [store, devStore] = scopeContainer
-  const [atoms, setAtoms] = useState(devStore.atoms)
+const useDebugState = (scopeContainer: ScopeContainer) => {
+  const store = scopeContainer.s
+  const [atoms, setAtoms] = useState<Atom<unknown>[]>([])
   useEffect(() => {
-    // HACK creating a new reference for useDebugValue to update
-    const callback = () => setAtoms([...devStore.atoms])
-    const unsubscribe = devStore.subscribe(callback)
+    const callback = () => {
+      setAtoms(Array.from(store[DEV_GET_MOUNTED_ATOMS]?.() || []))
+    }
+    const unsubscribe = store[DEV_SUBSCRIBE_STATE]?.(callback)
     callback()
     return unsubscribe
-  }, [devStore])
+  }, [store])
   useDebugValue([store, atoms], stateToPrintable)
 }

src/core/contexts.ts

@@ -2,66 +2,19 @@ import { createContext } from 'react'
 import type { Context } from 'react'
 import type { Atom, Scope } from './atom'
 import { createStore } from './store'
+import type { Store } from './store'
 
-const createScopeContainerForProduction = (
-  initialValues?: Iterable<readonly [Atom<unknown>, unknown]>
-) => {
-  const store = createStore(initialValues)
-  return [store] as const
+export type ScopeContainer = {
+  s: Store
 }
 
-const createScopeContainerForDevelopment = (
+export const createScopeContainer = (
   initialValues?: Iterable<readonly [Atom<unknown>, unknown]>
-) => {
-  const devStore = {
-    listeners: new Set<() => void>(),
-    subscribe: (callback: () => void) => {
-      devStore.listeners.add(callback)
-      return () => {
-        devStore.listeners.delete(callback)
-      }
-    },
-    atoms: Array.from(initialValues ?? []).map(([a]) => a),
-  }
-  const stateListener = (updatedAtom: Atom<unknown>, isNewAtom: boolean) => {
-    if (isNewAtom) {
-      // FIXME memory leak
-      // we should probably remove unmounted atoms eventually
-      devStore.atoms = [...devStore.atoms, updatedAtom]
-    }
-    Promise.resolve().then(() => {
-      devStore.listeners.forEach((listener) => listener())
-    })
-  }
-  const store = createStore(initialValues, stateListener)
-  return [store, devStore] as const
-}
-
-export const isDevScopeContainer = (
-  scopeContainer: ScopeContainer
-): scopeContainer is ScopeContainerForDevelopment => {
-  return scopeContainer.length > 1
+): ScopeContainer => {
+  const store = createStore(initialValues)
+  return { s: store }
 }
 
-type ScopeContainerForProduction = ReturnType<
-  typeof createScopeContainerForProduction
->
-export type ScopeContainerForDevelopment = ReturnType<
-  typeof createScopeContainerForDevelopment
->
-export type ScopeContainer =
-  | ScopeContainerForProduction
-  | ScopeContainerForDevelopment
-
-type CreateScopeContainer = (
-  initialValues?: Iterable<readonly [Atom<unknown>, unknown]>
-) => ScopeContainer
-
-export const createScopeContainer: CreateScopeContainer =
-  typeof process === 'object' && process.env.NODE_ENV !== 'production'
-    ? createScopeContainerForDevelopment
-    : createScopeContainerForProduction
-
 type ScopeContext = Context<ScopeContainer>
 
 const ScopeContextMap = new Map<Scope | undefined, ScopeContext>()

src/core/store.ts

@@ -64,23 +64,36 @@ type Mounted = {
 
 // for debugging purpose only
 type StateListener = (updatedAtom: AnyAtom, isNewAtom: boolean) => void
+type MountedAtoms = Set<AnyAtom>
 
 // store methods
 export const READ_ATOM = 'r'
 export const WRITE_ATOM = 'w'
 export const FLUSH_PENDING = 'f'
 export const SUBSCRIBE_ATOM = 's'
 export const RESTORE_ATOMS = 'h'
+
+// store dev methods (these are tentative and subject to change)
+export const DEV_SUBSCRIBE_STATE = 'n'
+export const DEV_GET_MOUNTED_ATOMS = 'l'
 export const DEV_GET_ATOM_STATE = 'a'
 export const DEV_GET_MOUNTED = 'm'
 
 export const createStore = (
-  initialValues?: Iterable<readonly [AnyAtom, unknown]>,
-  stateListener?: StateListener
+  initialValues?: Iterable<readonly [AnyAtom, unknown]>
 ) => {
   const atomStateMap = new WeakMap<AnyAtom, AtomState>()
   const mountedMap = new WeakMap<AnyAtom, Mounted>()
-  const pendingMap = new Map<AnyAtom, ReadDependencies | undefined>()
+  const pendingMap = new Map<
+    AnyAtom,
+    [dependencies: ReadDependencies | undefined, isNewAtom: boolean]
+  >()
+  let stateListeners: Set<StateListener>
+  let mountedAtoms: MountedAtoms
+  if (typeof process === 'object' && process.env.NODE_ENV !== 'production') {
+    stateListeners = new Set()
+    mountedAtoms = new Set()
+  }
 
   if (initialValues) {
     for (const [atom, value] of initialValues) {
@@ -144,7 +157,7 @@ export const createStore = (
         atomState.d.set(atom, atomState.r)
       }
     }
-    commitAtomState(atom, atomState, dependencies && prevDependencies)
+    setAtomState(atom, atomState, dependencies && prevDependencies)
   }
 
   const setAtomReadError = <Value>(
@@ -163,7 +176,7 @@ export const createStore = (
     delete atomState.c // cancel read promise
     delete atomState.i // invalidated revision
     atomState.e = error // read error
-    commitAtomState(atom, atomState, prevDependencies)
+    setAtomState(atom, atomState, prevDependencies)
   }
 
   const setAtomReadPromise = <Value>(
@@ -186,13 +199,13 @@ export const createStore = (
       atomState.p = interruptablePromise // read promise
       atomState.c = interruptablePromise[INTERRUPT_PROMISE]
     }
-    commitAtomState(atom, atomState, prevDependencies)
+    setAtomState(atom, atomState, prevDependencies)
   }
 
   const setAtomInvalidated = <Value>(atom: Atom<Value>): void => {
     const [atomState] = wipAtomState(atom)
     atomState.i = atomState.r // invalidated revision
-    commitAtomState(atom, atomState)
+    setAtomState(atom, atomState)
   }
 
   const setAtomWritePromise = <Value>(
@@ -207,7 +220,7 @@ export const createStore = (
       // delete it only if it's not overwritten
       delete atomState.w // write promise
     }
-    commitAtomState(atom, atomState)
+    setAtomState(atom, atomState)
   }
 
   const scheduleReadAtomState = <Value>(
@@ -483,6 +496,9 @@ export const createStore = (
       u: undefined,
     }
     mountedMap.set(atom, mounted)
+    if (typeof process === 'object' && process.env.NODE_ENV !== 'production') {
+      mountedAtoms.add(atom)
+    }
     if (isActuallyWritableAtom(atom) && atom.onMount) {
       const setAtom = (update: unknown) => writeAtom(atom, update)
       mounted.u = atom.onMount(setAtom)
@@ -497,6 +513,9 @@ export const createStore = (
       onUnmount()
     }
     mountedMap.delete(atom)
+    if (typeof process === 'object' && process.env.NODE_ENV !== 'production') {
+      mountedAtoms.delete(atom)
+    }
     // unmount read dependencies afterward
     const atomState = getAtomState(atom)
     if (atomState) {
@@ -550,7 +569,7 @@ export const createStore = (
     })
   }
 
-  const commitAtomState = <Value>(
+  const setAtomState = <Value>(
     atom: Atom<Value>,
     atomState: AtomState<Value>,
     prevDependencies?: ReadDependencies
@@ -560,31 +579,29 @@ export const createStore = (
     }
     const isNewAtom = !atomStateMap.has(atom)
     atomStateMap.set(atom, atomState)
-    if (stateListener) {
-      stateListener(atom, isNewAtom)
-    }
     if (!pendingMap.has(atom)) {
-      pendingMap.set(atom, prevDependencies)
+      pendingMap.set(atom, [prevDependencies, isNewAtom])
     }
   }
 
   const flushPending = (): void => {
     const pending = Array.from(pendingMap)
     pendingMap.clear()
-    pending.forEach(([atom, prevDependencies]) => {
-      const atomState = getAtomState(atom)
-      if (atomState) {
-        if (prevDependencies) {
+    pending.forEach(([atom, [prevDependencies, isNewAtom]]) => {
+      if (prevDependencies) {
+        const atomState = getAtomState(atom)
+        if (atomState) {
           mountDependencies(atom, atomState, prevDependencies)
         }
-      } else if (
+      }
+      const mounted = mountedMap.get(atom)
+      mounted?.l.forEach((listener) => listener())
+      if (
         typeof process === 'object' &&
         process.env.NODE_ENV !== 'production'
       ) {
-        console.warn('[Bug] atom state not found in flush', atom)
+        stateListeners.forEach((l) => l(atom, isNewAtom))
       }
-      const mounted = mountedMap.get(atom)
-      mounted?.l.forEach((listener) => listener())
     })
   }
 
@@ -617,6 +634,13 @@ export const createStore = (
       [FLUSH_PENDING]: flushPending,
       [SUBSCRIBE_ATOM]: subscribeAtom,
       [RESTORE_ATOMS]: restoreAtoms,
+      [DEV_SUBSCRIBE_STATE]: (l: StateListener) => {
+        stateListeners.add(l)
+        return () => {
+          stateListeners.delete(l)
+        }
+      },
+      [DEV_GET_MOUNTED_ATOMS]: () => mountedAtoms.values(),
       [DEV_GET_ATOM_STATE]: (a: AnyAtom) => atomStateMap.get(a),
       [DEV_GET_MOUNTED]: (a: AnyAtom) => mountedMap.get(a),
     }

src/core/useAtom.ts

@@ -53,7 +53,7 @@ export function useAtom<Value, Update>(
   }
 
   const ScopeContext = getScopeContext(scope)
-  const [store] = useContext(ScopeContext)
+  const store = useContext(ScopeContext).s
 
   const getAtomValue = useCallback(() => {
     const atomState = store[READ_ATOM](atom)

src/devtools/useAtomsSnapshot.ts

@@ -1,40 +1,51 @@
 import { useContext, useEffect, useState } from 'react'
 import { SECRET_INTERNAL_getScopeContext as getScopeContext } from 'jotai'
 import type { Atom, Scope } from '../core/atom'
-// NOTE importing from '../core/contexts' is across bundles and actually copying code
-import { isDevScopeContainer } from '../core/contexts'
-import { DEV_GET_ATOM_STATE, DEV_GET_MOUNTED } from '../core/store'
-import type { AtomState } from '../core/store'
+import {
+  DEV_GET_ATOM_STATE,
+  DEV_GET_MOUNTED_ATOMS,
+  DEV_SUBSCRIBE_STATE,
+} from '../core/store'
+import type { AtomState, Store } from '../core/store'
 
 type AtomsSnapshot = Map<Atom<unknown>, unknown>
 
+const createAtomsSnapshot = (
+  store: Store,
+  atoms: Atom<unknown>[]
+): AtomsSnapshot => {
+  const tuples = atoms.map<[Atom<unknown>, unknown]>((atom) => {
+    const atomState = store[DEV_GET_ATOM_STATE]?.(atom) ?? ({} as AtomState)
+    return [atom, atomState.v]
+  })
+  return new Map(tuples)
+}
+
 export function useAtomsSnapshot(scope?: Scope): AtomsSnapshot {
   const ScopeContext = getScopeContext(scope)
   const scopeContainer = useContext(ScopeContext)
+  const store = scopeContainer.s
 
-  if (!isDevScopeContainer(scopeContainer)) {
-    throw Error('useAtomsSnapshot can only be used in dev mode.')
+  if (!store[DEV_SUBSCRIBE_STATE]) {
+    throw new Error('useAtomsSnapshot can only be used in dev mode.')
   }
 
-  const [store, devStore] = scopeContainer
+  const [atomsSnapshot, setAtomsSnapshot] = useState<AtomsSnapshot>(
+    () => new Map()
+  )
 
-  const [atomsSnapshot, setAtomsSnapshot] = useState<AtomsSnapshot>(new Map())
   useEffect(() => {
-    const callback = () => {
-      const { atoms } = devStore
-      const atomToAtomValueTuples = atoms
-        .filter((atom) => !!store[DEV_GET_MOUNTED]?.(atom))
-        .map<[Atom<unknown>, unknown]>((atom) => {
-          const atomState =
-            store[DEV_GET_ATOM_STATE]?.(atom) ?? ({} as AtomState)
-          return [atom, atomState.v]
-        })
-      setAtomsSnapshot(new Map(atomToAtomValueTuples))
+    const callback = (updatedAtom?: Atom<unknown>, isNewAtom?: boolean) => {
+      const atoms = Array.from(store[DEV_GET_MOUNTED_ATOMS]?.() || [])
+      if (updatedAtom && isNewAtom && !atoms.includes(updatedAtom)) {
+        atoms.push(updatedAtom)
+      }
+      setAtomsSnapshot(createAtomsSnapshot(store, atoms))
     }
-    const unsubscribe = devStore.subscribe(callback)
+    const unsubscribe = store[DEV_SUBSCRIBE_STATE]?.(callback)
     callback()
     return unsubscribe
-  }, [store, devStore])
+  }, [store])
 
   return atomsSnapshot
 }