refactor to use atom.init and store.unstable_onAfterFlushPending

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>(
@@ -48,9 +50,9 @@ export interface Atom<Value> {
    */
   debugPrivate?: boolean
   /**
-   * Fires after all atoms have been updated in a transaction.
+   * Fires after atom is referenced by the store for the first time
    */
-  onAfterFlushPending?: (value: Value) => void
+  unstable_onInit?: (store: Store) => void
 }
 
 export interface WritableAtom<Value, Args extends unknown[], Result>

src/vanilla/store.ts

@@ -165,23 +165,16 @@ type Pending = readonly [
   dependents: Map<AnyAtom, Set<AnyAtom>>,
   atomStates: Map<AnyAtom, AtomState>,
   functions: Set<() => void>,
-  finalizers: Set<AnyAtom>,
 ]
 
-const createPending = (): Pending => [
-  new Map(),
-  new Map(),
-  new Set(),
-  new Set(),
-]
+const createPending = (): Pending => [new Map(), new Map(), new Set()]
 
 const addPendingAtom = (
   pending: Pending,
   atom: AnyAtom,
   atomState: AtomState,
 ) => {
   if (!pending[0].has(atom)) {
-    addPendingFinalizer(pending, atom)
     pending[0].set(atom, new Set())
   }
   pending[1].set(atom, atomState)
@@ -194,7 +187,6 @@ const addPendingDependent = (
 ) => {
   const dependents = pending[0].get(atom)
   if (dependents) {
-    addPendingFinalizer(pending, dependent)
     dependents.add(dependent)
   }
 }
@@ -206,10 +198,24 @@ const addPendingFunction = (pending: Pending, fn: () => void) => {
   pending[2].add(fn)
 }
 
-const addPendingFinalizer = (pending: Pending, atom: AnyAtom) => {
-  if ('onAfterFlushPending' in atom) {
-    pending[3].add(atom)
-  }
+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)
 }
 
 type GetAtomState = <Value>(
@@ -240,9 +246,10 @@ type PrdStore = {
   ) => Result
   sub: (atom: AnyAtom, listener: () => void) => () => void
   unstable_derive: (fn: (...args: StoreArgs) => StoreArgs) => Store
+  unstable_onAfterFlushPending: (listener: () => void) => () => void
 }
 
-type Store = PrdStore | (PrdStore & DevStoreRev4)
+export type Store = PrdStore | (PrdStore & DevStoreRev4)
 
 export type INTERNAL_DevStoreRev4 = DevStoreRev4
 export type INTERNAL_PrdStore = PrdStore
@@ -255,34 +262,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) {
-        const pendingFinalizers = Array.from(pending[3])
-        for (const finalizerAtom of pendingFinalizers) {
-          const atomState = getAtomState(finalizerAtom)
-          if (atomState.m) {
-            const get = function get<Value>(a: Atom<Value>) {
-              return getAtomState(a).v!
-            }
-            finalizerAtom.onAfterFlushPending!(atomState.v)
-            pending[3].delete(finalizerAtom)
-          }
-        }
-      }
-    } while (pending[1].size || pending[2].size)
-  }
-
   const setAtomStateValueOrPromise = (
     atom: AnyAtom,
     atomState: AtomState,
@@ -621,7 +600,6 @@ const buildStore = (getAtomState: StoreArgs[0]): Store => {
           }
         })
       }
-      addPendingFinalizer(pending, atom)
     }
     return atomState.m
   }
@@ -675,11 +653,20 @@ const buildStore = (getAtomState: StoreArgs[0]): Store => {
   const unstable_derive = (fn: (...args: StoreArgs) => StoreArgs) =>
     buildStore(...fn(getAtomState))
 
+  const finalizers = new Set<() => void>()
+  const unstable_onAfterFlushPending = (listener: () => void) => {
+    finalizers.add(listener)
+    return () => {
+      finalizers.delete(listener)
+    }
+  }
+
   const store: Store = {
     get: readAtom,
     set: writeAtom,
     sub: subscribeAtom,
     unstable_derive,
+    unstable_onAfterFlushPending,
   }
   if (import.meta.env?.MODE !== 'production') {
     const devStore: DevStoreRev4 = {
@@ -725,10 +712,14 @@ 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(getAtomState)
+  const store = buildStore(getAtomState)
+  return store
 }
 
 let defaultStore: Store | undefined

tests/vanilla/effect.test.ts

@@ -67,24 +67,27 @@ function atomSyncEffect(effect: Effect) {
     ref.isPending = true
     return ref
   })
-  internalAtom.onAfterFlushPending = (ref) => {
-    if (!ref.isPending || ref.inProgress > 0) {
-      return
-    }
-    ref.isPending = false
-    ref.cleanup?.()
-    const cleanup = effectAtom.effect(ref.get!, ref.set!)
-    ref.cleanup =
-      typeof cleanup === 'function'
-        ? () => {
-            try {
-              ref.fromCleanup = true
-              cleanup()
-            } finally {
-              ref.fromCleanup = false
+  internalAtom.unstable_onInit = (store) => {
+    store.unstable_onAfterFlushPending(() => {
+      const ref = store.get(refAtom)
+      if (!ref.isPending || ref.inProgress > 0) {
+        return
+      }
+      ref.isPending = false
+      ref.cleanup?.()
+      const cleanup = effectAtom.effect(ref.get!, ref.set!)
+      ref.cleanup =
+        typeof cleanup === 'function'
+          ? () => {
+              try {
+                ref.fromCleanup = true
+                cleanup()
+              } finally {
+                ref.fromCleanup = false
+              }
             }
-          }
-        : null
+          : null
+    })
   }
   if (process.env.NODE_ENV !== 'production') {
     refAtom.debugPrivate = true