add batch to internal callback

src/vanilla/store.ts

@@ -62,6 +62,8 @@ const isPromiseLike = (
 ): x is PromiseLike<unknown> & { onCancel?: (fn: CancelHandler) => void } =>
   typeof (x as any)?.then === 'function'
 
+type BatchListener = (batch: Batch) => void
+
 /**
  * State tracked for mounted atoms. An atom is considered "mounted" if it has a
  * subscriber, or is a transitive dependency of another atom that has a
@@ -77,7 +79,7 @@ type Mounted = {
   /** Set of mounted atoms that depends on the atom. */
   readonly t: Set<AnyAtom>
   /** Function to run when the atom is unmounted. */
-  u?: (batch: Batch) => void
+  u?: BatchListener
 }
 
 /**
@@ -91,7 +93,7 @@ export type AtomState<Value = AnyValue> = {
    */
   readonly d: Map<AnyAtom, number>
   /** Set of priority listeners to run when the atom value changes. */
-  l?: Set<readonly [listener: () => void, priority?: BatchPriority]>
+  l?: Set<readonly [listener: BatchListener, priority?: BatchPriority]>
   /**
    * Set of atoms with pending promise that depend on the atom.
    *
@@ -171,11 +173,11 @@ type Batch = Readonly<{
   /** Atom dependents map */
   D: Map<AnyAtom, Set<AnyAtom>>
   /** High priority functions */
-  H: Set<() => void>
+  H: Set<BatchListener>
   /** Medium priority functions */
-  M: Set<() => void>
+  M: Set<BatchListener>
   /** Low priority functions */
-  L: Set<() => void>
+  L: Set<BatchListener>
 }>
 
 const createBatch = (): Batch => ({
@@ -187,7 +189,7 @@ const createBatch = (): Batch => ({
 
 const addBatchFunc = (
   batch: Batch,
-  fn: () => void,
+  fn: BatchListener,
   priority: BatchPriority,
 ) => {
   batch[priority].add(fn)
@@ -205,7 +207,7 @@ const registerBatchAtom = (
         addBatchFunc(batch, listener, priority)
       }
       for (const listener of atomState.m?.l || []) {
-        addBatchFunc(batch, listener, 'M')
+        addBatchFunc(batch, () => listener(), 'M')
       }
     }
     addBatchFunc(batch, scheduleListeners, 'H')
@@ -229,9 +231,9 @@ const getBatchAtomDependents = (batch: Batch, atom: AnyAtom) =>
 const flushBatch = (batch: Batch) => {
   let error: AnyError
   let hasError = false
-  const call = (fn: () => void) => {
+  const call = (fn: BatchListener) => {
     try {
-      fn()
+      fn(batch)
     } catch (e) {
       if (!hasError) {
         error = e

tests/vanilla/effect.test.ts

@@ -19,6 +19,13 @@ type Ref = {
   epoch: number
 }
 
+type INTERNAL_onInit = NonNullable<AnyAtom['INTERNAL_onInit']>
+type AtomState = Parameters<INTERNAL_onInit>[1]
+type BatchListeners = NonNullable<AtomState['l']>
+type BatchEntry = BatchListeners extends Set<infer U> ? U : never
+type BatchListener = BatchEntry[0]
+type BatchPriority = NonNullable<BatchEntry[1]>
+
 function atomSyncEffect(effect: Effect) {
   const refAtom = atom(
     () => ({ deps: new Set(), inProgress: 0, epoch: 0 }) as Ref,
@@ -58,31 +65,32 @@ function atomSyncEffect(effect: Effect) {
   )
   refAtom.onMount = (mount) => mount()
   const refreshAtom = atom(0)
-  const internalAtom = atom((get) => {
-    get(refreshAtom)
-    const ref = get(refAtom)
-    if (!ref.get) {
-      ref.get = ((a) => {
-        ref.deps.add(a)
-        return get(a)
-      }) as Getter & { peak: Getter }
-    }
-    ref.deps.forEach(get)
-    ref.isPending = true
-    return ++ref.epoch
-  })
-  const bridgeAtom = atom(
-    (get) => get(internalAtom),
+  const internalAtom = atom(
+    (get) => {
+      get(refreshAtom)
+      const ref = get(refAtom)
+      if (!ref.get) {
+        ref.get = ((a) => {
+          ref.deps.add(a)
+          return get(a)
+        }) as Getter & { peak: Getter }
+      }
+      ref.deps.forEach(get)
+      ref.isPending = true
+      return ++ref.epoch
+    },
     (get, set) => {
       set(refreshAtom, (v) => ++v)
       return get(refAtom).sub()
     },
   )
-  bridgeAtom.onMount = (mount) => mount()
-  bridgeAtom.INTERNAL_onInit = (store, atomState) => {
-    atomState.l ||= new Set()
+  internalAtom.onMount = (mount) => mount()
+  internalAtom.INTERNAL_onInit = (store, atomState) => {
+    if (!('l' in atomState)) {
+      atomState.l = new Set()
+    }
     store.get(refAtom).sub = function subscribe() {
-      function listener() {
+      const batchListener: BatchListener = (_batch) => {
         const ref = store.get(refAtom)
         if (!ref.isPending || ref.inProgress > 0) {
           return
@@ -102,13 +110,14 @@ function atomSyncEffect(effect: Effect) {
               }
             : null
       }
-      const entry = [listener, 'H'] as const
+      const priority: BatchPriority = 'H'
+      const entry = [batchListener, priority] as const
       atomState.l!.add(entry)
       return () => atomState.l!.delete(entry)
     }
   }
   const effectAtom = Object.assign(
-    atom((get) => void get(bridgeAtom)),
+    atom((get) => void get(internalAtom)),
     { effect },
   )
   return effectAtom

tests/vanilla/store.test.tsx

@@ -1034,13 +1034,13 @@ it('should call onInit only once per store', () => {
   const a = atom(0)
   type AtomState = Parameters<NonNullable<Atom<unknown>['INTERNAL_onInit']>>[1]
   let aAtomState: AtomState
-  const aOnInit = vi.fn((_store, atomState) => {
+  const aOnInit = vi.fn((_store: Store, atomState: AtomState) => {
     aAtomState = atomState
   })
   a.INTERNAL_onInit = aOnInit
   const b = atom(0)
   let bAtomState: AtomState
-  const bOnInit = vi.fn((_store, atomState) => {
+  const bOnInit = vi.fn((_store: Store, atomState: AtomState) => {
     bAtomState = atomState
   })
   b.INTERNAL_onInit = bOnInit