From 51a742df2c5893cc22096b6b145c26c62c94f7cd Mon Sep 17 00:00:00 2001
From: David Maskasky <dmaskasky@gmail.com>
Date: Thu, 19 Dec 2024 03:22:07 -0800
Subject: [PATCH] refactor: batch priority (#2875)

---
 src/vanilla/store.ts | 28 +++++++++++++---------------
 1 file changed, 13 insertions(+), 15 deletions(-)

diff --git a/src/vanilla/store.ts b/src/vanilla/store.ts
index 341f81a351..0a7d7e368a 100644
--- a/src/vanilla/store.ts
+++ b/src/vanilla/store.ts
@@ -163,6 +163,8 @@ const addDependency = <Value>(
 // Batch
 //
 
+type BatchPriority = 'H' | 'M' | 'L'
+
 type Batch = Readonly<{
   /** Atom dependents map */
   D: Map<AnyAtom, Set<AnyAtom>>
@@ -181,16 +183,12 @@ const createBatch = (): Batch => ({
   L: new Set(),
 })
 
-const addBatchFuncHigh = (batch: Batch, fn: () => void) => {
-  batch.H.add(fn)
-}
-
-const addBatchFuncMedium = (batch: Batch, fn: () => void) => {
-  batch.M.add(fn)
-}
-
-const addBatchFuncLow = (batch: Batch, fn: () => void) => {
-  batch.L.add(fn)
+const addBatchFunc = (
+  batch: Batch,
+  priority: BatchPriority,
+  fn: () => void,
+) => {
+  batch[priority].add(fn)
 }
 
 const registerBatchAtom = (
@@ -200,8 +198,8 @@ const registerBatchAtom = (
 ) => {
   if (!batch.D.has(atom)) {
     batch.D.set(atom, new Set())
-    addBatchFuncMedium(batch, () => {
-      atomState.m?.l.forEach((listener) => addBatchFuncMedium(batch, listener))
+    addBatchFunc(batch, 'M', () => {
+      atomState.m?.l.forEach((listener) => addBatchFunc(batch, 'M', listener))
     })
   }
 }
@@ -514,7 +512,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
-    addBatchFuncHigh(batch, () => {
+    addBatchFunc(batch, 'H', () => {
       const changedAtoms = new Set<AnyAtom>([atom])
       for (let i = topSortedReversed.length - 1; i >= 0; --i) {
         const [a, aState, prevEpochNumber] = topSortedReversed[i]!
@@ -659,7 +657,7 @@ const buildStore = (
             isSync = false
           }
         }
-        addBatchFuncLow(batch, () => {
+        addBatchFunc(batch, 'L', () => {
           const onUnmount = createInvocationContext(batch, () =>
             atomOnMount(atom, (...args) => setAtom(...args)),
           )
@@ -685,7 +683,7 @@ const buildStore = (
       // unmount self
       const onUnmount = atomState.m.u
       if (onUnmount) {
-        addBatchFuncLow(batch, () => onUnmount(batch))
+        addBatchFunc(batch, 'L', () => onUnmount(batch))
       }
       delete atomState.m
       if (import.meta.env?.MODE !== 'production') {