refactor(store): Replace runWithTransactions to flushCallbacks (#2946)

* simplify store * prevent recursive flushCallbacks * short-circuit invalidateDependents an already invalidated --------- Co-authored-by: Daishi Kato <[email protected]>
pmndrs · Jan 19, 2025 · 81131dc · 81131dc
1 parent e98b58a
commit 81131dc
Showing 1 changed file with 74 additions and 48 deletions.
diff --git a/src/vanilla/store.ts b/src/vanilla/store.ts
@@ -247,9 +247,9 @@ const buildStore = (...storeArgs: StoreArgs): Store => {
   const changedAtoms = new Map<AnyAtom, AtomState>()
   const unmountCallbacks = new Set<() => void>()
   const mountCallbacks = new Set<() => void>()
-
   let inTransaction = 0
-  const runWithTransaction = <T>(fn: () => T): T => {
+
+  const flushCallbacks = () => {
     const errors: unknown[] = []
     const call = (fn: () => void) => {
       try {
@@ -258,36 +258,28 @@ const buildStore = (...storeArgs: StoreArgs): Store => {
         errors.push(e)
       }
     }
-    let result: T
     ++inTransaction
-    try {
-      result = fn()
-    } finally {
-      if (inTransaction === 1) {
-        while (
-          changedAtoms.size ||
-          unmountCallbacks.size ||
-          mountCallbacks.size
-        ) {
-          recomputeInvalidatedAtoms()
-          ;(store as any)[INTERNAL_flushStoreHook]?.()
-          const callbacks = new Set<() => void>()
-          const add = callbacks.add.bind(callbacks)
-          changedAtoms.forEach((atomState) => atomState.m?.l.forEach(add))
-          changedAtoms.clear()
-          unmountCallbacks.forEach(add)
-          unmountCallbacks.clear()
-          mountCallbacks.forEach(add)
-          mountCallbacks.clear()
-          callbacks.forEach(call)
-        }
+    while (changedAtoms.size || unmountCallbacks.size || mountCallbacks.size) {
+      recomputeInvalidatedAtoms()
+      if (inTransaction > 1) {
+        --inTransaction
+        return
       }
-      --inTransaction
+      ;(store as any)[INTERNAL_flushStoreHook]?.()
+      const callbacks = new Set<() => void>()
+      const add = callbacks.add.bind(callbacks)
+      changedAtoms.forEach((atomState) => atomState.m?.l.forEach(add))
+      changedAtoms.clear()
+      unmountCallbacks.forEach(add)
+      unmountCallbacks.clear()
+      mountCallbacks.forEach(add)
+      mountCallbacks.clear()
+      callbacks.forEach(call)
     }
+    --inTransaction
     if (errors.length) {
       throw errors[0]
     }
-    return result
   }
 
   const setAtomStateValueOrPromise = (
@@ -344,7 +336,8 @@ const buildStore = (...storeArgs: StoreArgs): Store => {
     let isSync = true
     const mountDependenciesIfAsync = () => {
       if (atomState.m) {
-        runWithTransaction(() => mountDependencies(atom, atomState))
+        mountDependencies(atom, atomState)
+        flushCallbacks()
       }
     }
     const getter: Getter = <V>(a: Atom<V>) => {
@@ -440,16 +433,18 @@ const buildStore = (...storeArgs: StoreArgs): Store => {
     return dependents
   }
 
-  const invalidateDependents = <Value>(atomState: AtomState<Value>) => {
+  const invalidateDependents = (atomState: AtomState) => {
     const visited = new WeakSet<AtomState>()
     const stack: AtomState[] = [atomState]
     while (stack.length) {
       const aState = stack.pop()!
       if (!visited.has(aState)) {
         visited.add(aState)
         for (const [d, s] of getMountedOrPendingDependents(aState)) {
-          invalidatedAtoms.set(d, s.n)
-          stack.push(s)
+          if (!invalidatedAtoms.has(d)) {
+            invalidatedAtoms.set(d, s.n)
+            stack.push(s)
+          }
         }
       }
     }
@@ -529,13 +524,14 @@ const buildStore = (...storeArgs: StoreArgs): Store => {
     atom: WritableAtom<Value, Args, Result>,
     ...args: Args
   ): Result => {
+    let isSync = true
     const getter: Getter = <V>(a: Atom<V>) => returnAtomValue(readAtomState(a))
     const setter: Setter = <V, As extends unknown[], R>(
       a: WritableAtom<V, As, R>,
       ...args: As
     ) => {
       const aState = ensureAtomState(a)
-      return runWithTransaction(() => {
+      try {
         if (isSelfAtom(atom, a)) {
           if (!hasInitialValue(a)) {
             // NOTE technically possible but restricted as it may cause bugs
@@ -554,15 +550,29 @@ const buildStore = (...storeArgs: StoreArgs): Store => {
         } else {
           return writeAtomState(a, ...args)
         }
-      })
+      } finally {
+        if (!isSync) {
+          flushCallbacks()
+        }
+      }
+    }
+    try {
+      return atomWrite(atom, getter, setter, ...args)
+    } finally {
+      isSync = false
     }
-    return atomWrite(atom, getter, setter, ...args)
   }
 
   const writeAtom = <Value, Args extends unknown[], Result>(
     atom: WritableAtom<Value, Args, Result>,
     ...args: Args
-  ): Result => runWithTransaction(() => writeAtomState(atom, ...args))
+  ): Result => {
+    try {
+      return writeAtomState(atom, ...args)
+    } finally {
+      flushCallbacks()
+    }
+  }
 
   const mountDependencies = (atom: AnyAtom, atomState: AtomState) => {
     if (atomState.m && !isPendingPromise(atomState.v)) {
@@ -604,12 +614,30 @@ const buildStore = (...storeArgs: StoreArgs): Store => {
       atomState.h?.()
       if (isActuallyWritableAtom(atom)) {
         const mounted = atomState.m
+        let setAtom: (...args: unknown[]) => unknown
+        const createInvocationContext = <T>(fn: () => T) => {
+          let isSync = true
+          setAtom = (...args: unknown[]) => {
+            try {
+              return writeAtomState(atom, ...args)
+            } finally {
+              if (!isSync) {
+                flushCallbacks()
+              }
+            }
+          }
+          try {
+            return fn()
+          } finally {
+            isSync = false
+          }
+        }
         const processOnMount = () => {
-          const onUnmount = atomOnMount(atom, (...args) =>
-            runWithTransaction(() => writeAtomState(atom, ...args)),
+          const onUnmount = createInvocationContext(() =>
+            atomOnMount(atom, (...args) => setAtom(...args)),
           )
           if (onUnmount) {
-            mounted.u = onUnmount
+            mounted.u = () => createInvocationContext(onUnmount)
           }
         }
         mountCallbacks.add(processOnMount)
@@ -646,17 +674,15 @@ const buildStore = (...storeArgs: StoreArgs): Store => {
 
   const subscribeAtom = (atom: AnyAtom, listener: () => void) => {
     const atomState = ensureAtomState(atom)
-    return runWithTransaction(() => {
-      const mounted = mountAtom(atom, atomState)
-      const listeners = mounted.l
-      listeners.add(listener)
-      return () => {
-        runWithTransaction(() => {
-          listeners.delete(listener)
-          unmountAtom(atom, atomState)
-        })
-      }
-    })
+    const mounted = mountAtom(atom, atomState)
+    const listeners = mounted.l
+    listeners.add(listener)
+    flushCallbacks()
+    return () => {
+      listeners.delete(listener)
+      unmountAtom(atom, atomState)
+      flushCallbacks()
+    }
   }
 
   const unstable_derive: Store['unstable_derive'] = (fn) =>