Snapshot-based devtools (1st iteration) (#425)

* snapshot wip

* useAtomsSnapshot wip

* Add type

* Add first test

* Return map from useAtomsSnapshot

* wip snapshots

* Change registeredAtoms to state
Add tests

* Add key

* Update src/core/contexts.ts

* Fixes after review

* Merges

* Fix test

* add FIXME comments

Co-authored-by: Daishi Kato <[email protected]>
Co-authored-by: daishi <[email protected]>

.size-snapshot.json

@@ -1,15 +1,15 @@
 {
   "devtools.js": {
-    "bundled": 1989,
-    "minified": 1051,
-    "gzipped": 594,
+    "bundled": 19585,
+    "minified": 9255,
+    "gzipped": 3131,
     "treeshaked": {
       "rollup": {
-        "code": 28,
-        "import_statements": 28
+        "code": 58,
+        "import_statements": 52
       },
       "webpack": {
-        "code": 1045
+        "code": 1357
       }
     }
   },
@@ -84,9 +84,9 @@
     }
   },
   "query.js": {
-    "bundled": 2957,
-    "minified": 1267,
-    "gzipped": 612,
+    "bundled": 2359,
+    "minified": 1046,
+    "gzipped": 518,
     "treeshaked": {
       "rollup": {
         "code": 57,
@@ -126,16 +126,16 @@
     }
   },
   "index.js": {
-    "bundled": 21044,
-    "minified": 9971,
-    "gzipped": 3185,
+    "bundled": 21378,
+    "minified": 10119,
+    "gzipped": 3243,
     "treeshaked": {
       "rollup": {
-        "code": 14,
-        "import_statements": 14
+        "code": 44,
+        "import_statements": 38
       },
       "webpack": {
-        "code": 1284
+        "code": 1318
       }
     }
   }

src/core/Provider.ts

@@ -1,30 +1,47 @@
-import React, { createElement, useCallback, useRef, useDebugValue } from 'react'
+import React, {
+  createElement,
+  useCallback,
+  useRef,
+  useDebugValue,
+  useState,
+} from 'react'
 
 import type { AnyAtom, Scope } from './types'
 import { subscribeAtom } from './vanilla'
 import type { AtomState, State } from './vanilla'
-import { createStore, getStoreContext } from './contexts'
+import {
+  createStore,
+  getStoreContext,
+  RegisteredAtomsContext,
+} from './contexts'
 import type { Store } from './contexts'
 import { useMutableSource } from './useMutableSource'
 
 export const Provider: React.FC<{
   initialValues?: Iterable<readonly [AnyAtom, unknown]>
   scope?: Scope
-}> = ({ initialValues, scope, children }) => {
+}> = ({ initialValues, scope, children: baseChildren }) => {
   const storeRef = useRef<ReturnType<typeof createStore> | null>(null)
+  let children = baseChildren
 
   if (typeof process === 'object' && process.env.NODE_ENV !== 'production') {
     /* eslint-disable react-hooks/rules-of-hooks */
-    const atomsRef = useRef<AnyAtom[]>([])
+    const [registeredAtoms, setRegisteredAtoms] = useState<AnyAtom[]>([])
     if (storeRef.current === null) {
       // lazy initialization
       storeRef.current = createStore(initialValues, (newAtom) => {
-        atomsRef.current.push(newAtom)
+        // FIXME find a proper way to handle registered atoms
+        setTimeout(() => setRegisteredAtoms((atoms) => [...atoms, newAtom]), 0)
       })
     }
     useDebugState(
       storeRef.current as ReturnType<typeof createStore>,
-      atomsRef.current
+      registeredAtoms
+    )
+    children = createElement(
+      RegisteredAtomsContext.Provider,
+      { value: registeredAtoms },
+      baseChildren
     )
     /* eslint-enable react-hooks/rules-of-hooks */
   } else {
@@ -67,7 +84,7 @@ const stateToPrintable = ([state, atoms]: [State, AnyAtom[]]) =>
 
 const getState = (state: State) => ({ ...state }) // shallow copy
 
-// We keep a reference to the atoms in Provider's atomsRef in dev mode,
+// 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 = (store: Store, atoms: AnyAtom[]) => {
   const subscribe = useCallback(

src/core/contexts.ts

@@ -39,3 +39,6 @@ export const getStoreContext = (scope?: Scope) => {
   }
   return StoreContextMap.get(scope) as StoreContext
 }
+
+// only for DEV purpose
+export const RegisteredAtomsContext = createContext<AnyAtom[]>([])

src/devtools.ts

@@ -1 +1,2 @@
 export { useAtomDevtools } from './devtools/useAtomDevtools'
+export { useAtomsSnapshot } from './devtools/useAtomsSnapshot'

src/devtools/useAtomsSnapshot.ts

@@ -0,0 +1,38 @@
+import { useCallback, useContext, useMemo } from 'react'
+import { SECRET_INTERNAL_getStoreContext as getStoreContext } from 'jotai'
+import { AtomState, State, subscribeAtom } from '../core/vanilla'
+import { RegisteredAtomsContext } from '../core/contexts'
+import { useMutableSource } from '../core/useMutableSource'
+import { AnyAtom } from '../core/types'
+
+type AtomsSnapshot = Map<AnyAtom, unknown>
+
+export function useAtomsSnapshot(): AtomsSnapshot {
+  const StoreContext = getStoreContext()
+  const [mutableSource] = useContext(StoreContext)
+  const atoms = useContext(RegisteredAtomsContext)
+
+  const subscribe = useCallback(
+    (state: State, callback: () => void) => {
+      // FIXME we don't need to resubscribe, just need to subscribe for new one
+      const unsubs = atoms.map((atom) => subscribeAtom(state, atom, callback))
+      return () => {
+        unsubs.forEach((unsub) => unsub())
+      }
+    },
+    [atoms]
+  )
+  const state: State = useMutableSource(mutableSource, getState, subscribe)
+
+  return useMemo(() => {
+    const atomToAtomValueTuples = atoms
+      .filter((atom) => !!state.m.get(atom))
+      .map<[AnyAtom, unknown]>((atom) => {
+        const atomState = state.a.get(atom) ?? ({} as AtomState)
+        return [atom, atomState.e || atomState.p || atomState.w || atomState.v]
+      })
+    return new Map(atomToAtomValueTuples)
+  }, [atoms, state])
+}
+
+const getState = (state: State) => ({ ...state }) // shallow copy

tests/devtools/useAtomsSnapshot.test.tsx

@@ -0,0 +1,86 @@
+import React, { useState } from 'react'
+import { fireEvent, render } from '@testing-library/react'
+import { Provider, atom, useAtom } from '../../src/index'
+import { useAtomsSnapshot } from '../../src/devtools'
+
+beforeEach(() => {
+  process.env.NODE_ENV = 'development'
+})
+afterEach(() => {
+  process.env.NODE_ENV = 'test'
+})
+
+it('should register newly added atoms', async () => {
+  const countAtom = atom(1)
+  const petAtom = atom('cat')
+
+  const DisplayCount: React.FC = () => {
+    const [clicked, setClicked] = useState(false)
+    const [count] = useAtom(countAtom)
+
+    return (
+      <>
+        <p>count: {count}</p>
+        <button onClick={() => setClicked(true)}>click</button>
+        {clicked && <DisplayPet />}
+      </>
+    )
+  }
+
+  const DisplayPet: React.FC = () => {
+    const [pet] = useAtom(petAtom)
+    return <p>pet: {pet}</p>
+  }
+
+  const RegisteredAtomsCount: React.FC = () => {
+    const atoms = useAtomsSnapshot()
+
+    return <p>atom count: {atoms.size}</p>
+  }
+
+  const { findByText, getByText } = render(
+    <Provider>
+      <DisplayCount />
+      <RegisteredAtomsCount />
+    </Provider>
+  )
+
+  await findByText('atom count: 1')
+  fireEvent.click(getByText('click'))
+  await findByText('atom count: 2')
+})
+
+it('should let you access atoms and their state', async () => {
+  const countAtom = atom(1)
+  countAtom.debugLabel = 'countAtom'
+  const petAtom = atom('cat')
+  petAtom.debugLabel = 'petAtom'
+
+  const Displayer: React.FC = () => {
+    useAtom(countAtom)
+    useAtom(petAtom)
+    return null
+  }
+
+  const SimpleDevtools: React.FC = () => {
+    const atoms = useAtomsSnapshot()
+
+    return (
+      <div>
+        {Array.from(atoms).map(([atom, atomValue]) => (
+          <p key={atom.debugLabel}>{`${atom.debugLabel}: ${atomValue}`}</p>
+        ))}
+      </div>
+    )
+  }
+
+  const { findByText } = render(
+    <Provider>
+      <Displayer />
+      <SimpleDevtools />
+    </Provider>
+  )
+
+  await findByText('countAtom: 1')
+  await findByText('petAtom: cat')
+})

tests/onmount.test.tsx

@@ -5,6 +5,7 @@ import { getTestProvider } from './testUtils'
 
 const Provider = getTestProvider()
 
+// FIXME the tests should also work on DEV
 let savedNodeEnv: string | undefined
 beforeEach(() => {
   savedNodeEnv = process.env.NODE_ENV

tests/query/atomWithQuery.test.tsx

@@ -7,6 +7,16 @@ import { getTestProvider } from '../testUtils'
 
 const Provider = getTestProvider()
 
+// FIXME the tests should also work on DEV
+let savedNodeEnv: string | undefined
+beforeEach(() => {
+  savedNodeEnv = process.env.NODE_ENV
+  process.env.NODE_ENV = 'production'
+})
+afterEach(() => {
+  process.env.NODE_ENV = savedNodeEnv
+})
+
 it('query basic test', async () => {
   const countAtom = atomWithQuery(() => ({
     queryKey: 'count1',