add atom.unstable_onInit and store.unstable_onChange as a store.sync …

…callback mechanism after flushPending

src/vanilla/atom.ts

@@ -1,3 +1,5 @@
+import type { Store } from './store'
+
 type Getter = <Value>(atom: Atom<Value>) => Value
 
 type Setter = <Value, Args extends unknown[], Result>(
@@ -47,6 +49,10 @@ export interface Atom<Value> {
    * @private
    */
   debugPrivate?: boolean
+  /**
+   * Fires after atom is referenced by the store for the first time
+   */
+  unstable_onInit?: (store: Store) => void
 }
 
 export interface WritableAtom<Value, Args extends unknown[], Result>

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,29 +209,12 @@ const addPendingFunction = (pending: Pending, fn: () => void) => {
   pending[2].add(fn)
 }
 
-const flushPending = (pending: Pending) => {
-  let error: unknown | undefined
-  const call = (fn: () => void) => {
-    try {
-      fn()
-    } catch (e) {
-      if (!error) {
-        error = e
-      }
-    }
-  }
-  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(call))
-    functions.forEach(call)
-  }
-  if (error) {
-    throw error
-  }
+const addMountedAtom = (pending: Pending, atom: AnyAtom) => {
+  pending[3][1].add(atom)
+}
+
+const addUnmountedAtom = (pending: Pending, atom: AnyAtom) => {
+  pending[3][2].add(atom)
 }
 
 // internal & unstable type
@@ -257,13 +251,20 @@ type PrdStore = {
   ) => Result
   sub: (atom: AnyAtom, listener: () => void) => () => void
   unstable_derive: (fn: (...args: StoreArgs) => StoreArgs) => Store
+  unstable_onChange: (handler: OnChangeHandler) => () => void
 }
 
-type Store = PrdStore | (PrdStore & DevStoreRev4)
+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, atomRead, atomWrite, atomOnMount]: StoreArgs
 ): Store => {
@@ -598,6 +599,7 @@ const buildStore = (
       if (import.meta.env?.MODE !== 'production') {
         debugMountedAtoms.add(atom)
       }
+      addMountedAtom(pending, atom)
       if (isActuallyWritableAtom(atom)) {
         const mounted = atomState.m
         addPendingFunction(pending, () => {
@@ -650,6 +652,7 @@ const buildStore = (
       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))
@@ -678,11 +681,56 @@ const buildStore = (
   const unstable_derive = (fn: (...args: StoreArgs) => StoreArgs) =>
     buildStore(...fn(getAtomState, atomRead, atomWrite, atomOnMount))
 
+  const onChangeHandlers = new Set<OnChangeHandler>()
+
+  const unstable_onChange = (handler: OnChangeHandler) => {
+    onChangeHandlers.add(handler)
+    return () => {
+      onChangeHandlers.delete(handler)
+    }
+  }
+
+  const flushPending = (pending: Pending) => {
+    let ref: { e?: AnyError } = {}
+    const call = (fn: () => void) => {
+      try {
+        fn()
+      } catch (e) {
+        if (!('e' in ref)) {
+          ref.e = e
+        }
+      }
+    }
+    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(call))
+        functions.forEach(call)
+      }
+
+      // Process onChange handlers after all atoms are updated
+      if (pending[3].some((s) => s.size)) {
+        for (const handler of onChangeHandlers) {
+          handler(...pending[3])
+        }
+        pending[3].forEach((s) => s.clear())
+      }
+    } while (pending[1].size || pending[2].size)
+    if ('e' in ref) {
+      throw ref.e
+    }
+  }
+
   const store: Store = {
     get: readAtom,
     set: writeAtom,
     sub: subscribeAtom,
     unstable_derive,
+    unstable_onChange,
   }
   if (import.meta.env?.MODE !== 'production') {
     const devStore: DevStoreRev4 = {
@@ -731,15 +779,19 @@ export const createStore = (): Store => {
     if (!atomState) {
       atomState = { d: new Map(), p: new Set(), n: 0 }
       atomStateMap.set(atom, atomState)
+      if (typeof atom.unstable_onInit === 'function') {
+        atom.unstable_onInit(store)
+      }
     }
     return atomState
   }
-  return buildStore(
+  const store = buildStore(
     getAtomState,
     (atom, ...params) => atom.read(...params),
     (atom, ...params) => atom.write(...params),
     (atom, ...params) => atom.onMount?.(...params),
   )
+  return store
 }
 
 let defaultStore: Store | undefined