refactor to store.unstable_onChange api

src/vanilla/store.ts

@@ -165,9 +165,19 @@ type Pending = readonly [
   dependents: Map<AnyAtom, Set<AnyAtom>>,
   atomStates: Map<AnyAtom, AtomState>,
   functions: Set<() => void>,
+  updates: readonly [
+    changed: Set<AnyAtom>,
+    mounted: Set<AnyAtom>,
+    unmounted: Set<AnyAtom>,
+  ],
 ]
 
-const createPending = (): Pending => [new Map(), new Map(), new Set()]
+const createPending = (): Pending => [
+  new Map(),
+  new Map(),
+  new Set(),
+  [new Set(), new Set(), new Set()],
+]
 
 const addPendingAtom = (
   pending: Pending,
@@ -178,6 +188,7 @@ const addPendingAtom = (
     pending[0].set(atom, new Set())
   }
   pending[1].set(atom, atomState)
+  pending[3][0].add(atom)
 }
 
 const addPendingDependent = (
@@ -198,6 +209,26 @@ const addPendingFunction = (pending: Pending, fn: () => void) => {
   pending[2].add(fn)
 }
 
+const addMountedAtom = (pending: Pending, atom: AnyAtom) => {
+  pending[3][1].add(atom)
+}
+
+const addUnmountedAtom = (pending: Pending, atom: AnyAtom) => {
+  pending[3][2].add(atom)
+}
+
+const flushPending = (pending: Pending) => {
+  while (pending[1].size || pending[2].size) {
+    pending[0].clear()
+    const atomStates = new Set(pending[1].values())
+    pending[1].clear()
+    const functions = new Set(pending[2])
+    pending[2].clear()
+    atomStates.forEach((atomState) => atomState.m?.l.forEach((l) => l()))
+    functions.forEach((fn) => fn())
+  }
+}
+
 type GetAtomState = <Value>(
   atom: Atom<Value>,
   originAtomState?: AtomState,
@@ -226,14 +257,20 @@ type PrdStore = {
   ) => Result
   sub: (atom: AnyAtom, listener: () => void) => () => void
   unstable_derive: (fn: (...args: StoreArgs) => StoreArgs) => Store
-  unstable_onAfterFlushPending: (listener: () => void) => () => void
+  unstable_onChange: (handler: OnChangeHandler) => () => void
 }
 
 export type Store = PrdStore | (PrdStore & DevStoreRev4)
 
 export type INTERNAL_DevStoreRev4 = DevStoreRev4
 export type INTERNAL_PrdStore = PrdStore
 
+type OnChangeHandler = (
+  changedAtoms: Set<AnyAtom>,
+  mountedAtoms: Set<AnyAtom>,
+  unmountedAtoms: Set<AnyAtom>,
+) => void
+
 const buildStore = (getAtomState: StoreArgs[0]): Store => {
   // for debugging purpose only
   let debugMountedAtoms: Set<AnyAtom>
@@ -242,26 +279,6 @@ const buildStore = (getAtomState: StoreArgs[0]): Store => {
     debugMountedAtoms = new Set()
   }
 
-  const flushPending = (pending: Pending, shouldProcessFinalizers = true) => {
-    do {
-      while (pending[1].size || pending[2].size) {
-        pending[0].clear()
-        const atomStates = new Set(pending[1].values())
-        pending[1].clear()
-        const functions = new Set(pending[2])
-        pending[2].clear()
-        atomStates.forEach((atomState) => atomState.m?.l.forEach((l) => l()))
-        functions.forEach((fn) => fn())
-      }
-      // Process finalizers after all atoms are updated
-      if (shouldProcessFinalizers) {
-        for (const finalizer of finalizers) {
-          finalizer()
-        }
-      }
-    } while (pending[1].size || pending[2].size)
-  }
-
   const setAtomStateValueOrPromise = (
     atom: AnyAtom,
     atomState: AtomState,
@@ -351,6 +368,7 @@ const buildStore = (getAtomState: StoreArgs[0]): Store => {
         addDependency(pending, atom, atomState, a, aState)
         mountDependencies(pending, atom, atomState)
         flushPending(pending)
+        processOnChangeHandlers(pending)
       }
       return returnAtomValue(aState)
     }
@@ -393,6 +411,7 @@ const buildStore = (getAtomState: StoreArgs[0]): Store => {
             const pending = createPending()
             mountDependencies(pending, atom, atomState)
             flushPending(pending)
+            processOnChangeHandlers(pending)
           }
         }
         valueOrPromise.then(complete, complete)
@@ -523,7 +542,10 @@ const buildStore = (getAtomState: StoreArgs[0]): Store => {
       } else {
         r = writeAtomState(pending, a, aState, ...args) as R
       }
-      flushPending(pending, !isSync)
+      flushPending(pending)
+      if (!isSync) {
+        processOnChangeHandlers(pending)
+      }
       return r as R
     }
     try {
@@ -540,6 +562,7 @@ const buildStore = (getAtomState: StoreArgs[0]): Store => {
     const pending = createPending()
     const result = writeAtomState(pending, atom, getAtomState(atom), ...args)
     flushPending(pending)
+    processOnChangeHandlers(pending)
     return result
   }
 
@@ -588,6 +611,7 @@ const buildStore = (getAtomState: StoreArgs[0]): Store => {
       if (import.meta.env?.MODE !== 'production') {
         debugMountedAtoms.add(atom)
       }
+      addMountedAtom(pending, atom)
       if (isActuallyWritableAtom(atom) && atom.onMount) {
         const mounted = atomState.m
         const { onMount } = atom
@@ -625,6 +649,7 @@ const buildStore = (getAtomState: StoreArgs[0]): Store => {
       if (import.meta.env?.MODE !== 'production') {
         debugMountedAtoms.delete(atom)
       }
+      addUnmountedAtom(pending, atom)
       // unmount dependencies
       for (const a of atomState.d.keys()) {
         const aMounted = unmountAtom(pending, a, getAtomState(a, atomState))
@@ -640,33 +665,49 @@ const buildStore = (getAtomState: StoreArgs[0]): Store => {
     const atomState = getAtomState(atom)
     const mounted = mountAtom(pending, atom, atomState)
     flushPending(pending)
+    processOnChangeHandlers(pending)
     const listeners = mounted.l
     listeners.add(listener)
     return () => {
       listeners.delete(listener)
       const pending = createPending()
       unmountAtom(pending, atom, atomState)
       flushPending(pending)
+      processOnChangeHandlers(pending)
     }
   }
 
   const unstable_derive = (fn: (...args: StoreArgs) => StoreArgs) =>
     buildStore(...fn(getAtomState))
 
-  const finalizers = new Set<() => void>()
-  const unstable_onAfterFlushPending = (listener: () => void) => {
-    finalizers.add(listener)
+  const onChangeHandlers = new Set<OnChangeHandler>()
+
+  const unstable_onChange = (handler: OnChangeHandler) => {
+    onChangeHandlers.add(handler)
     return () => {
-      finalizers.delete(listener)
+      onChangeHandlers.delete(handler)
     }
   }
 
+  const processOnChangeHandlers = (pending: Pending) => {
+    do {
+      flushPending(pending)
+      // Process finalizers after all atoms are updated
+      if (pending[3].every((s) => s.size)) {
+        for (const handler of onChangeHandlers) {
+          handler(...pending[3])
+        }
+        pending[3].forEach((s) => s.clear())
+      }
+    } while (pending[1].size || pending[2].size)
+  }
+
   const store: Store = {
     get: readAtom,
     set: writeAtom,
     sub: subscribeAtom,
     unstable_derive,
-    unstable_onAfterFlushPending,
+    unstable_onChange,
   }
   if (import.meta.env?.MODE !== 'production') {
     const devStore: DevStoreRev4 = {
@@ -698,6 +739,7 @@ const buildStore = (getAtomState: StoreArgs[0]): Store => {
           }
         }
         flushPending(pending)
+        processOnChangeHandlers(pending)
       },
     }
     Object.assign(store, devStore)

tests/vanilla/effect.test.ts

@@ -68,7 +68,7 @@ function atomSyncEffect(effect: Effect) {
     return ref
   })
   internalAtom.unstable_onInit = (store) => {
-    store.unstable_onAfterFlushPending(() => {
+    store.unstable_onChange(() => {
       const ref = store.get(refAtom)
       if (!ref.isPending || ref.inProgress > 0) {
         return