fix(core): update atom value with conditional dependencies (regressio…

…n with #2363) (#2534)

* chore(tests): add case against stale values to dependents

* Skip stale conditional dependents test if USE_STORE2

* Simplify out setting dataAtom

* fix store.ts

* fix store2.ts

* refactor

---------

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

src/vanilla/store.ts

@@ -397,11 +397,11 @@ export const createStore = (): Store => {
 
   const readAtomState = <Value>(
     atom: Atom<Value>,
-    force?: boolean,
+    force?: (a: AnyAtom) => boolean,
   ): AtomState<Value> => {
     // See if we can skip recomputing this atom.
     const atomState = getAtomState(atom)
-    if (!force && atomState) {
+    if (!force?.(atom) && atomState) {
       // If the atom is mounted, we can use the cache.
       // because it should have been updated by dependencies.
       if (mountedMap.has(atom)) {
@@ -415,7 +415,7 @@ export const createStore = (): Store => {
           if (a === atom) {
             return true
           }
-          const aState = readAtomState(a)
+          const aState = readAtomState(a, force)
           // Check if the atom state is unchanged, or
           // check the atom value in case only dependencies are changed
           return aState === s || isEqualAtomValue(aState, s)
@@ -442,7 +442,7 @@ export const createStore = (): Store => {
         throw new Error('no atom init')
       }
       // a !== atom
-      const aState = readAtomState(a)
+      const aState = readAtomState(a, force)
       nextDependencies.set(a, aState)
       return returnAtomValue(aState)
     }
@@ -531,6 +531,7 @@ export const createStore = (): Store => {
     // Step 2: use the topsorted 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])
+    const isMarked = (a: AnyAtom) => markedAtoms.has(a)
     for (let i = topsortedAtoms.length - 1; i >= 0; --i) {
       const a = topsortedAtoms[i]!
       const prevAtomState = getAtomState(a)
@@ -545,11 +546,12 @@ export const createStore = (): Store => {
         }
       }
       if (hasChangedDeps) {
-        const nextAtomState = readAtomState(a, true)
+        const nextAtomState = readAtomState(a, isMarked)
         if (!isEqualAtomValue(prevAtomState, nextAtomState)) {
           changedAtoms.add(a)
         }
       }
+      markedAtoms.delete(a)
     }
   }
 

src/vanilla/store2.ts

@@ -347,11 +347,11 @@ export const createStore = (): Store => {
   const readAtomState = <Value>(
     pending: Pending | undefined,
     atom: Atom<Value>,
-    force?: true,
+    force?: (a: AnyAtom) => boolean,
   ): AtomState<Value> => {
     // See if we can skip recomputing this atom.
     const atomState = getAtomState(atom)
-    if (!force && isAtomStateInitialized(atomState)) {
+    if (!force?.(atom) && isAtomStateInitialized(atomState)) {
       // If the atom is mounted, we can use the cache.
       // because it should have been updated by dependencies.
       if (atomState.m) {
@@ -364,7 +364,7 @@ export const createStore = (): Store => {
           ([a, n]) =>
             // Recursively, read the atom state of the dependency, and
             // check if the atom epoch number is unchanged
-            readAtomState(pending, a).n === n,
+            readAtomState(pending, a, force).n === n,
         )
       ) {
         return atomState
@@ -387,7 +387,7 @@ export const createStore = (): Store => {
         return returnAtomValue(aState)
       }
       // a !== atom
-      const aState = readAtomState(pending, a)
+      const aState = readAtomState(pending, a, force)
       if (isSync) {
         addDependency(pending, atom, a, aState)
       } else {
@@ -499,6 +499,7 @@ export const createStore = (): Store => {
     // Step 2: use the topsorted 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])
+    const isMarked = (a: AnyAtom) => markedAtoms.has(a)
     for (let i = topsortedAtoms.length - 1; i >= 0; --i) {
       const a = topsortedAtoms[i]!
       const aState = getAtomState(a)
@@ -512,13 +513,14 @@ export const createStore = (): Store => {
         }
       }
       if (hasChangedDeps) {
-        readAtomState(pending, a, true)
+        readAtomState(pending, a, isMarked)
         mountDependencies(pending, a, aState)
         if (!hasPrevValue || !Object.is(prevValue, aState.v)) {
           addPendingAtom(pending, a, aState)
           changedAtoms.add(a)
         }
       }
+      markedAtoms.delete(a)
     }
   }
 

tests/vanilla/dependency.test.tsx

@@ -153,3 +153,34 @@ it('keeps atoms mounted between recalculations', async () => {
     unmounted: 0,
   })
 })
+
+it('should not provide stale values to conditional dependents', () => {
+  const dataAtom = atom<number[]>([100])
+  const hasFilterAtom = atom(false)
+  const filteredAtom = atom((get) => {
+    const data = get(dataAtom)
+    const hasFilter = get(hasFilterAtom)
+    if (hasFilter) {
+      return []
+    } else {
+      return data
+    }
+  })
+  const stageAtom = atom((get) => {
+    const hasFilter = get(hasFilterAtom)
+    if (hasFilter) {
+      const filtered = get(filteredAtom)
+      return filtered.length === 0 ? 'is-empty' : 'has-data'
+    } else {
+      return 'no-filter'
+    }
+  })
+
+  const store = createStore()
+  store.sub(filteredAtom, () => undefined)
+  store.sub(stageAtom, () => undefined)
+
+  expect(store.get(stageAtom), 'should start without filter').toBe('no-filter')
+  store.set(hasFilterAtom, true)
+  expect(store.get(stageAtom), 'should update').toBe('is-empty')
+})