integrate comments

src/vanilla/atom.ts

@@ -1,4 +1,4 @@
-import type { AtomState, INTERNAL_PrdStore as Store } from './store'
+import type { 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, atomState: AtomState) => void
+  unstable_onInit?: (store: Store) => 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.
  */
-export type AtomState<Value = AnyValue> = {
+type AtomState<Value = AnyValue> = {
   /**
    * Map of atoms that the atom depends on.
    * The map value is the epoch number of the dependency.
@@ -200,11 +200,11 @@ const registerBatchAtom = (
 ) => {
   if (!batch.D.has(atom)) {
     batch.D.set(atom, new Set())
+    atomState.u?.(batch)
     const scheduleListeners = () => {
-      atomState.u?.(batch)
       atomState.m?.l.forEach((listener) => addBatchFunc(batch, 'M', listener))
     }
-    addBatchFunc(batch, 'H', scheduleListeners)
+    addBatchFunc(batch, 'M', scheduleListeners)
   }
 }
 
@@ -262,11 +262,7 @@ type StoreArgs = readonly [
     atom: WritableAtom<Value, Args, Result>,
     ...params: Parameters<WritableAtom<Value, Args, Result>['write']>
   ) => Result,
-  atomOnInit: <Value>(
-    atom: Atom<Value>,
-    store: Store,
-    atomState: AtomState<Value>,
-  ) => void,
+  atomOnInit: <Value>(atom: Atom<Value>, store: Store) => void,
   atomOnMount: <Value, Args extends unknown[], Result>(
     atom: WritableAtom<Value, Args, Result>,
     setAtom: (...args: Args) => Result,
@@ -312,7 +308,7 @@ const buildStore = (...storeArgs: StoreArgs): Store => {
     if (!atomState) {
       atomState = { d: new Map(), p: new Set(), n: 0 }
       atomState = setAtomState(atom, atomState)
-      atomOnInit?.(atom, store, atomState)
+      atomOnInit?.(atom, store)
     }
     return atomState
   }

tests/vanilla/atomSyncEffect.ts

@@ -53,7 +53,7 @@ export function atomSyncEffect(effect: Effect): Atom<void> {
     return ++ref.epoch
   })
   internalAtom.debugLabel = 'internal'
-  internalAtom.unstable_onInit = (store, atomState) => {
+  internalAtom.unstable_onInit = (store) => {
     console.log(':onInit')
     const ref = store.get(refAtom)
     ref.refresh = () => {
@@ -65,7 +65,7 @@ export function atomSyncEffect(effect: Effect): Atom<void> {
         ++ref.isRefreshing
       }
     }
-    ref.atomState = atomState
+    ref.atomState = getAtomState(store, internalAtom)
     ref.init = (get) => {
       console.log(':init')
       const currDeps = new Map<AnyAtom, unknown>()
@@ -131,6 +131,7 @@ export function atomSyncEffect(effect: Effect): Atom<void> {
     }
     const originalMountHook = ref.atomState.h
     ref.atomState.h = (batch) => {
+      originalMountHook?.(batch)
       console.log(':mountHook')
       if (ref.atomState.m) {
         // mount
@@ -145,10 +146,10 @@ export function atomSyncEffect(effect: Effect): Atom<void> {
           ref.cleanup = null
         })
       }
-      originalMountHook?.(batch)
     }
     const originalUpdateHook = ref.atomState.u
     ref.atomState.u = (batch) => {
+      originalUpdateHook?.(batch)
       console.log(':updateHook', ref.isRefreshing, !ref.isMounted)
       // update
       if (ref.isRefreshing || !ref.isMounted) {
@@ -161,7 +162,6 @@ export function atomSyncEffect(effect: Effect): Atom<void> {
         syncEffectChannel.add(function scheduledEffect() {
           runEffect()
         })
-        originalUpdateHook?.(batch)
       }
     }
   }
@@ -184,3 +184,22 @@ function ensureBatchChannel(ref: Ref, batch: Batch) {
   }
   return ref.batches.get(batch)!
 }
+
+/**
+ * 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((getAtomState) => {
+      atomState = getAtomState(atom)!
+      return null as any
+    })
+  } catch {
+    // expect error
+  }
+  return atomState!
+}

tests/vanilla/store.test.tsx

@@ -1012,15 +1012,8 @@ it('should call onInit only once per atom', () => {
   const onInit = vi.fn()
   a.unstable_onInit = onInit
   store.get(a)
-  type Store = ReturnType<typeof createStore>
-  let getAtomState: Parameters<Parameters<Store['unstable_derive']>[0]>[0]
-  store.unstable_derive((...storeArgs) => {
-    ;[getAtomState] = storeArgs
-    return storeArgs
-  })
-  const atomState = getAtomState!(a)
   expect(onInit).toHaveBeenCalledTimes(1)
-  expect(onInit).toHaveBeenCalledWith(store, atomState)
+  expect(onInit).toHaveBeenCalledWith(store)
   onInit.mockClear()
   store.get(a)
   store.set(a, 1)