From 644ba5eaac5f474fde7990f9c15e85e693a19e0f Mon Sep 17 00:00:00 2001
From: David Maskasky <dmaskasky@gmail.com>
Date: Fri, 31 Jan 2025 18:30:28 -0800
Subject: [PATCH] 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 <daishi@axlight.com>
---
 src/vanilla/store.ts         |  9 +++++++++
 tests/vanilla/store.test.tsx | 19 +++++++++++++++++++
 2 files changed, 28 insertions(+)

diff --git a/src/vanilla/store.ts b/src/vanilla/store.ts
index 15276af798..e2907a85b9 100644
--- a/src/vanilla/store.ts
+++ b/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?.()
+      }
     }
   }
 
diff --git a/tests/vanilla/store.test.tsx b/tests/vanilla/store.test.tsx
index 062938dce5..25c9e29f8b 100644
--- a/tests/vanilla/store.test.tsx
+++ b/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)
+})