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

tests/vanilla/unstable_derive.test.tsx

@@ -2,6 +2,9 @@ import { describe, expect, it, vi } from 'vitest'
 import { atom, createStore } from 'jotai/vanilla'
 import type { Atom, Getter } from 'jotai/vanilla'
 
+type Store = ReturnType<typeof createStore>
+type StoreArgs = ReturnType<Store['unstable_derive']>
+
 describe('unstable_derive for scoping atoms', () => {
   /**
    * a
@@ -13,23 +16,28 @@ describe('unstable_derive for scoping atoms', () => {
     const scopedAtoms = new Set<Atom<unknown>>([a])
 
     const store = createStore()
-    const derivedStore = store.unstable_derive((getAtomState, ...args) => {
-      const scopedAtomStateMap = new WeakMap()
-      return [
-        (atom) => {
-          if (scopedAtoms.has(atom)) {
-            let atomState = scopedAtomStateMap.get(atom)
-            if (!atomState) {
-              atomState = { d: new Map(), p: new Set(), n: 0 }
-              scopedAtomStateMap.set(atom, atomState)
+    const derivedStore = store.unstable_derive(
+      (getAtomState, atomRead, atomWrite, atomOnMount, atomOnInit) => {
+        const scopedAtomStateMap = new WeakMap()
+        return [
+          (atom) => {
+            if (scopedAtoms.has(atom)) {
+              let atomState = scopedAtomStateMap.get(atom)
+              if (!atomState) {
+                atomState = { d: new Map(), p: new Set(), n: 0 }
+                scopedAtomStateMap.set(atom, atomState)
+              }
+              return atomState
             }
-            return atomState
-          }
-          return getAtomState(atom)
-        },
-        ...args,
-      ]
-    })
+            return getAtomState(atom)
+          },
+          atomRead,
+          atomWrite,
+          atomOnMount,
+          atomOnInit,
+        ]
+      },
+    )
 
     expect(store.get(a)).toBe('a')
     expect(derivedStore.get(a)).toBe('a')
@@ -54,23 +62,28 @@ describe('unstable_derive for scoping atoms', () => {
     const scopedAtoms = new Set<Atom<unknown>>([a])
 
     const store = createStore()
-    const derivedStore = store.unstable_derive((getAtomState, ...args) => {
-      const scopedAtomStateMap = new WeakMap()
-      return [
-        (atom) => {
-          if (scopedAtoms.has(atom)) {
-            let atomState = scopedAtomStateMap.get(atom)
-            if (!atomState) {
-              atomState = { d: new Map(), p: new Set(), n: 0 }
-              scopedAtomStateMap.set(atom, atomState)
+    const derivedStore = store.unstable_derive(
+      (getAtomState, atomRead, atomWrite, atomOnMount, atomOnInit) => {
+        const scopedAtomStateMap = new WeakMap()
+        return [
+          (atom) => {
+            if (scopedAtoms.has(atom)) {
+              let atomState = scopedAtomStateMap.get(atom)
+              if (!atomState) {
+                atomState = { d: new Map(), p: new Set(), n: 0 }
+                scopedAtomStateMap.set(atom, atomState)
+              }
+              return atomState
             }
-            return atomState
-          }
-          return getAtomState(atom)
-        },
-        ...args,
-      ]
-    })
+            return getAtomState(atom)
+          },
+          atomRead,
+          atomWrite,
+          atomOnMount,
+          atomOnInit,
+        ]
+      },
+    )
 
     expect(store.get(c)).toBe('ab')
     expect(derivedStore.get(c)).toBe('ab')
@@ -95,7 +108,7 @@ describe('unstable_derive for scoping atoms', () => {
     function makeStores() {
       const store = createStore()
       const derivedStore = store.unstable_derive(
-        (getAtomState, atomRead, ...args) => {
+        (getAtomState, atomRead, atomWrite, atomOnMount, atomOnInit) => {
           const scopedAtomStateMap = new WeakMap()
           return [
             (atom) => {
@@ -118,8 +131,10 @@ describe('unstable_derive for scoping atoms', () => {
               }
               return atomRead(a, myGet, options)
             },
-            ...args,
-          ]
+            atomWrite,
+            atomOnMount,
+            atomOnInit,
+          ] as const
         },
       )
       expect(store.get(b)).toBe('a')
@@ -173,14 +188,19 @@ describe('unstable_derive for scoping atoms', () => {
 it('should pass the correct store instance to the atom initializer', () => {
   expect.assertions(2)
   const baseStore = createStore()
-  const derivedStore = baseStore.unstable_derive((getAtomState, ...args) => [
-    (a, atomOnInit) => {
-      const atomState = getAtomState(a)
-      atomOnInit?.(a, atomState)
-      return atomState
-    },
-    ...args,
-  ])
+  const derivedStore = baseStore.unstable_derive(
+    (getAtomState, atomRead, atomWrite, atomOnMount, atomOnInit) => [
+      (a, atomOnInit) => {
+        const atomState = getAtomState(a)
+        atomOnInit?.(a, atomState)
+        return atomState
+      },
+      atomRead,
+      atomWrite,
+      atomOnMount,
+      atomOnInit,
+    ],
+  )
   const a = atom(null)
   a.INTERNAL_onInit = (store) => {
     expect(store).toBe(baseStore)