fix: use singly-linked WeakRefs to clean up React 18 StrictMode trash

package.json

@@ -70,5 +70,6 @@
   },
   "workspaces": [
     "packages/*"
-  ]
+  ],
+  "packageManager": "yarn@1.22.4+sha512.a1833b862fe52169bd6c2a033045a07df5bc6a23595c259e675fed1b2d035ab37abe6ce309720abb6636d68f03615054b6292dc0a70da31c8697fda228b50d18"
 }

packages/atoms/src/classes/Selectors.ts

@@ -4,6 +4,7 @@ import {
   AtomSelectorOrConfig,
   Cleanup,
   DependentCallback,
+  DependentEdge,
   EvaluationReason,
   Selectable,
 } from '../types/index'
@@ -16,10 +17,13 @@ const defaultResultsComparator = (a: any, b: any) => a === b
 export class SelectorCache<T = any, Args extends any[] = any[]> {
   public static $$typeof = Symbol.for(`${prefix}/SelectorCache`)
   public isDestroyed?: boolean
+  public isMaterialized?: boolean
   public nextReasons: EvaluationReason[] = []
   public prevReasons?: EvaluationReason[]
   public result?: T
   public task?: () => void
+  public _lastEdge?: WeakRef<DependentEdge>
+  public _prevCache?: WeakRef<SelectorCache>
 
   constructor(
     public id: string,
@@ -41,6 +45,11 @@ export class Selectors {
    */
   public _items: Record<string, SelectorCache<any, any>> = {}
 
+  /**
+   * A workaround for React StrictMode
+   */
+  public _lastCache?: WeakRef<SelectorCache<any, any>>
+
   /**
    * Map selectors (or selector config objects) to a base selectorKey that can
    * be used to predictably create selectorKey+params ids to look up the cache

packages/atoms/src/classes/instances/AtomInstance.ts

@@ -14,6 +14,7 @@ import {
   AtomGenerics,
   AtomGenericsToAtomApiGenerics,
   Cleanup,
+  DependentEdge,
   EvaluationReason,
   EvaluationSourceType,
   ExportsInfusedSetter,
@@ -87,6 +88,7 @@ export class AtomInstance<
   public _createdAt: number
   public _injectors?: InjectorDescriptor[]
   public _isEvaluating?: boolean
+  public _lastEdge?: WeakRef<DependentEdge>
   public _nextInjectors?: InjectorDescriptor[]
   public _promiseError?: Error
   public _promiseStatus?: PromiseStatus

packages/atoms/src/types/index.ts

@@ -174,8 +174,11 @@ export type DependentCallback = (
 export interface DependentEdge {
   callback?: DependentCallback
   createdAt: number
+  dependentKey?: string
   flags: number // calculated from the EdgeFlags
+  isMaterialized?: boolean
   operation: string
+  prevEdge?: WeakRef<DependentEdge>
   task?: () => void // for external edges - so they can unschedule jobs
 }
 

packages/react/src/hooks/useAtomInstance.ts

@@ -3,6 +3,7 @@ import {
   AnyAtomTemplate,
   AtomInstanceType,
   AtomParamsType,
+  DependentEdge,
   ParamlessTemplate,
 } from '@zedux/atoms'
 import { useEffect, useState } from 'react'
@@ -75,12 +76,14 @@ export const useAtomInstance: {
 
   // It should be fine for this to run every render. It's possible to change
   // approaches if it is too heavy sometimes. But don't memoize this call:
-  const instance = ecosystem.getInstance(atom as A, params as AtomParamsType<A>)
+  let instance = ecosystem.getInstance(atom as A, params as AtomParamsType<A>)
   const renderedState = instance.getState()
 
+  let edge: DependentEdge | undefined
+
   const addEdge = () => {
     if (!ecosystem._graph.nodes[instance.id]?.dependents.get(dependentKey)) {
-      ecosystem._graph.addEdge(
+      edge = ecosystem._graph.addEdge(
         dependentKey,
         instance.id,
         operation,
@@ -89,6 +92,15 @@ export const useAtomInstance: {
           if ((render as any).mounted) render({})
         }
       )
+
+      if (edge) {
+        edge.dependentKey = dependentKey
+
+        if (instance._lastEdge) {
+          edge.prevEdge = instance._lastEdge
+        }
+        instance._lastEdge = new WeakRef(edge)
+      }
     }
   }
 
@@ -99,6 +111,21 @@ export const useAtomInstance: {
   // Only remove the graph edge when the instance id changes or on component
   // destruction.
   useEffect(() => {
+    // re-get the instance in case StrictMode destroys it
+    instance = ecosystem.getInstance(atom as A, params as AtomParamsType<A>)
+
+    if (edge) {
+      let prevEdge = edge.prevEdge?.deref()
+
+      // clear out any junk edges added by StrictMode
+      while (prevEdge && !prevEdge.isMaterialized) {
+        ecosystem._graph.removeEdge(prevEdge.dependentKey!, instance.id)
+        prevEdge = prevEdge.prevEdge?.deref()
+      }
+
+      edge.isMaterialized = true
+    }
+
     // Try adding the edge again (will be a no-op unless React's StrictMode ran
     // this effect's cleanup unnecessarily)
     addEdge()

packages/react/src/hooks/useAtomSelector.ts

@@ -1,6 +1,7 @@
 import {
   AtomSelectorConfig,
   AtomSelectorOrConfig,
+  DependentEdge,
   haveDepsChanged,
   SelectorCache,
 } from '@zedux/atoms'
@@ -71,15 +72,32 @@ export const useAtomSelector = <T, Args extends any[]>(
     ;(render as any).mounted = true
   }
 
-  const cache = isSwappingRefs
+  let cache = isSwappingRefs
     ? (existingCache as SelectorCache<T, Args>)
     : selectors.getCache(selectorOrConfig, resolvedArgs)
 
+  let edge: DependentEdge | undefined
+
   const addEdge = () => {
     if (!_graph.nodes[cache.id]?.dependents.get(dependentKey)) {
-      _graph.addEdge(dependentKey, cache.id, OPERATION, External, () => {
+      edge = _graph.addEdge(dependentKey, cache.id, OPERATION, External, () => {
         if ((render as any).mounted) render({})
       })
+
+      if (edge) {
+        edge.dependentKey = dependentKey
+
+        if (cache._lastEdge) {
+          edge.prevEdge = cache._lastEdge
+        }
+        cache._lastEdge = new WeakRef(edge)
+      }
+
+      if (selectors._lastCache && selectors._lastCache.deref() !== cache) {
+        cache._prevCache = selectors._lastCache
+      }
+
+      selectors._lastCache = new WeakRef(cache)
     }
   }
 
@@ -90,6 +108,32 @@ export const useAtomSelector = <T, Args extends any[]>(
   ;(render as any).cache = cache as SelectorCache<any, any[]>
 
   useEffect(() => {
+    cache = isSwappingRefs
+      ? (existingCache as SelectorCache<T, Args>)
+      : selectors.getCache(selectorOrConfig, resolvedArgs)
+
+    if (edge) {
+      let prevEdge = edge.prevEdge?.deref()
+
+      // clear out any junk edges added by StrictMode
+      while (prevEdge && !prevEdge.isMaterialized) {
+        ecosystem._graph.removeEdge(prevEdge.dependentKey!, cache.id)
+        prevEdge = prevEdge.prevEdge?.deref()
+      }
+
+      edge.isMaterialized = true
+    }
+
+    let prevCache = cache._prevCache?.deref()
+
+    // clear out any junk caches created by StrictMode
+    while (prevCache && !prevCache.isMaterialized) {
+      selectors.destroyCache(prevCache, [], true)
+      prevCache = prevCache._prevCache?.deref()
+    }
+
+    cache.isMaterialized = true
+
     // Try adding the edge again (will be a no-op unless React's StrictMode ran
     // this effect's cleanup unnecessarily)
     addEdge()

packages/react/test/units/__snapshots__/useAtomSelector.test.tsx.snap

@@ -24,7 +24,9 @@ exports[`useAtomSelector inline selector that returns a different object referen
       "Test-:r0:" => {
         "callback": [Function],
         "createdAt": 123456789,
+        "dependentKey": "Test-:r0:",
         "flags": 2,
+        "isMaterialized": true,
         "operation": "useAtomSelector",
       },
     },
@@ -59,7 +61,9 @@ exports[`useAtomSelector inline selector that returns a different object referen
       "Test-:r0:" => {
         "callback": [Function],
         "createdAt": 123456789,
+        "dependentKey": "Test-:r0:",
         "flags": 2,
+        "isMaterialized": true,
         "operation": "useAtomSelector",
         "task": undefined,
       },
@@ -95,6 +99,7 @@ exports[`useAtomSelector inline selector that returns a different object referen
       "Test-:r0:" => {
         "callback": [Function],
         "createdAt": 123456789,
+        "dependentKey": "Test-:r0:",
         "flags": 2,
         "operation": "useAtomSelector",
       },
@@ -130,6 +135,7 @@ exports[`useAtomSelector inline selector that returns a different object referen
       "Test-:r0:" => {
         "callback": [Function],
         "createdAt": 123456789,
+        "dependentKey": "Test-:r0:",
         "flags": 2,
         "operation": "useAtomSelector",
         "task": undefined,

packages/react/test/units/useAtomSelector.test.tsx

@@ -443,8 +443,11 @@ describe('useAtomSelector', () => {
     expect(ecosystem.selectors.findAll()).toMatchInlineSnapshot(`
       {
         "@@selector-unnamed-1": SelectorCache {
+          "_lastEdge": WeakRef {},
+          "_prevCache": WeakRef {},
           "args": [],
           "id": "@@selector-unnamed-1",
+          "isMaterialized": true,
           "nextReasons": [],
           "prevReasons": [],
           "result": 1,
@@ -464,8 +467,11 @@ describe('useAtomSelector', () => {
     expect(ecosystem.selectors.findAll()).toMatchInlineSnapshot(`
       {
         "@@selector-unnamed-1": SelectorCache {
+          "_lastEdge": WeakRef {},
+          "_prevCache": WeakRef {},
           "args": [],
           "id": "@@selector-unnamed-1",
+          "isMaterialized": true,
           "nextReasons": [],
           "prevReasons": [],
           "result": 1,

yarn.lock

@@ -1385,22 +1385,16 @@
   resolved "https://registry.yarnpkg.com/@types/prettier/-/prettier-2.4.2.tgz#4c62fae93eb479660c3bd93f9d24d561597a8281"
   integrity sha512-ekoj4qOQYp7CvjX8ZDBgN86w3MqQhLE1hczEJbEIjgFEumDy+na/4AJAbLXfgEWFNB2pKadM5rPFtuSGMWK7xA==
 
-"@types/prop-types@*":
-  version "15.7.14"
-  resolved "https://registry.yarnpkg.com/@types/prop-types/-/prop-types-15.7.14.tgz#1433419d73b2a7ebfc6918dcefd2ec0d5cd698f2"
-  integrity sha512-gNMvNH49DJ7OJYv+KAKn0Xp45p8PLl6zo2YnvDIbTd4J6MER2BmWN49TG7n9LvkyihINxeKW8+3bfS2yDC9dzQ==
-
-"@types/react-dom@^18.3.0":
-  version "18.3.5"
-  resolved "https://registry.yarnpkg.com/@types/react-dom/-/react-dom-18.3.5.tgz#45f9f87398c5dcea085b715c58ddcf1faf65f716"
-  integrity sha512-P4t6saawp+b/dFrUr2cvkVsfvPguwsxtH6dNIYRllMsefqFzkZk5UIjzyDOv5g1dXIPdG4Sp1yCR4Z6RCUsG/Q==
-
-"@types/react@^18.3.0":
-  version "18.3.16"
-  resolved "https://registry.yarnpkg.com/@types/react/-/react-18.3.16.tgz#5326789125fac98b718d586ad157442ceb44ff28"
-  integrity sha512-oh8AMIC4Y2ciKufU8hnKgs+ufgbA/dhPTACaZPM86AbwX9QwnFtSoPWEeRUj8fge+v6kFt78BXcDhAU1SrrAsw==
-  dependencies:
-    "@types/prop-types" "*"
+"@types/react-dom@^19.0.2":
+  version "19.0.3"
+  resolved "https://registry.yarnpkg.com/@types/react-dom/-/react-dom-19.0.3.tgz#0804dfd279a165d5a0ad8b53a5b9e65f338050a4"
+  integrity sha512-0Knk+HJiMP/qOZgMyNFamlIjw9OFCsyC2ZbigmEEyXXixgre6IQpm/4V+r3qH4GC1JPvRJKInw+on2rV6YZLeA==
+
+"@types/react@^19.0.1":
+  version "19.0.7"
+  resolved "https://registry.yarnpkg.com/@types/react/-/react-19.0.7.tgz#c451968b999d1cb2d9207dc5ff56496164cf511d"
+  integrity sha512-MoFsEJKkAtZCrC1r6CM8U22GzhG7u2Wir8ons/aCKH6MBdD1ibV24zOSSkdZVUKqN5i396zG5VKLYZ3yaUZdLA==
+  dependencies:
     csstype "^3.0.2"
 
 "@types/semver@^7.3.12":
@@ -3300,7 +3294,7 @@ ignore@^5.0.4, ignore@^5.2.0:
   resolved "https://registry.yarnpkg.com/ignore/-/ignore-5.2.4.tgz#a291c0c6178ff1b960befe47fcdec301674a6324"
   integrity sha512-MAb38BcSbH0eHNBxn7ql2NH/kX33OkB3lZ1BNdh7ENeRChHTYsTvWrMubiIAMNS2llXEEgZ1MUOBtXChP3kaFQ==
 
-immer@^10.0.0:
+immer@^10.1.1:
   version "10.1.1"
   resolved "https://registry.yarnpkg.com/immer/-/immer-10.1.1.tgz#206f344ea372d8ea176891545ee53ccc062db7bc"
   integrity sha512-s2MPrmjovJcoMaHtx6K11Ra7oD05NT97w1IC5zpMkT6Atjr7H8LjaDd81iIxUYpMKSRRNMJE703M1Fhr/TctHw==