swap listener and priority

src/vanilla/store.ts

@@ -90,8 +90,8 @@ export type AtomState<Value = AnyValue> = {
    * The map value is the epoch number of the dependency.
    */
   readonly d: Map<AnyAtom, number>
-  /** Set of listeners to run when the atom value changes. */
-  l?: Set<readonly [priority: BatchPriority, function: () => void]>
+  /** Set of priority listeners to run when the atom value changes. */
+  l?: Set<readonly [listener: () => void, priority: BatchPriority]>
   /**
    * Set of atoms with pending promise that depend on the atom.
    *
@@ -185,8 +185,12 @@ const createBatch = (): Batch => ({
   L: new Set(),
 })
 
-const addBatchFunc = (batch: Batch, p: BatchPriority, fn: () => void) => {
-  batch[p].add(fn)
+const addBatchFunc = (
+  batch: Batch,
+  fn: () => void,
+  priority: BatchPriority,
+) => {
+  batch[priority].add(fn)
 }
 
 const registerBatchAtom = (
@@ -196,14 +200,15 @@ const registerBatchAtom = (
 ) => {
   if (!batch.D.has(atom)) {
     batch.D.set(atom, new Set())
-    addBatchFunc(batch, 'H', () => {
-      for (const [priority, listener] of atomState.l || []) {
-        addBatchFunc(batch, priority, listener)
+    const scheduleListeners = () => {
+      for (const [listener, priority = 'M'] of atomState.l || []) {
+        addBatchFunc(batch, listener, priority)
       }
       for (const listener of atomState.m?.l || []) {
-        addBatchFunc(batch, 'M', listener)
+        addBatchFunc(batch, listener, 'M')
       }
-    })
+    }
+    addBatchFunc(batch, scheduleListeners, 'H')
   }
 }
 
@@ -248,11 +253,17 @@ const flushBatch = (batch: Batch) => {
   }
 }
 
+type AtomOnInit = <Value>(
+  atom: Atom<Value>,
+  atomState: AtomState<Value>,
+) => void
+
 // internal & unstable type
 type StoreArgs = readonly [
-  getAtomState: (
-    atomOnInit: ReturnType<StoreArgs[4]>,
-  ) => <Value>(atom: Atom<Value>) => AtomState<Value>,
+  getAtomState: <Value>(
+    atom: Atom<Value>,
+    atomOnInit?: ReturnType<StoreArgs[4]>,
+  ) => AtomState<Value>,
   atomRead: <Value>(
     atom: Atom<Value>,
     ...params: Parameters<Atom<Value>['read']>
@@ -265,9 +276,7 @@ type StoreArgs = readonly [
     atom: WritableAtom<Value, Args, Result>,
     setAtom: (...args: Args) => Result,
   ) => OnUnmount | void,
-  atomOnInit: (
-    store: Store,
-  ) => <Value>(atom: Atom<Value>, atomState: AtomState<Value>) => void,
+  createAtomOnInit: (store: Store) => AtomOnInit,
 ]
 
 // for debugging purpose only
@@ -294,9 +303,16 @@ export type Store = PrdStore | (PrdStore & DevStoreRev4)
 export type INTERNAL_DevStoreRev4 = DevStoreRev4
 export type INTERNAL_PrdStore = PrdStore
 
-const buildStore = (
-  ...[baseGetAtomState, atomRead, atomWrite, atomOnMount, atomOnInit]: StoreArgs
-): Store => {
+const buildStore = (...storeArgs: StoreArgs): Store => {
+  const [
+    createGetAtomState,
+    atomRead,
+    atomWrite,
+    atomOnMount,
+    createAtomOnInit,
+  ] = storeArgs
+  const getAtomState = <Value>(atom: Atom<Value>) =>
+    createGetAtomState(atom, createAtomOnInit(store))
   // for debugging purpose only
   let debugMountedAtoms: Set<AnyAtom>
 
@@ -520,7 +536,7 @@ const buildStore = (
 
     // Step 2: use the topSortedReversed atom list to recompute all affected atoms
     // Track what's changed, so that we can short circuit when possible
-    addBatchFunc(batch, 'H', () => {
+    const finishRecompute = () => {
       const changedAtoms = new Set<AnyAtom>([atom])
       for (let i = topSortedReversed.length - 1; i >= 0; --i) {
         const [a, aState, prevEpochNumber] = topSortedReversed[i]!
@@ -541,7 +557,8 @@ const buildStore = (
         }
         delete aState.x
       }
-    })
+    }
+    addBatchFunc(batch, finishRecompute, 'H')
   }
 
   const writeAtomState = <Value, Args extends unknown[], Result>(
@@ -665,14 +682,15 @@ const buildStore = (
             isSync = false
           }
         }
-        addBatchFunc(batch, 'L', () => {
+        const processOnMount = () => {
           const onUnmount = createInvocationContext(batch, () =>
             atomOnMount(atom, (...args) => setAtom(...args)),
           )
           if (onUnmount) {
             mounted.u = (batch) => createInvocationContext(batch, onUnmount)
           }
-        })
+        }
+        addBatchFunc(batch, processOnMount, 'L')
       }
     }
     return atomState.m
@@ -691,7 +709,7 @@ const buildStore = (
       // unmount self
       const onUnmount = atomState.m.u
       if (onUnmount) {
-        addBatchFunc(batch, 'L', () => onUnmount(batch))
+        addBatchFunc(batch, () => onUnmount(batch), 'L')
       }
       delete atomState.m
       if (import.meta.env?.MODE !== 'production') {
@@ -722,18 +740,15 @@ const buildStore = (
     }
   }
 
-  const unstable_derive = (fn: (...args: StoreArgs) => StoreArgs) =>
-    buildStore(
-      ...fn(baseGetAtomState, atomRead, atomWrite, atomOnMount, atomOnInit),
-    )
+  const unstable_derive: Store['unstable_derive'] = (fn) =>
+    buildStore(...fn(...storeArgs))
 
   const store: Store = {
     get: readAtom,
     set: writeAtom,
     sub: subscribeAtom,
     unstable_derive,
   }
-  const getAtomState = baseGetAtomState(atomOnInit(store))
   if (import.meta.env?.MODE !== 'production') {
     const devStore: DevStoreRev4 = {
       // store dev methods (these are tentative and subject to change without notice)
@@ -771,21 +786,21 @@ const buildStore = (
 }
 
 export const createStore = (): Store => {
-  const atomStateMap = new WeakMap<AnyAtom, AtomState>()
+  const atomStateMap = new WeakMap()
+  const getAtomState = <Value>(atom: Atom<Value>, atomOnInit?: AtomOnInit) => {
+    if (import.meta.env?.MODE !== 'production' && !atom) {
+      throw new Error('Atom is undefined or null')
+    }
+    let atomState = atomStateMap.get(atom) as AtomState<Value> | undefined
+    if (!atomState) {
+      atomState = { d: new Map(), p: new Set(), n: 0 }
+      atomStateMap.set(atom, atomState)
+      atomOnInit?.(atom, atomState)
+    }
+    return atomState
+  }
   return buildStore(
-    (atomOnInit) =>
-      function getAtomState<Value>(atom: Atom<Value>) {
-        if (import.meta.env?.MODE !== 'production' && !atom) {
-          throw new Error('Atom is undefined or null')
-        }
-        let atomState = atomStateMap.get(atom) as AtomState<Value> | undefined
-        if (!atomState) {
-          atomState = { d: new Map(), p: new Set(), n: 0 }
-          atomStateMap.set(atom, atomState)
-          atomOnInit(atom, atomState)
-        }
-        return atomState
-      },
+    getAtomState,
     (atom, ...params) => atom.read(...params),
     (atom, ...params) => atom.write(...params),
     (atom, ...params) => atom.onMount?.(...params),

tests/vanilla/effect.test.ts

@@ -102,7 +102,7 @@ function atomSyncEffect(effect: Effect) {
               }
             : null
       }
-      const entry = ['H', listener] as const
+      const entry = [listener, 'H'] as const
       atomState.l!.add(entry)
       return () => atomState.l!.delete(entry)
     }

tests/vanilla/store.test.tsx

@@ -1013,7 +1013,12 @@ it('should call onInit only once per atom', () => {
   a.INTERNAL_onInit = onInit
   store.get(a)
   expect(onInit).toHaveBeenCalledTimes(1)
-  expect(onInit).toHaveBeenCalledWith(store)
+  const aAtomState = expect.objectContaining({
+    d: expect.any(Map),
+    p: expect.any(Set),
+    n: expect.any(Number),
+  })
+  expect(onInit).toHaveBeenCalledWith(store, aAtomState)
   onInit.mockClear()
   store.get(a)
   store.set(a, 1)
@@ -1027,19 +1032,32 @@ it('should call onInit only once per atom', () => {
 
 it('should call onInit only once per store', () => {
   const a = atom(0)
-  const aOnInit = vi.fn()
+  type AtomState = Parameters<NonNullable<Atom<unknown>['INTERNAL_onInit']>>[1]
+  let aAtomState: AtomState
+  const aOnInit = vi.fn((_store, atomState) => {
+    aAtomState = atomState
+  })
   a.INTERNAL_onInit = aOnInit
   const b = atom(0)
-  const bOnInit = vi.fn()
+  let bAtomState: AtomState
+  const bOnInit = vi.fn((_store, atomState) => {
+    bAtomState = atomState
+  })
   b.INTERNAL_onInit = bOnInit
   type Store = ReturnType<typeof createStore>
   function testInStore(store: Store) {
     store.get(a)
     store.get(b)
+    const mockAtomState = expect.objectContaining({
+      d: expect.any(Map),
+      p: expect.any(Set),
+      n: expect.any(Number),
+    })
     expect(aOnInit).toHaveBeenCalledTimes(1)
     expect(bOnInit).toHaveBeenCalledTimes(1)
-    expect(aOnInit).toHaveBeenCalledWith(store)
-    expect(bOnInit).toHaveBeenCalledWith(store)
+    expect(aOnInit).toHaveBeenCalledWith(store, mockAtomState)
+    expect(bOnInit).toHaveBeenCalledWith(store, mockAtomState)
+    expect(aAtomState).not.toBe(bAtomState)
     aOnInit.mockClear()
     bOnInit.mockClear()
     return store
@@ -1052,18 +1070,19 @@ it('should call onInit only once per store', () => {
         const initializedAtoms = new WeakSet()
         return [
           (a, atomOnInit) => {
+            const atomState = getAtomState(a)
             if (!initializedAtoms.has(a)) {
               initializedAtoms.add(a)
-              atomOnInit?.(a)
+              atomOnInit?.(a, atomState)
             }
-            return getAtomState(a, atomOnInit)
+            return atomState
           },
           readAtom,
           writeAtom,
           atomOnMount,
           atomOnInit,
         ]
       },
-    ),
+    ) as Store,
   )
 })

tests/vanilla/unstable_derive.test.tsx

@@ -26,7 +26,7 @@ describe('unstable_derive for scoping atoms', () => {
               }
               return atomState
             }
-            return getAtomState(atom, atomOnInit(derivedStore))
+            return getAtomState(atom)
           },
           atomRead,
           atomWrite,
@@ -72,7 +72,7 @@ describe('unstable_derive for scoping atoms', () => {
               }
               return atomState
             }
-            return getAtomState(atom, atomOnInit(derivedStore))
+            return getAtomState(atom)
           },
           atomRead,
           atomWrite,
@@ -117,7 +117,7 @@ describe('unstable_derive for scoping atoms', () => {
                 }
                 return atomState
               }
-              return getAtomState(atom, atomOnInit(derivedStore))
+              return getAtomState(atom)
             },
             (a, get, options) => {
               const myGet: Getter = (aa) => {
@@ -183,8 +183,21 @@ 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((...args) => args)
+  const derivedStore = baseStore.unstable_derive(
+    (getAtomState, readAtom, writeAtom, atomOnMount, atomOnInit) => [
+      (a, atomOnInit) => {
+        const atomState = getAtomState(a)
+        atomOnInit?.(a, atomState)
+        return atomState
+      },
+      readAtom,
+      writeAtom,
+      atomOnMount,
+      atomOnInit,
+    ],
+  )
   const a = atom(null)
   a.INTERNAL_onInit = (store) => {
     expect(store).toBe(baseStore)