fix: optimize useAtomSelector for React 19 (#106)

@affects atoms, react

.vscode/launch.json

@@ -0,0 +1,15 @@
+{
+  "version": "0.2.0",
+  "configurations": [
+    {
+      "name": "Debug Tests",
+      "type": "node",
+      "request": "launch",
+      "runtimeExecutable": "yarn",
+      "args": ["jest", "--runInBand"],
+      "console": "integratedTerminal",
+      "cwd": "${workspaceRoot}",
+      "internalConsoleOptions": "neverOpen"
+    },
+  ]
+}

jest.setup.ts

@@ -1 +1,36 @@
 import '@testing-library/jest-dom/extend-expect'
+
+let callstacks: Record<string, string> = {}
+let id = 0
+
+afterEach(() => {
+  callstacks = {}
+  id = 0
+})
+
+const generateId = () => {
+  const stack =
+    (new Error().stack || '')
+      .split('\n')
+      .find(line => /\.test\.tsx:/.test(line)) || ''
+
+  return (callstacks[stack] ||= `:r${id++}:`)
+}
+
+jest.mock('react', () => ({
+  ...jest.requireActual('react'),
+  useId: generateId,
+}))
+
+// React's `useId` gives new ids in the same callstack when a component tree is
+// destroyed/unmounted. Call this to manually force ids to be recreated in tests
+// to mimic React's behavior.
+;(globalThis as any).clearUseIdEntry = (idNum: number) => {
+  const key = Object.keys(callstacks).find(
+    key => callstacks[key] === `:r${idNum}:`
+  )
+
+  if (key) {
+    delete callstacks[key]
+  }
+}

packages/atoms/src/classes/Selectors.ts

@@ -48,18 +48,6 @@ export class Selectors {
    */
   public _refBaseKeys = new WeakMap<AtomSelectorOrConfig<any, any>, string>()
 
-  /**
-   * Used to work around React double-renders and double-effects.
-   */
-  public _storage: Record<
-    string,
-    {
-      cache?: SelectorCache
-      ignorePhase?: number
-      timeoutId?: ReturnType<typeof requestIdleCallback | typeof setTimeout>
-    }
-  > = {}
-
   constructor(private readonly ecosystem: Ecosystem) {}
 
   public addDependent(
@@ -182,9 +170,8 @@ export class Selectors {
       args as Args,
       true
     )
-    if (!id) return
 
-    return this._items[id]
+    return id && this._items[id]
   }
 
   /**
@@ -305,6 +292,7 @@ export class Selectors {
     _graph.removeDependencies(id)
     _graph.removeNode(id)
     delete this._items[id]
+    this._refBaseKeys.delete(cache.selectorRef)
     cache.isDestroyed = true
     // don't delete the ref from this._refBaseKeys; this selector cache isn't
     // necessarily the only one using it (if the selector takes params). Just
@@ -369,20 +357,19 @@ export class Selectors {
   /**
    * Should only be used internally
    */
-  public _swapRefs(
-    oldRef: AtomSelectorOrConfig<any, any[]>,
-    newRef: AtomSelectorOrConfig<any, any[]>,
+  public _swapRefs<T, Args extends any[]>(
+    oldCache: SelectorCache<T, Args>,
+    newRef: AtomSelectorOrConfig<T, Args>,
     args: any[] = []
   ) {
-    const existingCache = this.find(oldRef, args)
-    const baseKey = this._refBaseKeys.get(oldRef)
+    const baseKey = this._refBaseKeys.get(oldCache.selectorRef)
 
-    if (!existingCache || !baseKey) return
+    if (!baseKey) return
 
     this._refBaseKeys.set(newRef, baseKey)
-    this._refBaseKeys.delete(oldRef)
-    existingCache.selectorRef = newRef
-    this.runSelector(existingCache.id, args)
+    this._refBaseKeys.delete(oldCache.selectorRef)
+    oldCache.selectorRef = newRef
+    this.runSelector(oldCache.id, args, false, true)
   }
 
   /**
@@ -395,7 +382,6 @@ export class Selectors {
     })
 
     this._refBaseKeys = new WeakMap()
-    this._storage = {}
   }
 
   /**
@@ -429,7 +415,8 @@ export class Selectors {
   private runSelector<T = any, Args extends any[] = []>(
     id: string,
     args: Args,
-    isInitializing?: boolean
+    isInitializing?: boolean,
+    skipNotifyingDependents?: boolean
   ) {
     const { _evaluationStack, _graph, _mods, modBus } = this.ecosystem
     _graph.bufferUpdates(id)
@@ -451,7 +438,9 @@ export class Selectors {
       const result = selector(_evaluationStack.atomGetters, ...args)
 
       if (!isInitializing && !resultsComparator(result, cache.result as T)) {
-        _graph.scheduleDependents(id, cache.nextReasons, result, cache.result)
+        if (!skipNotifyingDependents) {
+          _graph.scheduleDependents(id, cache.nextReasons, result, cache.result)
+        }
 
         if (_mods.stateChanged) {
           modBus.dispatch(

packages/machines/test/integrations/state-machines.test.tsx

@@ -3,7 +3,7 @@ import {
   InjectMachineStoreParams,
   MachineState,
 } from '@zedux/machines'
-import { api, atom } from '@zedux/react'
+import { api, atom } from '@zedux/atoms'
 import { ecosystem } from '../../../react/test/utils/ecosystem'
 
 const injectMachine = <

packages/machines/test/snippets/api.tsx

@@ -5,7 +5,7 @@ import {
   ion,
   useAtomSelector,
   useAtomValue,
-} from '@zedux/react'
+} from '../../../react/src'
 import { injectMachineStore } from '@zedux/machines'
 import React, { Suspense, useState } from 'react'
 

packages/react/src/hooks/useAtomInstance.ts

@@ -117,8 +117,12 @@ export const useAtomInstance: {
     }
 
     return () => {
-      // no need to set the "ref"'s `.mounted` property to false here
+      // remove the edge immediately - no need for a delay here. When StrictMode
+      // double-invokes (invokes, then cleans up, then re-invokes) this effect,
+      // it's expected that any `ttl: 0` atoms get destroyed and recreated -
+      // that's part of what StrictMode is ensuring
       ecosystem._graph.removeEdge(dependentKey, instance.id)
+      // no need to set `render.mounted` to false here
     }
   }, [instance.id])
 

packages/react/src/hooks/useAtomSelector.ts

@@ -29,10 +29,10 @@ export const useAtomSelector = <T, Args extends any[]>(
   const { _graph, selectors } = ecosystem
   const dependentKey = useReactComponentId()
   const [, render] = useState<undefined | object>()
-  const storage =
-    (render as any).storage || (selectors._storage[dependentKey] ||= {})
 
-  const existingCache = storage.cache as SelectorCache<T, Args> | undefined
+  const existingCache = (render as any).cache as
+    | SelectorCache<T, Args>
+    | undefined
 
   const argsChanged =
     !existingCache ||
@@ -46,98 +46,74 @@ export const useAtomSelector = <T, Args extends any[]>(
       : haveDepsChanged(existingCache.args, args))
 
   const resolvedArgs = argsChanged ? args : (existingCache.args as Args)
-  const cache = selectors.getCache(selectorOrConfig, resolvedArgs)
-  const renderedResult = cache.result
-
-  if (cache !== existingCache) {
-    if (existingCache) {
-      // yes, remove this during render
-      _graph.removeEdge(dependentKey, existingCache.id)
-    }
 
-    storage.cache = cache as SelectorCache<any, any[]>
-  }
-
-  // When an inline selector returns a referentially unstable result every run,
-  // we have to ignore the subsequent update. Do that using a "state machine"
-  // that goes from 0 -> 1 -> 2. This machine ensures that the ignored update
-  // occurs after the component rerenders and the effect reruns after that
-  // render. This works with strict mode on or off. Use the stable `render`
-  // function as a "ref" :O
-  if (storage.ignorePhase === 1) {
-    storage.ignorePhase = 2
+  // if the refs/args don't match, existingCache has refCount: 1, there is no
+  // cache yet for the new ref, and the new ref has the same name, assume it's
+  // an inline selector
+  const isSwappingRefs =
+    existingCache &&
+    existingCache.selectorRef !== selectorOrConfig &&
+    !argsChanged
+      ? _graph.nodes[existingCache.id]?.refCount === 1 &&
+        !selectors._refBaseKeys.has(selectorOrConfig) &&
+        selectors._getIdealCacheId(existingCache.selectorRef) ===
+          selectors._getIdealCacheId(selectorOrConfig)
+      : false
+
+  if (isSwappingRefs) {
+    // switch `mounted` to false temporarily to prevent circular rerenders
+    ;(render as any).mounted = false
+    selectors._swapRefs(
+      existingCache as SelectorCache<any, any[]>,
+      selectorOrConfig as AtomSelectorOrConfig<any, any[]>,
+      resolvedArgs
+    )
+    ;(render as any).mounted = false
   }
 
-  let cancelCleanup = false
+  const cache = isSwappingRefs
+    ? (existingCache as SelectorCache<T, Args>)
+    : selectors.getCache(selectorOrConfig, resolvedArgs)
 
-  useEffect(() => {
-    cancelCleanup = true
-    delete selectors._storage[dependentKey]
-    ;(render as any).storage = storage
-
-    // re-get the cache in case an unmounting component's effect cleanup
-    // destroyed it before we could add this dependent
-    const newCache = selectors.getCache(selectorOrConfig, resolvedArgs)
-
-    const cleanup = () => {
-      if (cancelCleanup) {
-        cancelCleanup = false
-
-        return
-      }
+  const addEdge = () => {
+    if (!_graph.nodes[cache.id]?.dependents.get(dependentKey)) {
+      _graph.addEdge(dependentKey, cache.id, OPERATION, External, () => {
+        if ((render as any).mounted) render({})
+      })
+    }
+  }
 
-      if (storage.ignorePhase !== 1) {
-        delete selectors._storage[dependentKey]
+  // Yes, subscribe during render. This operation is idempotent.
+  addEdge()
 
-        queueMicrotask(() => {
-          _graph.removeEdge(dependentKey, newCache.id)
-        })
-      }
-    }
+  const renderedResult = cache.result
+  ;(render as any).cache = cache as SelectorCache<any, any[]>
 
-    // Make this function idempotent to guard against React's double-invocation
-    if (_graph.nodes[newCache.id]?.dependents.get(dependentKey)) {
-      return cleanup
-    }
+  useEffect(() => {
+    // Try adding the edge again (will be a no-op unless React's StrictMode ran
+    // this effect's cleanup unnecessarily)
+    addEdge()
 
-    _graph.addEdge(dependentKey, newCache.id, OPERATION, External, () =>
-      render({})
-    )
+    // use the referentially stable render function as a ref :O
+    ;(render as any).mounted = true
 
     // an unmounting component's effect cleanup can force-destroy the selector
-    // or update its dependencies before this component is mounted. If that
-    // happened, trigger a rerender to recache the selector and/or get its new
-    // result. On the rerender, ignore changes
-    if (newCache.result !== renderedResult && !storage.ignorePhase) {
-      storage.ignorePhase = 1
+    // or update the state of its dependencies (causing it to rerun) before we
+    // set `render.mounted`. If that happened, trigger a rerender to recreate
+    // the selector and/or get its new state
+    if (cache.isDestroyed || cache.result !== renderedResult) {
       render({})
     }
 
-    if (storage.ignorePhase === 2) {
-      storage.ignorePhase = 0
+    return () => {
+      // remove the edge immediately - no need for a delay here. When StrictMode
+      // double-invokes (invokes, then cleans up, then re-invokes) this effect,
+      // it's expected that selectors and `ttl: 0` atoms with no other
+      // dependents get destroyed and recreated - that's part of what StrictMode
+      // is ensuring
+      _graph.removeEdge(dependentKey, cache.id)
+      // no need to set `render.mounted` to false here
     }
-
-    // React StrictMode's double renders can wreak havoc on the selector cache.
-    // Clean up havoc
-    if (storage.timeoutId == null) {
-      const removeCruft = () => {
-        storage.timeoutId = null
-        cancelCleanup = false
-
-        Object.values(selectors._storage).forEach(storageItem => {
-          if (storageItem.cache?.id) {
-            selectors._destroySelector(storageItem.cache.id)
-          }
-        })
-      }
-
-      storage.timeoutId =
-        typeof requestIdleCallback !== 'undefined'
-          ? requestIdleCallback(removeCruft, { timeout: 500 })
-          : setTimeout(removeCruft, 500)
-    }
-
-    return cleanup
   }, [cache])
 
   return renderedResult as T