use atomState change and mount hooks to schedule effect and cleanup

src/vanilla/atom.ts

@@ -1,4 +1,4 @@
-import type { INTERNAL_PrdStore as Store } from './store'
+import type { AtomState, INTERNAL_PrdStore as Store } from './store'
 
 type Getter = <Value>(atom: Atom<Value>) => Value
 
@@ -53,7 +53,7 @@ export interface Atom<Value> {
    * Fires after atom is referenced by the store for the first time
    * For advanced use only and subject to change without notice.
    */
-  unstable_onInit?: (store: Store) => void
+  unstable_onInit?: (store: Store, atomState: AtomState) => void
 }
 
 export interface WritableAtom<Value, Args extends unknown[], Result>

src/vanilla/store.ts

@@ -84,7 +84,7 @@ type Mounted = {
  * Mutable atom state,
  * tracked for both mounted and unmounted atoms in a store.
  */
-type AtomState<Value = AnyValue> = {
+export type AtomState<Value = AnyValue> = {
   /**
    * Map of atoms that the atom depends on.
    * The map value is the epoch number of the dependency.
@@ -103,7 +103,7 @@ type AtomState<Value = AnyValue> = {
   /** Listener to notify when the atom value is updated. */
   u?: (batch: Batch) => void
   /** Listener to notify when the atom is mounted or unmounted. */
-  h?: () => void
+  h?: (batch: Batch) => void
   /** Atom value */
   v?: Value
   /** Atom error */
@@ -258,10 +258,10 @@ const flushBatch = (batch: Batch) => {
 // internal & unstable type
 type StoreArgs = readonly [
   getAtomState: <Value>(atom: Atom<Value>) => AtomState<Value> | undefined,
-  setAtomState: <Value, CustomAtomState extends AtomState<Value>>(
+  setAtomState: <Value>(
     atom: Atom<Value>,
-    atomState: CustomAtomState,
-  ) => CustomAtomState,
+    atomState: AtomState<Value>,
+  ) => AtomState<Value>,
   atomRead: <Value>(
     atom: Atom<Value>,
     ...params: Parameters<Atom<Value>['read']>
@@ -270,7 +270,11 @@ type StoreArgs = readonly [
     atom: WritableAtom<Value, Args, Result>,
     ...params: Parameters<WritableAtom<Value, Args, Result>['write']>
   ) => Result,
-  atomOnInit: <Value>(atom: Atom<Value>, store: Store) => void,
+  atomOnInit: <Value>(
+    atom: Atom<Value>,
+    store: Store,
+    atomState: AtomState<Value>,
+  ) => void,
   atomOnMount: <Value, Args extends unknown[], Result>(
     atom: WritableAtom<Value, Args, Result>,
     setAtom: (...args: Args) => Result,
@@ -316,7 +320,7 @@ const buildStore = (...storeArgs: StoreArgs): Store => {
     if (!atomState) {
       atomState = { d: new Map(), p: new Set(), n: 0 }
       atomState = setAtomState(atom, atomState)
-      atomOnInit?.(atom, store)
+      atomOnInit?.(atom, store, atomState)
     }
     return atomState
   }
@@ -660,7 +664,7 @@ const buildStore = (...storeArgs: StoreArgs): Store => {
         d: new Set(atomState.d.keys()),
         t: new Set(),
       }
-      atomState.h?.()
+      atomState.h?.(batch)
       if (isActuallyWritableAtom(atom)) {
         const mounted = atomState.m
         let setAtom: (...args: unknown[]) => unknown
@@ -711,7 +715,7 @@ const buildStore = (...storeArgs: StoreArgs): Store => {
         addBatchFunc(batch, 'L', () => onUnmount(batch))
       }
       delete atomState.m
-      atomState.h?.()
+      atomState.h?.(batch)
       // unmount dependencies
       for (const a of atomState.d.keys()) {
         const aMounted = unmountAtom(batch, a, ensureAtomState(a))
@@ -768,8 +772,8 @@ const deriveDevStoreRev4 = (store: Store): Store & DevStoreRev4 => {
         (atom, atomState) => {
           const newAtomState = setAtomState(atom, atomState)
           let originalMounted = newAtomState.h
-          newAtomState.h = () => {
-            originalMounted?.()
+          newAtomState.h = (batch) => {
+            originalMounted?.(batch)
             if (newAtomState.m) {
               debugMountedAtoms.add(atom)
             } else {

tests/vanilla/atomSyncEffect.ts

@@ -6,97 +6,134 @@ type AtomState = NonNullable<
   ReturnType<Parameters<Parameters<Store['unstable_derive']>[0]>[0]>
 >
 type AnyAtom = Atom<unknown>
+type Batch = Parameters<NonNullable<AtomState['u']>>[0]
 type GetterWithPeek = Getter & { peek: Getter }
 type SetterWithRecurse = Setter & { recurse: Setter }
 type Cleanup = () => void
 type Effect = (get: GetterWithPeek, set: SetterWithRecurse) => void | Cleanup
 type Ref = {
-  get?: GetterWithPeek
-  set?: SetterWithRecurse
-  cleanup?: Cleanup | null
-  fromCleanup?: boolean
-  inProgress: number
-  init?: () => void
+  get: GetterWithPeek
+  set: SetterWithRecurse
   atomState: AtomState
-  unsub?: () => void
+  batches: Map<Batch, Set<() => void>>
+  inProgress: number
+  fromCleanup?: boolean
+  isMounted: boolean
+  isRecursing: boolean
+  isRefreshing: boolean
+  init?: (get: Getter) => void
+  runEffect: () => void
+  refresh: () => void
+  cleanup?: Cleanup | null
+  epoch: number
+  error?: unknown
 }
 
-export function atomSyncEffect(effect: Effect) {
+export function atomSyncEffect(effect: Effect): Atom<void> {
   const refAtom = atom(
-    () => ({ inProgress: 0 }) as Ref,
-    (get) => {
-      const ref = get(refAtom)
-      return () => {
-        ref.cleanup?.()
-        ref.cleanup = null
-      }
-    },
+    () => ({ batches: new WeakMap(), inProgress: 0, epoch: 0 }) as Ref,
   )
   refAtom.debugLabel = 'ref'
-  refAtom.onMount = (mount) => mount()
-  const runAtom = atom({ runEffect: () => {} })
-  runAtom.debugLabel = 'run'
+  const refreshAtom = atom(0)
+  refreshAtom.debugLabel = 'refresh'
   const internalAtom = atom(function internalAtomRead(get) {
+    get(refreshAtom)
     const ref = get(refAtom)
-    ref.get = ((a) => {
-      return get(a)
-    }) as Getter & { peek: Getter }
-    ref.init!()
+    if (ref.error) {
+      const { error } = ref
+      delete ref.error
+      throw error
+    }
     if (ref.inProgress > 0) {
       return
     }
-    const runEffect = () => {
-      ref.cleanup?.()
-      const cleanup = effectAtom.effect(ref.get!, ref.set!)
-      ref.cleanup = () => {
-        try {
-          ref.fromCleanup = true
-          cleanup?.()
-        } finally {
-          ref.fromCleanup = false
-        }
-      }
+    if (!ref.isRecursing && !ref.isRefreshing) {
+      ref.init!(get)
     }
-    const tmp = atom(undefined)
-    tmp.debugLabel = 'tmp'
-    tmp.unstable_onInit = (store) => {
-      store.set(runAtom, { runEffect })
+    ref.runEffect = () => {
+      ++ref.inProgress
+      try {
+        ref.cleanup?.()
+        const cleanup = effectAtom.effect(ref.get!, ref.set!)
+        ref.cleanup = cleanup
+          ? () => {
+              try {
+                ref.fromCleanup = true
+                cleanup?.()
+              } finally {
+                ref.fromCleanup = false
+              }
+            }
+          : null
+      } catch (e) {
+        ref.error = e
+        ref.refresh()
+        --ref.inProgress
+      }
     }
-    get(tmp)
+    return ++ref.epoch
   })
   internalAtom.debugLabel = 'internal'
-  internalAtom.unstable_onInit = (store) => {
+  internalAtom.unstable_onInit = (store, atomState) => {
     const ref = store.get(refAtom)
-    ref.unsub = store.sub(runAtom, () => {
-      const { runEffect } = store.get(runAtom)
-      runEffect()
-    })
+    ref.refresh = () => {
+      try {
+        ref.isRefreshing = true
+        ref.set(refreshAtom, (v) => v + 1)
+      } finally {
+        ref.isRefreshing = false
+      }
+    }
     if (!ref.atomState) {
-      ref.atomState = getAtomState(store, internalAtom)!
-      const originalHook = ref.atomState.h
-      ref.atomState.h = () => {
+      ref.atomState = atomState
+      const originalMountHook = ref.atomState.h
+      ref.atomState.h = (batch) => {
         if (ref.atomState.m) {
-          // TODO - schedule effect
+          // mount
+          ref.isMounted = true
         } else {
-          // TODO - schedule cleanup
+          // unmount
+          ref.isMounted = false
+          const syncEffectChannel = ensureBatchChannel(ref, batch)
+          syncEffectChannel.add(() => {
+            ref.cleanup?.()
+            ref.cleanup = null
+          })
         }
-        originalHook?.()
+        originalMountHook?.(batch)
       }
-    }
-    const get = store.get
-    const set = store.set
-    ref.init = () => {
-      if (!ref.get!.peek) {
-        ref.get!.peek = get
+      const originalUpdateHook = ref.atomState.u
+      ref.atomState.u = (batch) => {
+        // update
+        if (ref.isRefreshing || !ref.isMounted) {
+          return
+        }
+        if (ref.isRecursing) {
+          ref.runEffect()
+        } else {
+          const syncEffectChannel = ensureBatchChannel(ref, batch)
+          syncEffectChannel.add(ref.runEffect)
+          originalUpdateHook?.(batch)
+        }
       }
+    }
+    ref.init = (get) => {
+      const currDeps = new Map<AnyAtom, unknown>()
+      const getter = ((a) => {
+        const value = get(a)
+        currDeps.set(a, value)
+        return value
+      }) as GetterWithPeek
+      const peek = store.get
+      ref.get = Object.assign(getter, { peek })
       if (!ref.set) {
         const setter: Setter = (a, ...args) => {
           try {
             ++ref.inProgress
-            return set(a, ...args)
+            return store.set(a, ...args)
           } finally {
+            Array.from(currDeps.keys(), get) // FIXME: why do we need this?
             --ref.inProgress
-            ref.get!(a) // FIXME why do we need this?
           }
         }
         const recurse: Setter = (a, ...args) => {
@@ -106,7 +143,18 @@ export function atomSyncEffect(effect: Effect) {
             }
             return undefined as never
           }
-          return set(a, ...args)
+          try {
+            ref.isRecursing = true
+            return store.set(a, ...args)
+          } finally {
+            ref.isRecursing = false
+            const depsChanged = Array.from(currDeps).some(
+              ([a, v]) => get(a) !== v,
+            )
+            if (depsChanged) {
+              ref.refresh()
+            }
+          }
         }
         ref.set = Object.assign(setter, { recurse })
       }
@@ -121,21 +169,13 @@ export function atomSyncEffect(effect: Effect) {
   return effectAtom
 }
 
-/**
- * HACK: steal atomState to synchronously determine if
- * the atom is mounted
- * We return void 0 to cause the buildStore(...args) to throw
- * to abort creating a derived store
- */
-function getAtomState(store: Store, atom: AnyAtom): AtomState {
-  let atomState: AtomState
-  try {
-    store.unstable_derive(function deriveExtractAtomState(getAtomState) {
-      atomState = getAtomState(atom)!
-      return null as any
-    })
-  } catch {
-    // expect error
+function ensureBatchChannel(ref: Ref, batch: Batch) {
+  if (!ref.batches.has(batch)) {
+    const syncEffectChannel = new Set<() => void>()
+    ref.batches.set(batch, syncEffectChannel)
+    const syncEffectIndex =
+      batch.findIndex((channel) => channel === batch.H) + 1
+    batch.splice(syncEffectIndex, 0, syncEffectChannel)
   }
-  return atomState!
+  return ref.batches.get(batch)!
 }