refactor

src/vanilla/store.ts

@@ -467,21 +467,26 @@ const buildStore = (
     return dependents
   }
 
-  // This is a topological sort via depth-first search, slightly modified from
-  // what's described here for simplicity and performance reasons:
-  // https://en.wikipedia.org/wiki/Topological_sorting#Depth-first_search
-  function getSortedDependents(
+  const recomputeDependents = <Value>(
     batch: Batch,
-    rootAtom: AnyAtom,
-    rootAtomState: AtomState,
-  ): [[AnyAtom, AtomState, number][], Set<AnyAtom>] {
-    const sorted: [atom: AnyAtom, atomState: AtomState, epochNumber: number][] =
-      []
+    atom: Atom<Value>,
+    atomState: AtomState<Value>,
+  ) => {
+    // Step 1: traverse the dependency graph to build the topsorted atom list
+    // We don't bother to check for cycles, which simplifies the algorithm.
+    // This is a topological sort via depth-first search, slightly modified from
+    // what's described here for simplicity and performance reasons:
+    // https://en.wikipedia.org/wiki/Topological_sorting#Depth-first_search
+    const topSortedReversed: [
+      atom: AnyAtom,
+      atomState: AtomState,
+      epochNumber: number,
+    ][] = []
     const visiting = new Set<AnyAtom>()
     const visited = new Set<AnyAtom>()
     // Visit the root atom. This is the only atom in the dependency graph
     // without incoming edges, which is one reason we can simplify the algorithm
-    const stack: [a: AnyAtom, aState: AtomState][] = [[rootAtom, rootAtomState]]
+    const stack: [a: AnyAtom, aState: AtomState][] = [[atom, atomState]]
     while (stack.length > 0) {
       const [a, aState] = stack[stack.length - 1]!
       if (visited.has(a)) {
@@ -493,9 +498,11 @@ const buildStore = (
         // The algorithm calls for pushing onto the front of the list. For
         // performance, we will simply push onto the end, and then will iterate in
         // reverse order later.
-        sorted.push([a, aState, aState.n])
+        topSortedReversed.push([a, aState, aState.n])
         // Atom has been visited but not yet processed
         visited.add(a)
+        // Mark atom dirty
+        aState.x = true
         stack.pop()
         continue
       }
@@ -507,32 +514,13 @@ const buildStore = (
         }
       }
     }
-    return [sorted, visited]
-  }
-
-  const recomputeDependents = <Value>(
-    batch: Batch,
-    atom: Atom<Value>,
-    atomState: AtomState<Value>,
-  ) => {
-    // Step 1: traverse the dependency graph to build the topsorted atom list
-    // We don't bother to check for cycles, which simplifies the algorithm.
-    const [topsortedAtoms, markedAtoms] = getSortedDependents(
-      batch,
-      atom,
-      atomState,
-    )
-    // TODO this should be refactored
-    markedAtoms.forEach((a) => {
-      getAtomState(a).x = true
-    })
 
-    // Step 2: use the topsorted atom list to recompute all affected atoms
+    // Step 2: use the topSortedReversed atom list to recompute all affected atoms
     // Track what's changed, so that we can short circuit when possible
     addBatchFuncHigh(batch, () => {
       const changedAtoms = new Set<AnyAtom>([atom])
-      for (let i = topsortedAtoms.length - 1; i >= 0; --i) {
-        const [a, aState, prevEpochNumber] = topsortedAtoms[i]!
+      for (let i = topSortedReversed.length - 1; i >= 0; --i) {
+        const [a, aState, prevEpochNumber] = topSortedReversed[i]!
         let hasChangedDeps = false
         for (const dep of aState.d.keys()) {
           if (dep !== a && changedAtoms.has(dep)) {