diff --git a/src/vanilla/store.ts b/src/vanilla/store.ts index 53f9e40156..00939d24c8 100644 --- a/src/vanilla/store.ts +++ b/src/vanilla/store.ts @@ -247,9 +247,9 @@ const buildStore = (...storeArgs: StoreArgs): Store => { const changedAtoms = new Map() const unmountCallbacks = new Set<() => void>() const mountCallbacks = new Set<() => void>() - let inTransaction = 0 - const runWithTransaction = (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 = (a: Atom) => { @@ -440,7 +433,7 @@ const buildStore = (...storeArgs: StoreArgs): Store => { return dependents } - const invalidateDependents = (atomState: AtomState) => { + const invalidateDependents = (atomState: AtomState) => { const visited = new WeakSet() const stack: AtomState[] = [atomState] while (stack.length) { @@ -448,8 +441,10 @@ const buildStore = (...storeArgs: StoreArgs): Store => { 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, ...args: Args ): Result => { + let isSync = true const getter: Getter = (a: Atom) => returnAtomValue(readAtomState(a)) const setter: Setter = ( a: WritableAtom, ...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 = ( atom: WritableAtom, ...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 = (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) => diff --git a/tests/syncEffect/syncEffect.test.tsx b/tests/syncEffect/syncEffect.test.tsx index e994041ed7..6f59a2ca10 100644 --- a/tests/syncEffect/syncEffect.test.tsx +++ b/tests/syncEffect/syncEffect.test.tsx @@ -358,7 +358,6 @@ it('should disallow synchronous set.recurse in cleanup', () => { }) const store = createDebugStore() store.sub(effectAtom, () => {}) - store.set(anotherAtom, increment) expect(() => store.set(anotherAtom, increment)).toThrowError( 'set.recurse is not allowed in cleanup', ) diff --git a/tests/syncEffect/syncEffect.ts b/tests/syncEffect/syncEffect.ts index 3e15de6f8f..4656bf9abf 100644 --- a/tests/syncEffect/syncEffect.ts +++ b/tests/syncEffect/syncEffect.ts @@ -53,6 +53,7 @@ export function syncEffect(effect: Effect): Atom & { effect: Effect } { }, () => {}, ) + internalAtom.onMount = () => () => {} internalAtom.unstable_onInit = (store) => { const ref = store.get(refAtom) @@ -116,19 +117,16 @@ export function syncEffect(effect: Effect): Atom & { effect: Effect } { try { fromCleanup = true return cleanup() - } /* catch (error) { - ref.error = error - refresh() - } */ finally { + } finally { fromCleanup = false runCleanup = undefined } } } - } /* catch (error) { + } catch (error) { ref.error = error - refresh() - } */ finally { + throw error + } finally { Array.from(deps.keys(), ref.get!) --inProgress } @@ -169,7 +167,6 @@ export function syncEffect(effect: Effect): Atom & { effect: Effect } { syncEffectChannel.add(runEffect) }) } - internalAtom.onMount = () => () => {} if (process.env.NODE_ENV !== 'production') { function setLabel(atom: Atom, label: string) { @@ -209,7 +206,7 @@ function ensureSyncEffectChannel(store: Store) { () => void >() hookInto(storeWithHooks, INTERNAL_flushStoreHook, () => { - syncEffectChannel!.forEach((fn: () => void) => fn()) + syncEffectChannel!.forEach((fn) => fn()) syncEffectChannel!.clear() }) }