feat: add hook for atom mount and unmount

pmndrs · Dec 30, 2024 · 15151a4 · 15151a4
1 parent 65a8500
commit 15151a4
Showing 1 changed file with 19 additions and 20 deletions.
diff --git a/src/vanilla/store.ts b/src/vanilla/store.ts
@@ -100,6 +100,8 @@ type AtomState<Value = AnyValue> = {
   n: number
   /** Object to store mounted state of the atom. */
   m?: Mounted // only available if the atom is mounted
+  /** Listeners to notify when the atom is mounted or unmounted. */
+  ml?: Set<<Value>(atom: Atom<Value>, atomState: AtomState<Value>) => void>
   /** Atom value */
   v?: Value
   /** Atom error */
@@ -625,6 +627,7 @@ const buildStore = (
         d: new Set(atomState.d.keys()),
         t: new Set(),
       }
+      atomState.ml?.forEach((listener) => listener(atom, atomState))
       if (isActuallyWritableAtom(atom)) {
         const mounted = atomState.m
         let setAtom: (...args: unknown[]) => unknown
@@ -674,6 +677,7 @@ const buildStore = (
         addBatchFunc(batch, 'L', () => onUnmount(batch))
       }
       delete atomState.m
+      atomState.ml?.forEach((listener) => listener(atom, atomState))
       // unmount dependencies
       for (const a of atomState.d.keys()) {
         const aMounted = unmountAtom(batch, a, getAtomState(a))
@@ -712,35 +716,30 @@ const buildStore = (
 }
 
 const deriveDevStoreRev4 = (store: Store): Store & DevStoreRev4 => {
-  const proxyAtomStateMap = new WeakMap()
   const debugMountedAtoms = new Set<AnyAtom>()
   let savedGetAtomState: StoreArgs[0]
   let inRestoreAtom = 0
+  const updateDebugMountedAtoms = (atom: AnyAtom, atomState: AtomState) => {
+    if (atomState.m) {
+      debugMountedAtoms.add(atom)
+    } else {
+      debugMountedAtoms.delete(atom)
+    }
+  }
   const derivedStore = store.unstable_derive(
     (getAtomState, atomRead, atomWrite, atomOnMount) => {
       savedGetAtomState = getAtomState
       return [
         (atom) => {
-          let proxyAtomState = proxyAtomStateMap.get(atom)
-          if (!proxyAtomState) {
-            const atomState = getAtomState(atom)
-            proxyAtomState = new Proxy(atomState, {
-              set(target, prop, value) {
-                if (prop === 'm') {
-                  debugMountedAtoms.add(atom)
-                }
-                return Reflect.set(target, prop, value)
-              },
-              deleteProperty(target, prop) {
-                if (prop === 'm') {
-                  debugMountedAtoms.delete(atom)
-                }
-                return Reflect.deleteProperty(target, prop)
-              },
-            })
-            proxyAtomStateMap.set(atom, proxyAtomState)
+          const atomState = getAtomState(atom)
+          if (!atomState.ml) {
+            atomState.ml = new Set()
+          }
+          if (!atomState.ml.has(updateDebugMountedAtoms)) {
+            atomState.ml.add(updateDebugMountedAtoms)
           }
-          return proxyAtomState
+          updateDebugMountedAtoms(atom, atomState)
+          return atomState
         },
         atomRead,
         (atom, getter, setter, ...args) => {