lazy recomputeDependents (#2827)

* add failing test: batches sync writes

* defer recompute dependents if possible

* add task delay for some unknown reason to make this test pass

* sync with fix/write-batching-2

* add a test

---------

Co-authored-by: daishi <[email protected]>

src/vanilla/store.ts

@@ -104,6 +104,8 @@ type AtomState<Value = AnyValue> = {
   v?: Value
   /** Atom error */
   e?: AnyError
+  /** Indicates that the atom value has been changed */
+  x?: true
 }
 
 const isAtomStateInitialized = <Value>(atomState: AtomState<Value>) =>
@@ -164,6 +166,8 @@ const addDependency = <Value>(
 type Batch = Readonly<{
   /** Atom dependents map */
   D: Map<AnyAtom, Set<AnyAtom>>
+  /** High priority functions */
+  H: Set<() => void>
   /** Medium priority functions */
   M: Set<() => void>
   /** Low priority functions */
@@ -172,10 +176,15 @@ type Batch = Readonly<{
 
 const createBatch = (): Batch => ({
   D: new Map(),
+  H: new Set(),
   M: new Set(),
   L: new Set(),
 })
 
+const addBatchFuncHigh = (batch: Batch, fn: () => void) => {
+  batch.H.add(fn)
+}
+
 const addBatchFuncMedium = (batch: Batch, fn: () => void) => {
   batch.M.add(fn)
 }
@@ -232,6 +241,7 @@ const flushBatch = (batch: Batch) => {
   }
   while (batch.M.size || batch.L.size) {
     batch.D.clear()
+    copySetAndClear(batch.H).forEach(call)
     copySetAndClear(batch.M).forEach(call)
     copySetAndClear(batch.L).forEach(call)
   }
@@ -305,11 +315,11 @@ const buildStore = (
         addPendingPromiseToDependency(atom, valueOrPromise, getAtomState(a))
       }
       atomState.v = valueOrPromise
-      delete atomState.e
     } else {
       atomState.v = valueOrPromise
-      delete atomState.e
     }
+    delete atomState.e
+    delete atomState.x
     if (!hasPrevValue || !Object.is(prevValue, atomState.v)) {
       ++atomState.n
       if (pendingPromise) {
@@ -321,15 +331,14 @@ const buildStore = (
   const readAtomState = <Value>(
     batch: Batch | undefined,
     atom: Atom<Value>,
-    dirtyAtoms?: Set<AnyAtom>,
   ): AtomState<Value> => {
     const atomState = getAtomState(atom)
     // See if we can skip recomputing this atom.
     if (isAtomStateInitialized(atomState)) {
       // If the atom is mounted, we can use cached atom state.
       // because it should have been updated by dependencies.
       // We can't use the cache if the atom is dirty.
-      if (atomState.m && !dirtyAtoms?.has(atom)) {
+      if (atomState.m && !atomState.x) {
         return atomState
       }
       // Otherwise, check if the dependencies have changed.
@@ -339,7 +348,7 @@ const buildStore = (
           ([a, n]) =>
             // Recursively, read the atom state of the dependency, and
             // check if the atom epoch number is unchanged
-            readAtomState(batch, a, dirtyAtoms).n === n,
+            readAtomState(batch, a).n === n,
         )
       ) {
         return atomState
@@ -362,7 +371,7 @@ const buildStore = (
         return returnAtomValue(aState)
       }
       // a !== atom
-      const aState = readAtomState(batch, a, dirtyAtoms)
+      const aState = readAtomState(batch, a)
       try {
         return returnAtomValue(aState)
       } finally {
@@ -423,6 +432,7 @@ const buildStore = (
     } catch (error) {
       delete atomState.v
       atomState.e = error
+      delete atomState.x
       ++atomState.n
       return atomState
     } finally {
@@ -457,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)) {
@@ -483,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
       }
@@ -497,44 +514,31 @@ 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,
-    )
 
-    // 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
-    const changedAtoms = new Set<AnyAtom>([atom])
-    for (let i = topsortedAtoms.length - 1; i >= 0; --i) {
-      const [a, aState, prevEpochNumber] = topsortedAtoms[i]!
-      let hasChangedDeps = false
-      for (const dep of aState.d.keys()) {
-        if (dep !== a && changedAtoms.has(dep)) {
-          hasChangedDeps = true
-          break
+    addBatchFuncHigh(batch, () => {
+      const changedAtoms = new Set<AnyAtom>([atom])
+      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)) {
+            hasChangedDeps = true
+            break
+          }
         }
-      }
-      if (hasChangedDeps) {
-        readAtomState(batch, a, markedAtoms)
-        mountDependencies(batch, a, aState)
-        if (prevEpochNumber !== aState.n) {
-          registerBatchAtom(batch, a, aState)
-          changedAtoms.add(a)
+        if (hasChangedDeps) {
+          readAtomState(batch, a)
+          mountDependencies(batch, a, aState)
+          if (prevEpochNumber !== aState.n) {
+            registerBatchAtom(batch, a, aState)
+            changedAtoms.add(a)
+          }
         }
+        delete aState.x
       }
-      markedAtoms.delete(a)
-    }
+    })
   }
 
   const writeAtomState = <Value, Args extends unknown[], Result>(

tests/vanilla/dependency.test.tsx

@@ -405,3 +405,23 @@ it('can cache reading an atom in write function (with mounting)', () => {
   store.set(w)
   expect(aReadCount).toBe(1)
 })
+
+it('batches sync writes', () => {
+  const a = atom(0)
+  const b = atom((get) => get(a))
+  const fetch = vi.fn()
+  const c = atom((get) => fetch(get(a)))
+  const w = atom(null, (get, set) => {
+    set(a, 1)
+    expect(get(b)).toBe(1)
+    expect(fetch).toHaveBeenCalledTimes(0)
+  })
+  const store = createStore()
+  store.sub(b, () => {})
+  store.sub(c, () => {})
+  fetch.mockClear()
+  store.set(w)
+  expect(fetch).toHaveBeenCalledOnce()
+  expect(fetch).toBeCalledWith(1)
+  expect(store.get(a)).toBe(1)
+})

tests/vanilla/store.test.tsx

@@ -963,3 +963,23 @@ it('processes deep atom a graph beyond maxDepth', () => {
   expect(() => store.set(baseAtom, 1)).not.toThrow()
   // store.set(lastAtom) // FIXME: This is causing a stack overflow
 })
+
+it('mounted atom should be recomputed eagerly', () => {
+  const result: string[] = []
+  const a = atom(0)
+  const b = atom((get) => {
+    result.push('bRead')
+    return get(a)
+  })
+  const store = createStore()
+  store.sub(a, () => {
+    result.push('aCallback')
+  })
+  store.sub(b, () => {
+    result.push('bCallback')
+  })
+  expect(result).toEqual(['bRead'])
+  result.splice(0)
+  store.set(a, 1)
+  expect(result).toEqual(['bRead', 'aCallback', 'bCallback'])
+})