rework atomOnIn DI

dmaskasky · Dec 21, 2024 · 77d577e · 77d577e
1 parent a1e9433
commit 77d577e
Showing 1 changed file with 8 additions and 11 deletions.
diff --git a/src/vanilla/store.ts b/src/vanilla/store.ts
@@ -165,6 +165,8 @@ const addDependency = <Value>(
 // Batch
 //
 
+type BatchPriority = 'H' | 'M' | 'L'
+
 type Batch = Readonly<{
   /** Atom dependents map */
   D: Map<AnyAtom, Set<AnyAtom>>
@@ -176,8 +178,6 @@ type Batch = Readonly<{
   L: Set<() => void>
 }>
 
-type BatchPriority = 'H' | 'M' | 'L'
-
 const createBatch = (): Batch => ({
   D: new Map(),
   H: new Set(),
@@ -197,8 +197,8 @@ const registerBatchAtom = (
   if (!batch.D.has(atom)) {
     batch.D.set(atom, new Set())
     addBatchFunc(batch, 'H', () => {
-      for (const [p, f] of atomState.l || []) {
-        addBatchFunc(batch, p, f)
+      for (const [priority, listener] of atomState.l || []) {
+        addBatchFunc(batch, priority, listener)
       }
       for (const listener of atomState.m?.l || []) {
         addBatchFunc(batch, 'M', listener)
@@ -250,10 +250,9 @@ const flushBatch = (batch: Batch) => {
 
 // internal & unstable type
 type StoreArgs = readonly [
-  getAtomState: <Value>(
-    atom: Atom<Value>,
+  getAtomState: (
     atomOnInit: ReturnType<StoreArgs[4]>,
-  ) => AtomState<Value>,
+  ) => <Value>(atom: Atom<Value>) => AtomState<Value>,
   atomRead: <Value>(
     atom: Atom<Value>,
     ...params: Parameters<Atom<Value>['read']>
@@ -298,9 +297,6 @@ export type INTERNAL_PrdStore = PrdStore
 const buildStore = (
   ...[baseGetAtomState, atomRead, atomWrite, atomOnMount, atomOnInit]: StoreArgs
 ): Store => {
-  const getAtomState = <Value>(atom: Atom<Value>) =>
-    baseGetAtomState(atom, atomOnInit(store))
-
   // for debugging purpose only
   let debugMountedAtoms: Set<AnyAtom>
 
@@ -737,6 +733,7 @@ const buildStore = (
     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)
@@ -775,7 +772,7 @@ const buildStore = (
 
 export const createStore = (): Store => {
   const atomStateMap = new WeakMap()
-  const getAtomState: StoreArgs[0] = (atom, atomOnInit) => {
+  const getAtomState: StoreArgs[0] = (atomOnInit) => (atom) => {
     if (import.meta.env?.MODE !== 'production' && !atom) {
       throw new Error('Atom is undefined or null')
     }