Revert "fix: remove extra onChange store hook in recomputeInvalidatedAtoms (#2982)"

This reverts commit f5d843c.

Author: dmaskasky

src/vanilla/internals.ts

@@ -457,6 +457,7 @@ const recomputeInvalidatedAtoms: RecomputeInvalidatedAtoms = (store) => {
   const mountedMap = buildingBlocks[1]
   const invalidatedAtoms = buildingBlocks[2]
   const changedAtoms = buildingBlocks[3]
+  const storeHooks = buildingBlocks[6]
   const ensureAtomState = buildingBlocks[11]
   const readAtomState = buildingBlocks[14]
   const mountDependencies = buildingBlocks[17]
@@ -465,7 +466,11 @@ const recomputeInvalidatedAtoms: RecomputeInvalidatedAtoms = (store) => {
   // 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][] = []
+  const topSortedReversed: [
+    atom: AnyAtom,
+    atomState: AtomState,
+    epochNumber: EpochNumber,
+  ][] = []
   const visiting = new WeakSet<AnyAtom>()
   const visited = new WeakSet<AnyAtom>()
   // Visit the root atom. This is the only atom in the dependency graph
@@ -484,7 +489,7 @@ const recomputeInvalidatedAtoms: RecomputeInvalidatedAtoms = (store) => {
       // performance, we will simply push onto the end, and then will iterate in
       // reverse order later.
       if (invalidatedAtoms.get(a) === aState.n) {
-        topSortedReversed.push([a, aState])
+        topSortedReversed.push([a, aState, aState.n])
       } else if (
         import.meta.env?.MODE !== 'production' &&
         invalidatedAtoms.has(a)
@@ -507,7 +512,7 @@ const recomputeInvalidatedAtoms: RecomputeInvalidatedAtoms = (store) => {
   // Step 2: use the topSortedReversed atom list to recompute all affected atoms
   // Track what's changed, so that we can short circuit when possible
   for (let i = topSortedReversed.length - 1; i >= 0; --i) {
-    const [a, aState] = topSortedReversed[i]!
+    const [a, aState, prevEpochNumber] = topSortedReversed[i]!
     let hasChangedDeps = false
     for (const dep of aState.d.keys()) {
       if (dep !== a && changedAtoms.has(dep)) {
@@ -518,6 +523,10 @@ const recomputeInvalidatedAtoms: RecomputeInvalidatedAtoms = (store) => {
     if (hasChangedDeps) {
       readAtomState(store, a)
       mountDependencies(store, a)
+      if (prevEpochNumber !== aState.n) {
+        changedAtoms.add(a)
+        storeHooks.c?.(a)
+      }
     }
     invalidatedAtoms.delete(a)
   }