refactor(core): re-implement useAtom with useReducer (#687)

* wip: use-sync-external-store

* use reducer just like that

* fix: avoid reading atom value twice

* fix useAtomsSnapshot

* update size snapshot

* a workaround for react devtools

* update size snapshot

* empty commit

* Update src/devtools/useAtomsSnapshot.ts

Co-authored-by: Mathis Møller <mmm@cryosinternational.com>

* Update src/devtools/useAtomsSnapshot.ts

Co-authored-by: Mathis Møller <mmm@cryosinternational.com>

* Update src/devtools/useAtomsSnapshot.ts

Co-authored-by: Mathis Møller <mmm@cryosinternational.com>

Co-authored-by: Mathis Møller <mmm@cryosinternational.com>

Author: dai-shi

.size-snapshot.json

@@ -1,8 +1,8 @@
 {
   "index.js": {
-    "bundled": 22130,
-    "minified": 8884,
-    "gzipped": 3283,
+    "bundled": 19763,
+    "minified": 8215,
+    "gzipped": 3043,
     "treeshaked": {
       "rollup": {
         "code": 14,
@@ -14,9 +14,9 @@
     }
   },
   "index.mjs": {
-    "bundled": 22130,
-    "minified": 8884,
-    "gzipped": 3283,
+    "bundled": 19763,
+    "minified": 8215,
+    "gzipped": 3043,
     "treeshaked": {
       "rollup": {
         "code": 14,
@@ -56,9 +56,9 @@
     }
   },
   "devtools.js": {
-    "bundled": 20426,
-    "minified": 8380,
-    "gzipped": 3171,
+    "bundled": 20062,
+    "minified": 8222,
+    "gzipped": 3114,
     "treeshaked": {
       "rollup": {
         "code": 28,
@@ -70,9 +70,9 @@
     }
   },
   "devtools.mjs": {
-    "bundled": 20426,
-    "minified": 8380,
-    "gzipped": 3171,
+    "bundled": 20062,
+    "minified": 8222,
+    "gzipped": 3114,
     "treeshaked": {
       "rollup": {
         "code": 28,

package.json

@@ -170,6 +170,7 @@
       "tests/**/*.{js,ts,tsx}"
     ]
   },
+  "dependencies": {},
   "devDependencies": {
     "@babel/core": "^7.15.0",
     "@babel/plugin-transform-react-jsx": "^7.14.9",

src/core/Provider.ts

@@ -1,4 +1,10 @@
-import { createElement, useCallback, useDebugValue, useRef } from 'react'
+import {
+  createElement,
+  useDebugValue,
+  useEffect,
+  useRef,
+  useState,
+} from 'react'
 import type { PropsWithChildren } from 'react'
 import type { Atom, Scope } from './atom'
 import {
@@ -10,7 +16,6 @@ import {
 import type { ScopeContainerForDevelopment } from './contexts'
 import { DEV_GET_ATOM_STATE, DEV_GET_MOUNTED } from './store'
 import type { AtomState, Store } from './store'
-import { useMutableSource } from './useMutableSource'
 
 export const Provider = ({
   initialValues,
@@ -40,9 +45,7 @@ export const Provider = ({
   return createElement(
     ScopeContainerContext.Provider,
     {
-      value: scopeContainerRef.current as ReturnType<
-        typeof createScopeContainer
-      >,
+      value: scopeContainerRef.current,
     },
     children
   )
@@ -75,11 +78,14 @@ 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, , devMutableSource, devSubscribe] = scopeContainer
-  const atoms = useMutableSource(
-    devMutableSource,
-    useCallback((devContainer) => devContainer.atoms, []),
-    devSubscribe
-  )
+  const [store, devStore] = scopeContainer
+  const [atoms, setAtoms] = useState(devStore.atoms)
+  useEffect(() => {
+    // HACK creating a new reference for useDebugValue to update
+    const callback = () => setAtoms([...devStore.atoms])
+    const unsubscribe = devStore.subscribe(callback)
+    callback()
+    return unsubscribe
+  }, [devStore])
   useDebugValue([store, atoms], stateToPrintable)
 }

src/core/contexts.ts

@@ -1,50 +1,46 @@
 import { createContext } from 'react'
 import type { Context } from 'react'
 import type { Atom, Scope } from './atom'
-import { GET_VERSION, createStore } from './store'
-import { createMutableSource } from './useMutableSource'
+import { createStore } from './store'
 
 const createScopeContainerForProduction = (
   initialValues?: Iterable<readonly [Atom<unknown>, unknown]>
 ) => {
   const store = createStore(initialValues)
-  const mutableSource = createMutableSource(store, store[GET_VERSION])
-  return [store, mutableSource] as const
+  return [store] as const
 }
 
 const createScopeContainerForDevelopment = (
   initialValues?: Iterable<readonly [Atom<unknown>, unknown]>
 ) => {
-  let devVersion = 0
-  const devListeners = new Set<() => void>()
-  const devContainer = {
+  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) => {
-    ++devVersion
     if (isNewAtom) {
       // FIXME memory leak
       // we should probably remove unmounted atoms eventually
-      devContainer.atoms = [...devContainer.atoms, updatedAtom]
+      devStore.atoms = [...devStore.atoms, updatedAtom]
     }
     Promise.resolve().then(() => {
-      devListeners.forEach((listener) => listener())
+      devStore.listeners.forEach((listener) => listener())
     })
   }
   const store = createStore(initialValues, stateListener)
-  const mutableSource = createMutableSource(store, store[GET_VERSION])
-  const devMutableSource = createMutableSource(devContainer, () => devVersion)
-  const devSubscribe = (_: unknown, callback: () => void) => {
-    devListeners.add(callback)
-    return () => devListeners.delete(callback)
-  }
-  return [store, mutableSource, devMutableSource, devSubscribe] as const
+  return [store, devStore] as const
 }
 
 export const isDevScopeContainer = (
   scopeContainer: ScopeContainer
 ): scopeContainer is ScopeContainerForDevelopment => {
-  return scopeContainer.length > 2
+  return scopeContainer.length > 1
 }
 
 type ScopeContainerForProduction = ReturnType<

src/core/store.ts

@@ -66,7 +66,6 @@ type Mounted = {
 type StateListener = (updatedAtom: AnyAtom, isNewAtom: boolean) => void
 
 // store methods
-export const GET_VERSION = 'v'
 export const READ_ATOM = 'r'
 export const WRITE_ATOM = 'w'
 export const FLUSH_PENDING = 'f'
@@ -79,7 +78,6 @@ export const createStore = (
   initialValues?: Iterable<readonly [AnyAtom, unknown]>,
   stateListener?: StateListener
 ) => {
-  let version = 0
   const atomStateMap = new WeakMap<AnyAtom, AtomState>()
   const mountedMap = new WeakMap<AnyAtom, Mounted>()
   const pendingMap = new Map<AnyAtom, ReadDependencies | undefined>()
@@ -142,6 +140,9 @@ export const createStore = (
     if (!('v' in atomState) || !Object.is(atomState.v, value)) {
       atomState.v = value
       ++atomState.r // increment revision
+      if (atomState.d.has(atom)) {
+        atomState.d.set(atom, atomState.r)
+      }
     }
     commitAtomState(atom, atomState, dependencies && prevDependencies)
   }
@@ -569,7 +570,6 @@ export const createStore = (
     if (stateListener) {
       stateListener(atom, isNewAtom)
     }
-    ++version
     if (!pendingMap.has(atom)) {
       pendingMap.set(atom, prevDependencies)
     }
@@ -619,7 +619,6 @@ export const createStore = (
 
   if (typeof process === 'object' && process.env.NODE_ENV !== 'production') {
     return {
-      [GET_VERSION]: () => version,
       [READ_ATOM]: readAtom,
       [WRITE_ATOM]: writeAtom,
       [FLUSH_PENDING]: flushPending,
@@ -630,7 +629,6 @@ export const createStore = (
     }
   }
   return {
-    [GET_VERSION]: () => version,
     [READ_ATOM]: readAtom,
     [WRITE_ATOM]: writeAtom,
     [FLUSH_PENDING]: flushPending,

src/core/useAtom.ts

@@ -1,9 +1,13 @@
-import { useCallback, useContext, useDebugValue, useEffect } from 'react'
+import {
+  useCallback,
+  useContext,
+  useDebugValue,
+  useEffect,
+  useReducer,
+} from 'react'
 import type { Atom, Scope, SetAtom, WritableAtom } from './atom'
 import { getScopeContext } from './contexts'
 import { FLUSH_PENDING, READ_ATOM, SUBSCRIBE_ATOM, WRITE_ATOM } from './store'
-import type { Store } from './store'
-import { useMutableSource } from './useMutableSource'
 
 const isWritable = <Value, Update>(
   atom: Atom<Value> | WritableAtom<Value, Update>
@@ -41,32 +45,6 @@ export function useAtom<Value, Update>(
   atom: Atom<Value> | WritableAtom<Value, Update>,
   scope?: Scope
 ) {
-  const getAtomValue = useCallback(
-    (store: Store) => {
-      const atomState = store[READ_ATOM](atom)
-      if (atomState.e) {
-        throw atomState.e // read error
-      }
-      if (atomState.p) {
-        throw atomState.p // read promise
-      }
-      if (atomState.w) {
-        throw atomState.w // write promise
-      }
-      if ('v' in atomState) {
-        return atomState.v as Value
-      }
-      throw new Error('no atom value')
-    },
-    [atom]
-  )
-
-  const subscribe = useCallback(
-    (store: Store, callback: () => void) =>
-      store[SUBSCRIBE_ATOM](atom, callback),
-    [atom]
-  )
-
   if ('scope' in atom) {
     console.warn(
       'atom.scope is deprecated. Please do useAtom(atom, scope) instead.'
@@ -75,8 +53,33 @@ export function useAtom<Value, Update>(
   }
 
   const ScopeContext = getScopeContext(scope)
-  const [store, mutableSource] = useContext(ScopeContext)
-  const value = useMutableSource(mutableSource, getAtomValue, subscribe)
+  const [store] = useContext(ScopeContext)
+
+  const getAtomValue = useCallback(() => {
+    const atomState = store[READ_ATOM](atom)
+    if (atomState.e) {
+      throw atomState.e // read error
+    }
+    if (atomState.p) {
+      throw atomState.p // read promise
+    }
+    if (atomState.w) {
+      throw atomState.w // write promise
+    }
+    if ('v' in atomState) {
+      return atomState.v as Value
+    }
+    throw new Error('no atom value')
+  }, [store, atom])
+
+  const [value, forceUpdate] = useReducer(getAtomValue, undefined, getAtomValue)
+
+  useEffect(() => {
+    const unsubscribe = store[SUBSCRIBE_ATOM](atom, forceUpdate)
+    forceUpdate()
+    return unsubscribe
+  }, [store, atom])
+
   useEffect(() => {
     store[FLUSH_PENDING]()
   })