diff --git a/src/vanilla/store.ts b/src/vanilla/store.ts index 395ff345c2..81ef5702af 100644 --- a/src/vanilla/store.ts +++ b/src/vanilla/store.ts @@ -100,6 +100,8 @@ type AtomState = { 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() 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) + const 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) => {