refactor(store): avoid transaction

src/vanilla/store.ts

@@ -247,7 +247,6 @@ 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 flushCallbacks = () => {
     const errors: unknown[] = []
@@ -258,13 +257,7 @@ const buildStore = (...storeArgs: StoreArgs): Store => {
         errors.push(e)
       }
     }
-    ++inTransaction
-    while (changedAtoms.size || unmountCallbacks.size || mountCallbacks.size) {
-      recomputeInvalidatedAtoms()
-      if (inTransaction > 1) {
-        --inTransaction
-        return
-      }
+    do {
       ;(store as any)[INTERNAL_flushStoreHook]?.()
       const callbacks = new Set<() => void>()
       const add = callbacks.add.bind(callbacks)
@@ -275,8 +268,10 @@ const buildStore = (...storeArgs: StoreArgs): Store => {
       mountCallbacks.forEach(add)
       mountCallbacks.clear()
       callbacks.forEach(call)
-    }
-    --inTransaction
+      if (changedAtoms.size) {
+        recomputeInvalidatedAtoms()
+      }
+    } while (changedAtoms.size || unmountCallbacks.size || mountCallbacks.size)
     if (errors.length) {
       throw errors[0]
     }
@@ -434,17 +429,13 @@ const buildStore = (...storeArgs: StoreArgs): Store => {
   }
 
   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)) {
-          if (!invalidatedAtoms.has(d)) {
-            invalidatedAtoms.set(d, s.n)
-            stack.push(s)
-          }
+      for (const [d, s] of getMountedOrPendingDependents(aState)) {
+        if (!invalidatedAtoms.has(d)) {
+          invalidatedAtoms.set(d, s.n)
+          stack.push(s)
         }
       }
     }
@@ -552,6 +543,7 @@ const buildStore = (...storeArgs: StoreArgs): Store => {
         }
       } finally {
         if (!isSync) {
+          recomputeInvalidatedAtoms()
           flushCallbacks()
         }
       }
@@ -570,6 +562,7 @@ const buildStore = (...storeArgs: StoreArgs): Store => {
     try {
       return writeAtomState(atom, ...args)
     } finally {
+      recomputeInvalidatedAtoms()
       flushCallbacks()
     }
   }
@@ -614,32 +607,34 @@ const buildStore = (...storeArgs: StoreArgs): Store => {
       atomState.h?.()
       if (isActuallyWritableAtom(atom)) {
         const mounted = atomState.m
-        let setAtom: (...args: unknown[]) => unknown
-        const createInvocationContext = <T>(fn: () => T) => {
+        const processOnMount = () => {
           let isSync = true
-          setAtom = (...args: unknown[]) => {
+          const setAtom = (...args: unknown[]) => {
             try {
               return writeAtomState(atom, ...args)
             } finally {
               if (!isSync) {
+                recomputeInvalidatedAtoms()
                 flushCallbacks()
               }
             }
           }
           try {
-            return fn()
+            const onUnmount = atomOnMount(atom, setAtom)
+            if (onUnmount) {
+              mounted.u = () => {
+                isSync = true
+                try {
+                  onUnmount()
+                } finally {
+                  isSync = false
+                }
+              }
+            }
           } finally {
             isSync = false
           }
         }
-        const processOnMount = () => {
-          const onUnmount = createInvocationContext(() =>
-            atomOnMount(atom, (...args) => setAtom(...args)),
-          )
-          if (onUnmount) {
-            mounted.u = () => createInvocationContext(onUnmount)
-          }
-        }
         mountCallbacks.add(processOnMount)
       }
     }

tests/vanilla/store.test.tsx

@@ -645,6 +645,34 @@ describe('should invoke flushPending only after all atoms are updated (#2804)',
       'after store.sub',
     ])
   })
+
+  it('should flush only after all atoms are updated with unmount', () => {
+    const result: string[] = []
+    const a = atom(0)
+    const b = atom(null, (_get, set, value: number) => {
+      set(a, value)
+    })
+    b.onMount = (setAtom) => {
+      return () => {
+        result.push('onUmount: before setAtom')
+        setAtom(1)
+        result.push('onUmount: after setAtom')
+      }
+    }
+    const c = atom(true)
+    const d = atom((get) => get(c) && get(b))
+    store.sub(a, () => {
+      result.push('a value changed - ' + store.get(a))
+    })
+    store.sub(d, () => {})
+    expect(store.get(d)).toEqual(null)
+    store.set(c, false)
+    expect(result).toEqual([
+      'onUmount: before setAtom',
+      'onUmount: after setAtom',
+      'a value changed - 1',
+    ])
+  })
 })
 
 describe('should mount and trigger listeners even when an error is thrown', () => {