From 1e994b7bae524276394e738f45370abdb84e1b48 Mon Sep 17 00:00:00 2001
From: Arjun <14841132+arjunvegda@users.noreply.github.com>
Date: Sun, 2 Jun 2024 10:39:05 -0400
Subject: [PATCH 1/2] feat: add dev4_get_mounted_atoms

---
 src/vanilla/store2.ts           | 14 +++++++++++
 tests/vanilla/storedev.test.tsx | 44 +++++++++++++++++++++++++++++++++
 2 files changed, 58 insertions(+)

diff --git a/src/vanilla/store2.ts b/src/vanilla/store2.ts
index d83f7bd069..16431c71df 100644
--- a/src/vanilla/store2.ts
+++ b/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: () => IterableIterator<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.values(),
       dev4_restore_atoms: (values) => {
         const pending = createPending()
         for (const [atom, value] of values) {
diff --git a/tests/vanilla/storedev.test.tsx b/tests/vanilla/storedev.test.tsx
index 7e7402e829..5abfa4c497 100644
--- a/tests/vanilla/storedev.test.tsx
+++ b/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([])
+  })
 })

From b10f78126fa5a941aefa25284ea6d9950c46d7a2 Mon Sep 17 00:00:00 2001
From: Arjun <14841132+arjunvegda@users.noreply.github.com>
Date: Mon, 3 Jun 2024 08:03:38 -0400
Subject: [PATCH 2/2] refactor: return a set from dev4_get_mounted_atoms

---
 src/vanilla/store2.ts | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/vanilla/store2.ts b/src/vanilla/store2.ts
index 16431c71df..39d9697863 100644
--- a/src/vanilla/store2.ts
+++ b/src/vanilla/store2.ts
@@ -243,7 +243,7 @@ const flushPending = (pending: Pending) => {
 // for debugging purpose only
 type DevStoreRev4 = {
   dev4_get_internal_weak_map: () => WeakMap<AnyAtom, AtomState>
-  dev4_get_mounted_atoms: () => IterableIterator<AnyAtom>
+  dev4_get_mounted_atoms: () => Set<AnyAtom>
   dev4_restore_atoms: (values: Iterable<readonly [AnyAtom, AnyValue]>) => void
 }
 
@@ -688,7 +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.values(),
+      dev4_get_mounted_atoms: () => debugMountedAtoms,
       dev4_restore_atoms: (values) => {
         const pending = createPending()
         for (const [atom, value] of values) {