fix(store): recompute dependents on eager evaluation in writeAtomState (#2965)

* add failing test

* fix the test

* only recompute the atom we are getting

* include dirty dependencies (deep)

* only recompute if the top is dirty

* beautify

* alternate solution

---------

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

Author: dmaskasky

src/vanilla/store.ts

@@ -389,6 +389,7 @@ const buildStore = (...storeArgs: StoreArgs): Store => {
         return setSelf
       },
     }
+    const prevEpochNumber = atomState.n
     try {
       const valueOrPromise = atomRead(atom, getter, options as never)
       setAtomStateValueOrPromise(atom, atomState, valueOrPromise)
@@ -404,6 +405,14 @@ const buildStore = (...storeArgs: StoreArgs): Store => {
       return atomState
     } finally {
       isSync = false
+      if (
+        prevEpochNumber !== atomState.n &&
+        invalidatedAtoms.get(atom) === prevEpochNumber
+      ) {
+        invalidatedAtoms.set(atom, atomState.n)
+        changedAtoms.set(atom, atomState)
+        atomState.u?.()
+      }
     }
   }
 

tests/vanilla/store.test.tsx

@@ -1223,3 +1223,22 @@ it('updates with reading derived atoms (#2959)', () => {
   store.set(countUpAtom)
   expect(store.get(countDerivedAtom)).toBe(2)
 })
+
+it('updates dependents when it eagerly recomputes dirty atoms', () => {
+  const countAtom = atom(0)
+  const isActiveAtom = atom(false)
+  const activeCountAtom = atom((get) =>
+    get(isActiveAtom) ? get(countAtom) : undefined,
+  )
+  const activateAction = atom(null, (get, set, value: boolean) => {
+    set(isActiveAtom, value)
+    get(activeCountAtom)
+  })
+
+  const store = createStore()
+  store.sub(activeCountAtom, () => {})
+  store.set(activateAction, true)
+  store.set(countAtom, 1)
+
+  expect(store.get(activeCountAtom)).toBe(1)
+})