fix(core): Unmount an atom that is no longer dependent within a deriv…

…ed atom (#2660)

* Add test

* fix #2658

* refactor

---------

Co-authored-by: Daishi Kato <dai-shi@users.noreply.github.com>
Co-authored-by: daishi <daishi@axlight.com>

src/vanilla/store.ts

@@ -594,9 +594,9 @@ export const createStore = (): Store => {
       }
       for (const a of atomState.m.d || []) {
         if (!atomState.d.has(a)) {
+          atomState.m.d.delete(a)
           const aMounted = unmountAtom(pending, a)
           aMounted?.t.delete(atom)
-          atomState.m.d.delete(a)
         }
       }
     }
@@ -645,7 +645,7 @@ export const createStore = (): Store => {
     if (
       atomState.m &&
       !atomState.m.l.size &&
-      !Array.from(atomState.m.t).some((a) => getAtomState(a).m)
+      !Array.from(atomState.m.t).some((a) => getAtomState(a).m?.d.has(atom))
     ) {
       // unmount self
       const onUnmount = atomState.m.u

tests/vanilla/store.test.tsx

@@ -555,3 +555,20 @@ describe('aborting atoms', () => {
     expect(callAfterAbort).toHaveBeenCalledTimes(1)
   })
 })
+
+it('Unmount an atom that is no longer dependent within a derived atom (#2658)', async () => {
+  const condAtom = atom(true)
+
+  const baseAtom = atom(0)
+  const onUnmount = vi.fn()
+  baseAtom.onMount = () => onUnmount
+
+  const derivedAtom = atom((get) => {
+    if (get(condAtom)) get(baseAtom)
+  })
+
+  const store = createStore()
+  store.sub(derivedAtom, () => {})
+  store.set(condAtom, false)
+  expect(onUnmount).toHaveBeenCalledTimes(1)
+})