feat: add hook for atom mount and unmount

pmndrs · Dec 30, 2024 · a7ffc4e · a7ffc4e
1 parent 65a8500
commit a7ffc4e
Showing 1 changed file with 14 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
+  /** Listener to notify when the atom is mounted or unmounted. */
+  h?: () => void
   /** Atom value */
   v?: Value
   /** Atom error */
@@ -625,6 +627,7 @@ const buildStore = (
         d: new Set(atomState.d.keys()),
         t: new Set(),
       }
+      atomState.h?.()
       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.h?.()
       // unmount dependencies
       for (const a of atomState.d.keys()) {
         const aMounted = unmountAtom(batch, a, getAtomState(a))
@@ -712,7 +716,6 @@ const buildStore = (
 }
 
 const deriveDevStoreRev4 = (store: Store): Store & DevStoreRev4 => {
-  const proxyAtomStateMap = new WeakMap()
   const debugMountedAtoms = new Set<AnyAtom>()
   let savedGetAtomState: StoreArgs[0]
   let inRestoreAtom = 0
@@ -721,26 +724,17 @@ const deriveDevStoreRev4 = (store: Store): Store & DevStoreRev4 => {
       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)
+          let originalMounted = atomState.h
+          atomState.h = () => {
+            originalMounted?.()
+            if (atomState.m) {
+              debugMountedAtoms.add(atom)
+            } else {
+              debugMountedAtoms.delete(atom)
+            }
           }
-          return proxyAtomState
+          return atomState
         },
         atomRead,
         (atom, getter, setter, ...args) => {