fix(store): do not recompute unmounted atoms eagerly (#2849)

* test: it can read in write fns w/o changing deps

* add more tests

* a possible fix

* revert memory leak test

---------

Co-authored-by: David Maskasky <[email protected]>

src/vanilla/store.ts

@@ -418,14 +418,17 @@ const buildStore = (
   const readAtom = <Value>(atom: Atom<Value>): Value =>
     returnAtomValue(readAtomState(undefined, atom))
 
-  const getDependents = <Value>(
+  const getMountedOrPendingDependents = <Value>(
     pending: Pending,
     atom: Atom<Value>,
     atomState: AtomState<Value>,
   ): Map<AnyAtom, AtomState> => {
     const dependents = new Map<AnyAtom, AtomState>()
     for (const a of atomState.m?.t || []) {
-      dependents.set(a, getAtomState(a))
+      const aState = getAtomState(a)
+      if (aState.m) {
+        dependents.set(a, aState)
+      }
     }
     for (const atomWithPendingPromise of atomState.p) {
       dependents.set(
@@ -473,7 +476,7 @@ const buildStore = (
       }
       visiting.add(a)
       // Push unvisited dependents onto the stack
-      for (const [d, s] of getDependents(pending, a, aState)) {
+      for (const [d, s] of getMountedOrPendingDependents(pending, a, aState)) {
         if (a !== d && !visiting.has(d)) {
           stack.push([d, s])
         }

tests/vanilla/dependency.test.tsx

@@ -346,3 +346,62 @@ it('can read sync derived atom in write without initializing', () => {
   // note: this is why write get needs to update deps
   expect(store.get(a)).toBe(2)
 })
+
+it('can read in write function without changing dependencies', () => {
+  // https://github.com/pmndrs/jotai/discussions/2789
+  const a = atom(0)
+  let bReadCount = 0
+  const b = atom(
+    (get) => {
+      ++bReadCount
+      return get(a)
+    },
+    () => {},
+  )
+  let bIsMounted = false
+  b.onMount = () => {
+    bIsMounted = true
+  }
+  const c = atom((get) => get(a))
+  const w = atom(null, (get, set) => {
+    expect(bReadCount).toBe(0)
+    const bValue = get(b)
+    expect(bReadCount).toBe(1)
+    set(a, bValue + 1)
+    expect(bReadCount).toBe(1)
+  })
+
+  const store = createStore()
+  store.sub(c, () => {}) // mounts c,a
+  store.set(w)
+  expect(bIsMounted).toBe(false)
+})
+
+it('can cache reading an atom in write function (without mounting)', () => {
+  let aReadCount = 0
+  const a = atom(() => {
+    ++aReadCount
+    return 'a'
+  })
+  const w = atom(null, (get) => get(a))
+  const store = createStore()
+  store.set(w)
+  expect(aReadCount).toBe(1)
+  store.set(w)
+  expect(aReadCount).toBe(1)
+})
+
+it('can cache reading an atom in write function (with mounting)', () => {
+  let aReadCount = 0
+  const a = atom(() => {
+    ++aReadCount
+    return 'a'
+  })
+  const w = atom(null, (get) => get(a))
+  const store = createStore()
+  store.sub(a, () => {}) // mounts a
+  store.set(w)
+  expect(aReadCount).toBe(1)
+  store.set(w)
+  expect(aReadCount).toBe(1)
+})