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

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
 }