Improve performance of recomputeDependents

src/vanilla/store.ts

@@ -476,53 +476,82 @@ export const createStore = () => {
   }
 
   const recomputeDependents = (atom: AnyAtom): void => {
-    const dependencyMap = new Map<AnyAtom, Set<AnyAtom>>()
-    const dirtyMap = new WeakMap<AnyAtom, number>()
-    const getDependents = (a: AnyAtom): Dependents => {
+    // Returns an array rather than a set to work around an issue with
+    // transpilation for older environments. Otherwise, Array.prototype.push
+    // below will fail due to the transpiled set iterator implementation.
+    // https://github.com/pmndrs/jotai/discussions/2368#discussioncomment-8274991
+    const getDependents = (a: AnyAtom): Array<AnyAtom> => {
       const dependents = new Set(mountedMap.get(a)?.t)
       pendingMap.forEach((_, pendingAtom) => {
         if (getAtomState(pendingAtom)?.d.has(a)) {
           dependents.add(pendingAtom)
         }
       })
-      return dependents
-    }
-    const loop1 = (a: AnyAtom) => {
-      getDependents(a).forEach((dependent) => {
-        if (dependent !== a) {
-          dependencyMap.set(
-            dependent,
-            (dependencyMap.get(dependent) || new Set()).add(a),
-          )
-          dirtyMap.set(dependent, (dirtyMap.get(dependent) || 0) + 1)
-          loop1(dependent)
+      return Array.from(dependents)
+    }
+
+    // This is a topological sort via depth-first search, slightly modified from
+    // what's described here for performance reasons:
+    // https://en.wikipedia.org/wiki/Topological_sorting#Depth-first_search
+
+    // Step 1: find all atoms in the graph
+    const allAtomsInDependencyGraph = new Array<AnyAtom>()
+    const seen = new Set<AnyAtom>()
+    const queue = [atom]
+    let a: AnyAtom | undefined
+    while ((a = queue.pop())) {
+      if (!seen.has(a)) {
+        seen.add(a)
+        allAtomsInDependencyGraph.push(a)
+        queue.push(...getDependents(a))
+      }
+    }
+
+    // Step 2: traverse the dependency graph to build the topsorted atom list
+    // We don't bother to check for cycles, which simplifies the algorithm.
+    const topsortedAtoms = new Array<AnyAtom>(allAtomsInDependencyGraph.length)
+    let topsortedAtomsIndex = allAtomsInDependencyGraph.length
+    const markedAtoms = new Set<AnyAtom>()
+    const visit = (n: AnyAtom) => {
+      if (markedAtoms.has(n)) {
+        return
+      }
+      markedAtoms.add(n)
+      for (const m of getDependents(n)) {
+        if (n !== m) {
+          visit(m)
         }
-      })
+      }
+      // The algorithm calls for pushing onto the front of the list, so we fill
+      // the array in reverse order.
+      topsortedAtoms[--topsortedAtomsIndex] = n
     }
-    loop1(atom)
-    const loop2 = (a: AnyAtom) => {
-      getDependents(a).forEach((dependent) => {
-        if (dependent !== a) {
-          let dirtyCount = dirtyMap.get(dependent)
-          if (dirtyCount) {
-            dirtyMap.set(dependent, --dirtyCount)
-          }
-          if (!dirtyCount) {
-            let isChanged = !!dependencyMap.get(dependent)?.size
-            if (isChanged) {
-              const prevAtomState = getAtomState(dependent)
-              const nextAtomState = readAtomState(dependent, true)
-              isChanged = !isEqualAtomValue(prevAtomState, nextAtomState)
-            }
-            if (!isChanged) {
-              dependencyMap.forEach((s) => s.delete(dependent))
-            }
-          }
-          loop2(dependent)
+    while (allAtomsInDependencyGraph.length) {
+      visit(allAtomsInDependencyGraph.pop()!)
+    }
+
+    // Step 3: use the topsorted atom list to recompute all affected atoms
+    // Track what's changed, so that we can short circuit when possible
+    const changedAtoms = new Set<AnyAtom>([atom])
+    for (const a of topsortedAtoms) {
+      const prevAtomState = getAtomState(a)
+      if (!prevAtomState) {
+        continue
+      }
+      let hasChangedDeps = false
+      for (const dep of prevAtomState.d.keys()) {
+        if (dep !== a && changedAtoms.has(dep)) {
+          hasChangedDeps = true
+          break
         }
-      })
+      }
+      if (hasChangedDeps) {
+        const nextAtomState = readAtomState(a, true)
+        if (!isEqualAtomValue(prevAtomState, nextAtomState)) {
+          changedAtoms.add(a)
+        }
+      }
     }
-    loop2(atom)
   }
 
   const writeAtomState = <Value, Args extends unknown[], Result>(