feat(store2): add dev4_get_mounted_atoms (#2594)

* feat: add dev4_get_mounted_atoms

* refactor: return a set from dev4_get_mounted_atoms

---------

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

src/vanilla/store2.ts

@@ -243,6 +243,7 @@ const flushPending = (pending: Pending) => {
 // for debugging purpose only
 type DevStoreRev4 = {
   dev4_get_internal_weak_map: () => WeakMap<AnyAtom, AtomState>
+  dev4_get_mounted_atoms: () => Set<AnyAtom>
   dev4_restore_atoms: (values: Iterable<readonly [AnyAtom, AnyValue]>) => void
 }
 
@@ -261,6 +262,12 @@ export type INTERNAL_PrdStore = PrdStore
 
 export const createStore = (): Store => {
   const atomStateMap = new WeakMap<AnyAtom, AtomState>()
+  // for debugging purpose only
+  let debugMountedAtoms: Set<AnyAtom>
+
+  if (import.meta.env?.MODE !== 'production') {
+    debugMountedAtoms = new Set()
+  }
 
   const getAtomState = <Value>(atom: Atom<Value>) => {
     let atomState = atomStateMap.get(atom) as AtomState<Value> | undefined
@@ -606,6 +613,9 @@ export const createStore = (): Store => {
         d: new Set(atomState.d.keys()),
         t: new Set(),
       }
+      if (import.meta.env?.MODE !== 'production') {
+        debugMountedAtoms.add(atom)
+      }
       if (isActuallyWritableAtom(atom) && atom.onMount) {
         const mounted = atomState.m
         const { onMount } = atom
@@ -638,6 +648,9 @@ export const createStore = (): Store => {
         addPendingFunction(pending, onUnmount)
       }
       delete atomState.m
+      if (import.meta.env?.MODE !== 'production') {
+        debugMountedAtoms.delete(atom)
+      }
       // unmount dependencies
       for (const a of atomState.d.keys()) {
         const aMounted = unmountAtom(pending, a)
@@ -675,6 +688,7 @@ export const createStore = (): Store => {
       sub: subscribeAtom,
       // store dev methods (these are tentative and subject to change without notice)
       dev4_get_internal_weak_map: () => atomStateMap,
+      dev4_get_mounted_atoms: () => debugMountedAtoms,
       dev4_restore_atoms: (values) => {
         const pending = createPending()
         for (const [atom, value] of values) {

tests/vanilla/storedev.test.tsx

@@ -1,5 +1,6 @@
 import { describe, expect, it, vi } from 'vitest'
 import { atom, createStore } from 'jotai/vanilla'
+import { INTERNAL_DevStoreRev4, INTERNAL_PrdStore } from 'jotai/vanilla/store2'
 
 const IS_DEV_STORE = 'dev_subscribe_store' in createStore()
 const IS_DEV_STORE2 = 'dev4_get_internal_weak_map' in createStore()
@@ -172,4 +173,47 @@ describe.skipIf(!IS_DEV_STORE2)('[DEV-ONLY] dev-only methods rev4', () => {
     unsubCount()
     unsubDerived()
   })
+
+  it('should return all the mounted atoms correctly', () => {
+    const store = createStore() as INTERNAL_DevStoreRev4 & INTERNAL_PrdStore
+    if (!('dev4_get_mounted_atoms' in store)) {
+      throw new Error('dev methods are not available')
+    }
+    const countAtom = atom(0)
+    countAtom.debugLabel = 'countAtom'
+    const derivedAtom = atom((get) => get(countAtom) * 2)
+    const unsub = store.sub(derivedAtom, vi.fn())
+    store.set(countAtom, 1)
+    const result = store.dev4_get_mounted_atoms()
+    expect(
+      Array.from(result).sort(
+        (a, b) => Object.keys(a).length - Object.keys(b).length,
+      ),
+    ).toStrictEqual([
+      { toString: expect.any(Function), read: expect.any(Function) },
+      {
+        toString: expect.any(Function),
+        init: 0,
+        read: expect.any(Function),
+        write: expect.any(Function),
+        debugLabel: 'countAtom',
+      },
+    ])
+    unsub()
+  })
+
+  it("should return all the mounted atoms correctly after they're unsubscribed", () => {
+    const store = createStore() as INTERNAL_DevStoreRev4 & INTERNAL_PrdStore
+    if (!('dev4_get_mounted_atoms' in store)) {
+      throw new Error('dev methods are not available')
+    }
+    const countAtom = atom(0)
+    countAtom.debugLabel = 'countAtom'
+    const derivedAtom = atom((get) => get(countAtom) * 2)
+    const unsub = store.sub(derivedAtom, vi.fn())
+    store.set(countAtom, 1)
+    unsub()
+    const result = store.dev4_get_mounted_atoms()
+    expect(Array.from(result)).toStrictEqual([])
+  })
 })