refactor: replace store.unstable_onChange with store.sub

src/vanilla/store.ts

@@ -194,8 +194,14 @@ const registerBatchAtom = (
 ) => {
   if (!batch.D.has(atom)) {
     batch.D.set(atom, new Set())
-    addBatchFunc(batch, 'M', () => {
-      atomState.m?.l.forEach((listener) => addBatchFunc(batch, 'M', listener))
+    addBatchFunc(batch, 'H', () => {
+      for (const listener of atomState.m?.l || []) {
+        let priority: BatchPriority = 'M'
+        if ('INTERNAL_priority' in listener) {
+          priority = listener.INTERNAL_priority as BatchPriority
+        }
+        addBatchFunc(batch, priority, listener)
+      }
     })
   }
 }
@@ -214,9 +220,39 @@ const addBatchAtomDependent = (
 const getBatchAtomDependents = (batch: Batch, atom: AnyAtom) =>
   batch.D.get(atom)
 
+const flushBatch = (batch: Batch) => {
+  let error: AnyError
+  let hasError = false
+  const call = (fn: () => void) => {
+    try {
+      fn()
+    } catch (e) {
+      if (!hasError) {
+        error = e
+        hasError = true
+      }
+    }
+  }
+  while (batch.H.size || batch.M.size || batch.L.size) {
+    batch.D.clear()
+    batch.H.forEach(call)
+    batch.H.clear()
+    batch.M.forEach(call)
+    batch.M.clear()
+    batch.L.forEach(call)
+    batch.L.clear()
+  }
+  if (hasError) {
+    throw error
+  }
+}
+
 // internal & unstable type
 type StoreArgs = readonly [
-  getAtomState: <Value>(atom: Atom<Value>) => AtomState<Value>,
+  getAtomState: <Value>(
+    atom: Atom<Value>,
+    atomOnInit?: (atom: AnyAtom) => void,
+  ) => AtomState<Value>,
   atomRead: <Value>(
     atom: Atom<Value>,
     ...params: Parameters<Atom<Value>['read']>
@@ -229,6 +265,7 @@ type StoreArgs = readonly [
     atom: WritableAtom<Value, Args, Result>,
     setAtom: (...args: Args) => Result,
   ) => OnUnmount | void,
+  atomOnInit: (store: Store) => (atom: AnyAtom) => void,
 ]
 
 // for debugging purpose only
@@ -248,23 +285,19 @@ type PrdStore = {
   ) => Result
   sub: (atom: AnyAtom, listener: () => void) => () => void
   unstable_derive: (fn: (...args: StoreArgs) => StoreArgs) => Store
-  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, atomRead, atomWrite, atomOnMount]: StoreArgs
+  ...[baseGetAtomState, atomRead, atomWrite, atomOnMount, atomOnInit]: StoreArgs
 ): Store => {
+  const getAtomState = <Value>(atom: Atom<Value>) =>
+    baseGetAtomState(atom, atomOnInit(store))
+
   // for debugging purpose only
   let debugMountedAtoms: Set<AnyAtom>
 
@@ -691,68 +724,15 @@ 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 flushBatch = (batch: Batch) => {
-    let error: AnyError
-    let hasError = false
-    const call = (fn: () => void) => {
-      try {
-        fn()
-      } catch (e) {
-        if (!hasError) {
-          error = e
-          hasError = true
-        }
-      }
-    }
-    const shouldContinue = () => ([1, 2, 3] as const).some((p) => batch[p].size)
-    do {
-      const changedAtoms = new Set<AnyAtom>(batch.D.keys())
-      while (shouldContinue()) {
-        batch.D.clear()
-        batch.H.forEach(call)
-        batch.H.clear()
-        batch.M.forEach(call)
-        batch.M.clear()
-        batch.L.forEach(call)
-        batch.L.clear()
-        const moreChangedAtoms = new Set<AnyAtom>(batch.D.keys())
-        for (const atom of moreChangedAtoms) {
-          changedAtoms.add(atom)
-        }
-      }
-      // Process onChange handlers after all atoms are updated
-      if (changedAtoms.size || batch.M.size || batch.U.size) {
-        const mountedAtoms = new Set(batch.M)
-        const unmountedAtoms = new Set(batch.U)
-        batch.M.clear()
-        batch.U.clear()
-        for (const handler of onChangeHandlers) {
-          handler(changedAtoms, mountedAtoms, unmountedAtoms)
-        }
-      }
-    } while (shouldContinue())
-    if (hasError) {
-      throw error
-    }
-  }
+    buildStore(
+      ...fn(baseGetAtomState, atomRead, atomWrite, atomOnMount, atomOnInit),
+    )
 
   const store: Store = {
     get: readAtom,
     set: writeAtom,
     sub: subscribeAtom,
     unstable_derive,
-    unstable_onChange,
   }
   if (import.meta.env?.MODE !== 'production') {
     const devStore: DevStoreRev4 = {
@@ -792,27 +772,25 @@ const buildStore = (
 
 export const createStore = (): Store => {
   const atomStateMap = new WeakMap()
-  const getAtomState = <Value>(atom: Atom<Value>) => {
+  const getAtomState: StoreArgs[0] = (atom, onInit) => {
     if (import.meta.env?.MODE !== 'production' && !atom) {
       throw new Error('Atom is undefined or null')
     }
-    let atomState = atomStateMap.get(atom) as AtomState<Value> | undefined
+    let atomState = atomStateMap.get(atom) as AtomState<any> | undefined
     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)
-      }
+      onInit?.(atom)
     }
     return atomState
   }
-  const store = buildStore(
+  return buildStore(
     getAtomState,
     (atom, ...params) => atom.read(...params),
     (atom, ...params) => atom.write(...params),
     (atom, ...params) => atom.onMount?.(...params),
+    (store) => (atom) => atom.unstable_onInit?.(store),
   )
-  return store
 }
 
 let defaultStore: Store | undefined