From 590252f70ad1e6620bb12a8d08b8e6529f148ae8 Mon Sep 17 00:00:00 2001 From: daishi Date: Fri, 10 Jan 2025 19:37:25 +0900 Subject: [PATCH 01/23] refactor: eliminate batch --- src/vanilla/store.ts | 289 +++++++++++++++-------------------- tests/vanilla/effect.test.ts | 98 ++++++------ 2 files changed, 166 insertions(+), 221 deletions(-) diff --git a/src/vanilla/store.ts b/src/vanilla/store.ts index 6c20762127..929176a961 100644 --- a/src/vanilla/store.ts +++ b/src/vanilla/store.ts @@ -77,7 +77,7 @@ type Mounted = { /** Set of mounted atoms that depends on the atom. */ readonly t: Set /** Function to run when the atom is unmounted. */ - u?: (batch: Batch) => void + u?: () => void } /** @@ -102,19 +102,19 @@ type AtomState = { m?: Mounted // only available if the atom is mounted /** * Listener to notify when the atom value is updated. - * This is still an experimental API and subject to change without notice. + * This is an experimental API and will be changed in the next minor. */ - u?: (batch: Batch) => void + u?: () => void /** * Listener to notify when the atom is mounted or unmounted. - * This is still an experimental API and subject to change without notice. + * This is an experimental API and will be changed in the next minor. */ - h?: (batch: Batch) => void + h?: () => void /** Atom value */ v?: Value /** Atom error */ e?: AnyError - /** Indicates that the atom value has been changed */ + /** Indicates that the atom value has to be recomputed */ x?: true } @@ -165,75 +165,6 @@ const addDependency = ( aState.m?.t.add(atom) } -// -// Batch -// - -type BatchPriority = 0 | 1 | 2 - -type Batch = [ - /** finish recompute */ - priority0: Set<() => void>, - /** atom listeners */ - priority1: Set<() => void>, - /** atom mount hooks */ - priority2: Set<() => void>, -] & { - /** changed Atoms */ - C: Set -} - -const createBatch = (): Batch => - Object.assign([new Set(), new Set(), new Set()], { C: new Set() }) as Batch - -const addBatchFunc = ( - batch: Batch, - priority: BatchPriority, - fn: () => void, -) => { - batch[priority].add(fn) -} - -const registerBatchAtom = ( - batch: Batch, - atom: AnyAtom, - atomState: AtomState, -) => { - if (!batch.C.has(atom)) { - batch.C.add(atom) - atomState.u?.(batch) - const scheduleListeners = () => { - atomState.m?.l.forEach((listener) => addBatchFunc(batch, 1, listener)) - } - addBatchFunc(batch, 1, scheduleListeners) - } -} - -const flushBatch = (batch: Batch) => { - let error: AnyError - let hasError = false - const call = (fn: () => void) => { - try { - fn() - } catch (e) { - if (!hasError) { - error = e - hasError = true - } - } - } - while (batch.C.size || batch.some((channel) => channel.size)) { - batch.C.clear() - for (const channel of batch) { - channel.forEach(call) - channel.clear() - } - } - if (hasError) { - throw error - } -} - // internal & unstable type type StoreArgs = readonly [ getAtomState: (atom: Atom) => AtomState | undefined, @@ -275,6 +206,11 @@ type Store = { export type INTERNAL_DevStoreRev4 = DevStoreRev4 export type INTERNAL_PrdStore = Store +/** + * This is an experimental API and will be changed in the next minor. + */ +export const INTERNAL_flushStoreHook: unique symbol = Symbol() + const buildStore = (...storeArgs: StoreArgs): Store => { const [ getAtomState, @@ -297,6 +233,34 @@ const buildStore = (...storeArgs: StoreArgs): Store => { return atomState } + const changedAtoms = new Map() + const unmountCallbacks = new Set<() => void>() + const mountCallbacks = new Set<() => void>() + + const flushCallbacks = () => { + const errors: unknown[] = [] + const call = (fn: () => void) => { + try { + fn() + } catch (e) { + errors.push(e) + } + } + do { + ;(store as any)[INTERNAL_flushStoreHook]?.() + changedAtoms.forEach((atomState) => atomState.m?.l.forEach(call)) + changedAtoms.clear() + unmountCallbacks.forEach(call) + unmountCallbacks.clear() + mountCallbacks.forEach(call) + mountCallbacks.clear() + recomputeInvalidatedAtoms() + } while (changedAtoms.size) + if (errors.length) { + throw errors[0] + } + } + const setAtomStateValueOrPromise = ( atom: AnyAtom, atomState: AtomState, @@ -324,16 +288,13 @@ const buildStore = (...storeArgs: StoreArgs): Store => { } } - const readAtomState = ( - batch: Batch | undefined, - atom: Atom, - ): AtomState => { + const readAtomState = (atom: Atom): AtomState => { const atomState = ensureAtomState(atom) // See if we can skip recomputing this atom. if (isAtomStateInitialized(atomState)) { // If the atom is mounted, we can use cached atom state. // because it should have been updated by dependencies. - // We can't use the cache if the atom is dirty. + // We can't use the cache if the atom is invalidated. if (atomState.m && !atomState.x) { return atomState } @@ -344,7 +305,7 @@ const buildStore = (...storeArgs: StoreArgs): Store => { ([a, n]) => // Recursively, read the atom state of the dependency, and // check if the atom epoch number is unchanged - readAtomState(batch, a).n === n, + readAtomState(a).n === n, ) ) { return atomState @@ -367,17 +328,14 @@ const buildStore = (...storeArgs: StoreArgs): Store => { return returnAtomValue(aState) } // a !== atom - const aState = readAtomState(batch, a) + const aState = readAtomState(a) try { return returnAtomValue(aState) } finally { - if (isSync) { - addDependency(atom, atomState, a, aState) - } else { - const batch = createBatch() - addDependency(atom, atomState, a, aState) - mountDependencies(batch, atom, atomState) - flushBatch(batch) + addDependency(atom, atomState, a, aState) + if (!isSync) { + mountDependencies(atom, atomState) + flushCallbacks() } } } @@ -417,9 +375,8 @@ const buildStore = (...storeArgs: StoreArgs): Store => { valueOrPromise.onCancel?.(() => controller?.abort()) const complete = () => { if (atomState.m) { - const batch = createBatch() - mountDependencies(batch, atom, atomState) - flushBatch(batch) + mountDependencies(atom, atomState) + flushCallbacks() } } valueOrPromise.then(complete, complete) @@ -437,9 +394,9 @@ const buildStore = (...storeArgs: StoreArgs): Store => { } const readAtom = (atom: Atom): Value => - returnAtomValue(readAtomState(undefined, atom)) + returnAtomValue(readAtomState(atom)) - const getMountedOrBatchDependents = ( + const getMountedOrPendingDependents = ( atomState: AtomState, ): Map => { const dependents = new Map() @@ -458,11 +415,22 @@ const buildStore = (...storeArgs: StoreArgs): Store => { return dependents } - const recomputeDependents = ( - batch: Batch, - atom: Atom, - atomState: AtomState, - ) => { + const invalidateDependents = (atomState: AtomState) => { + const visited = new WeakSet() + const stack: AtomState[] = [atomState] + while (stack.length) { + const aState = stack.pop()! + if (!visited.has(aState)) { + visited.add(aState) + for (const s of getMountedOrPendingDependents(aState).values()) { + s.x = true + stack.push(s) + } + } + } + } + + const recomputeInvalidatedAtoms = () => { // Step 1: traverse the dependency graph to build the topsorted atom list // We don't bother to check for cycles, which simplifies the algorithm. // This is a topological sort via depth-first search, slightly modified from @@ -473,12 +441,12 @@ const buildStore = (...storeArgs: StoreArgs): Store => { atomState: AtomState, epochNumber: number, ][] = [] - const visiting = new Set() - const visited = new Set() + const visiting = new WeakSet() + const visited = new WeakSet() // Visit the root atom. This is the only atom in the dependency graph // without incoming edges, which is one reason we can simplify the algorithm - const stack: [a: AnyAtom, aState: AtomState][] = [[atom, atomState]] - while (stack.length > 0) { + const stack: [a: AnyAtom, aState: AtomState][] = Array.from(changedAtoms) + while (stack.length) { const [a, aState] = stack[stack.length - 1]! if (visited.has(a)) { // All dependents have been processed, now process this atom @@ -489,17 +457,17 @@ const buildStore = (...storeArgs: StoreArgs): Store => { // The algorithm calls for pushing onto the front of the list. For // performance, we will simply push onto the end, and then will iterate in // reverse order later. - topSortedReversed.push([a, aState, aState.n]) + if (aState.x) { + topSortedReversed.push([a, aState, aState.n]) + } // Atom has been visited but not yet processed visited.add(a) - // Mark atom dirty - aState.x = true stack.pop() continue } visiting.add(a) // Push unvisited dependents onto the stack - for (const [d, s] of getMountedOrBatchDependents(aState)) { + for (const [d, s] of getMountedOrPendingDependents(aState)) { if (a !== d && !visiting.has(d)) { stack.push([d, s]) } @@ -508,39 +476,33 @@ const buildStore = (...storeArgs: StoreArgs): Store => { // Step 2: use the topSortedReversed atom list to recompute all affected atoms // Track what's changed, so that we can short circuit when possible - const finishRecompute = () => { - const changedAtoms = new Set([atom]) - for (let i = topSortedReversed.length - 1; i >= 0; --i) { - const [a, aState, prevEpochNumber] = topSortedReversed[i]! - let hasChangedDeps = false - for (const dep of aState.d.keys()) { - if (dep !== a && changedAtoms.has(dep)) { - hasChangedDeps = true - break - } + for (let i = topSortedReversed.length - 1; i >= 0; --i) { + const [a, aState, prevEpochNumber] = topSortedReversed[i]! + let hasChangedDeps = false + for (const dep of aState.d.keys()) { + if (dep !== a && changedAtoms.has(dep)) { + hasChangedDeps = true + break } - if (hasChangedDeps) { - readAtomState(batch, a) - mountDependencies(batch, a, aState) - if (prevEpochNumber !== aState.n) { - registerBatchAtom(batch, a, aState) - changedAtoms.add(a) - } + } + if (hasChangedDeps) { + readAtomState(a) + mountDependencies(a, aState) + if (prevEpochNumber !== aState.n) { + changedAtoms.set(a, aState) + aState.u?.() } - delete aState.x } + delete aState.x } - addBatchFunc(batch, 0, finishRecompute) } const writeAtomState = ( - batch: Batch, atom: WritableAtom, ...args: Args ): Result => { let isSync = true - const getter: Getter = (a: Atom) => - returnAtomValue(readAtomState(batch, a)) + const getter: Getter = (a: Atom) => returnAtomValue(readAtomState(a)) const setter: Setter = ( a: WritableAtom, ...args: As @@ -555,18 +517,20 @@ const buildStore = (...storeArgs: StoreArgs): Store => { const prevEpochNumber = aState.n const v = args[0] as V setAtomStateValueOrPromise(a, aState, v) - mountDependencies(batch, a, aState) + mountDependencies(a, aState) if (prevEpochNumber !== aState.n) { - registerBatchAtom(batch, a, aState) - recomputeDependents(batch, a, aState) + changedAtoms.set(a, aState) + aState.u?.() + invalidateDependents(aState) } return undefined as R } else { - return writeAtomState(batch, a, ...args) + return writeAtomState(a, ...args) } } finally { if (!isSync) { - flushBatch(batch) + recomputeInvalidatedAtoms() + flushCallbacks() } } } @@ -581,23 +545,19 @@ const buildStore = (...storeArgs: StoreArgs): Store => { atom: WritableAtom, ...args: Args ): Result => { - const batch = createBatch() try { - return writeAtomState(batch, atom, ...args) + return writeAtomState(atom, ...args) } finally { - flushBatch(batch) + recomputeInvalidatedAtoms() + flushCallbacks() } } - const mountDependencies = ( - batch: Batch, - atom: AnyAtom, - atomState: AtomState, - ) => { + const mountDependencies = (atom: AnyAtom, atomState: AtomState) => { if (atomState.m && !isPendingPromise(atomState.v)) { for (const a of atomState.d.keys()) { if (!atomState.m.d.has(a)) { - const aMounted = mountAtom(batch, a, ensureAtomState(a)) + const aMounted = mountAtom(a, ensureAtomState(a)) aMounted.t.add(atom) atomState.m.d.add(a) } @@ -605,7 +565,7 @@ const buildStore = (...storeArgs: StoreArgs): Store => { for (const a of atomState.m.d || []) { if (!atomState.d.has(a)) { atomState.m.d.delete(a) - const aMounted = unmountAtom(batch, a, ensureAtomState(a)) + const aMounted = unmountAtom(a, ensureAtomState(a)) aMounted?.t.delete(atom) } } @@ -613,16 +573,15 @@ const buildStore = (...storeArgs: StoreArgs): Store => { } const mountAtom = ( - batch: Batch, atom: Atom, atomState: AtomState, ): Mounted => { if (!atomState.m) { // recompute atom state - readAtomState(batch, atom) + readAtomState(atom) // mount dependencies first for (const a of atomState.d.keys()) { - const aMounted = mountAtom(batch, a, ensureAtomState(a)) + const aMounted = mountAtom(a, ensureAtomState(a)) aMounted.t.add(atom) } // mount self @@ -631,18 +590,19 @@ const buildStore = (...storeArgs: StoreArgs): Store => { d: new Set(atomState.d.keys()), t: new Set(), } - atomState.h?.(batch) + atomState.h?.() if (isActuallyWritableAtom(atom)) { const mounted = atomState.m let setAtom: (...args: unknown[]) => unknown - const createInvocationContext = (batch: Batch, fn: () => T) => { + const createInvocationContext = (fn: () => T) => { let isSync = true setAtom = (...args: unknown[]) => { try { - return writeAtomState(batch, atom, ...args) + return writeAtomState(atom, ...args) } finally { if (!isSync) { - flushBatch(batch) + recomputeInvalidatedAtoms() + flushCallbacks() } } } @@ -653,21 +613,20 @@ const buildStore = (...storeArgs: StoreArgs): Store => { } } const processOnMount = () => { - const onUnmount = createInvocationContext(batch, () => + const onUnmount = createInvocationContext(() => atomOnMount(atom, (...args) => setAtom(...args)), ) if (onUnmount) { - mounted.u = (batch) => createInvocationContext(batch, onUnmount) + mounted.u = () => createInvocationContext(onUnmount) } } - addBatchFunc(batch, 2, processOnMount) + mountCallbacks.add(processOnMount) } } return atomState.m } const unmountAtom = ( - batch: Batch, atom: Atom, atomState: AtomState, ): Mounted | undefined => { @@ -679,13 +638,13 @@ const buildStore = (...storeArgs: StoreArgs): Store => { // unmount self const onUnmount = atomState.m.u if (onUnmount) { - addBatchFunc(batch, 2, () => onUnmount(batch)) + unmountCallbacks.add(onUnmount) } delete atomState.m - atomState.h?.(batch) + atomState.h?.() // unmount dependencies for (const a of atomState.d.keys()) { - const aMounted = unmountAtom(batch, a, ensureAtomState(a)) + const aMounted = unmountAtom(a, ensureAtomState(a)) aMounted?.t.delete(atom) } return undefined @@ -694,17 +653,15 @@ const buildStore = (...storeArgs: StoreArgs): Store => { } const subscribeAtom = (atom: AnyAtom, listener: () => void) => { - const batch = createBatch() const atomState = ensureAtomState(atom) - const mounted = mountAtom(batch, atom, atomState) + const mounted = mountAtom(atom, atomState) const listeners = mounted.l listeners.add(listener) - flushBatch(batch) + flushCallbacks() return () => { listeners.delete(listener) - const batch = createBatch() - unmountAtom(batch, atom, atomState) - flushBatch(batch) + unmountAtom(atom, atomState) + flushCallbacks() } } @@ -730,8 +687,8 @@ const deriveDevStoreRev4 = (store: Store): Store & DevStoreRev4 => { storeArgs[1] = function devSetAtomState(atom, atomState) { setAtomState(atom, atomState) const originalMounted = atomState.h - atomState.h = (batch) => { - originalMounted?.(batch) + atomState.h = () => { + originalMounted?.() if (atomState.m) { debugMountedAtoms.add(atom) } else { diff --git a/tests/vanilla/effect.test.ts b/tests/vanilla/effect.test.ts index 697d6a6733..9ce918c314 100644 --- a/tests/vanilla/effect.test.ts +++ b/tests/vanilla/effect.test.ts @@ -1,12 +1,13 @@ import { expect, it, vi } from 'vitest' import type { Atom, Getter, Setter } from 'jotai/vanilla' import { atom, createStore } from 'jotai/vanilla' +// For this test only +import { INTERNAL_flushStoreHook } from '../../src/vanilla/store.js' type Store = ReturnType type GetAtomState = Parameters[0]>[0] type AtomState = NonNullable> type AnyAtom = Atom -type Batch = Parameters>[0] type Cleanup = () => void type Effect = (get: Getter, set: Setter) => Cleanup | void @@ -17,8 +18,6 @@ type Ref = { cleanup?: Cleanup | undefined } -const syncEffectChannelSymbol = Symbol() - function syncEffect(effect: Effect): Atom { const refAtom = atom(() => ({ inProgress: 0, epoch: 0 })) const refreshAtom = atom(0) @@ -35,34 +34,37 @@ function syncEffect(effect: Effect): Atom { const ref = store.get(refAtom) const runEffect = () => { const deps = new Set() - ref.cleanup?.() - ref.cleanup = - effect( - (a) => { - deps.add(a) - return ref.get!(a) - }, - (a, ...args) => { - try { - ++ref.inProgress - return store.set(a, ...args) - } finally { - deps.forEach(ref.get!) - --ref.inProgress - } - }, - ) || undefined + try { + ref.cleanup?.() + ref.cleanup = + effect( + (a) => { + deps.add(a) + return store.get(a) + }, + (a, ...args) => { + try { + ++ref.inProgress + return store.set(a, ...args) + } finally { + --ref.inProgress + } + }, + ) || undefined + } finally { + deps.forEach(ref.get!) + } } const internalAtomState = getAtomState(store, internalAtom) const originalMountHook = internalAtomState.h - internalAtomState.h = (batch) => { - originalMountHook?.(batch) + internalAtomState.h = () => { + originalMountHook?.() if (internalAtomState.m) { // mount store.set(refreshAtom, (v) => v + 1) } else { // unmount - const syncEffectChannel = ensureBatchChannel(batch) + const syncEffectChannel = ensureSyncEffectChannel(store) syncEffectChannel.add(() => { ref.cleanup?.() delete ref.cleanup @@ -70,10 +72,10 @@ function syncEffect(effect: Effect): Atom { } } const originalUpdateHook = internalAtomState.u - internalAtomState.u = (batch) => { - originalUpdateHook?.(batch) + internalAtomState.u = () => { + originalUpdateHook?.() // update - const syncEffectChannel = ensureBatchChannel(batch) + const syncEffectChannel = ensureSyncEffectChannel(store) syncEffectChannel.add(runEffect) } } @@ -82,37 +84,23 @@ function syncEffect(effect: Effect): Atom { }) } -type BatchWithSyncEffect = Batch & { - [syncEffectChannelSymbol]?: Set<() => void> -} -function ensureBatchChannel(batch: BatchWithSyncEffect) { - // ensure continuation of the flushBatch while loop - const originalQueue = batch[1] - if (!originalQueue) { - throw new Error('batch[1] must be present') - } - if (!batch[syncEffectChannelSymbol]) { - batch[syncEffectChannelSymbol] = new Set<() => void>() - batch[1] = { - ...originalQueue, - add(item) { - originalQueue.add(item) - return this - }, - clear() { - batch[syncEffectChannelSymbol]!.clear() - originalQueue.clear() - }, - forEach(callback) { - batch[syncEffectChannelSymbol]!.forEach(callback) - originalQueue.forEach(callback) - }, - get size() { - return batch[syncEffectChannelSymbol]!.size + originalQueue.size - }, +const syncEffectChannelSymbol = Symbol() + +function ensureSyncEffectChannel(store: any) { + if (!store[syncEffectChannelSymbol]) { + store[syncEffectChannelSymbol] = new Set<() => void>() + const originalFlushHook = store[INTERNAL_flushStoreHook] + store[INTERNAL_flushStoreHook] = () => { + originalFlushHook?.() + const syncEffectChannel = store[syncEffectChannelSymbol] as Set< + () => void + > + const fns = Array.from(syncEffectChannel) + syncEffectChannel.clear() + fns.forEach((fn: () => void) => fn()) } } - return batch[syncEffectChannelSymbol]! + return store[syncEffectChannelSymbol] as Set<() => void> } const getAtomStateMap = new WeakMap() From 4f0baaa50ae75de406c6bc30e96ec0aa894d1da0 Mon Sep 17 00:00:00 2001 From: daishi Date: Fri, 10 Jan 2025 20:48:35 +0900 Subject: [PATCH 02/23] secret internal symbol --- src/vanilla/store.ts | 2 +- tests/vanilla/effect.test.ts | 3 +-- 2 files changed, 2 insertions(+), 3 deletions(-) diff --git a/src/vanilla/store.ts b/src/vanilla/store.ts index 929176a961..653e738a15 100644 --- a/src/vanilla/store.ts +++ b/src/vanilla/store.ts @@ -209,7 +209,7 @@ export type INTERNAL_PrdStore = Store /** * This is an experimental API and will be changed in the next minor. */ -export const INTERNAL_flushStoreHook: unique symbol = Symbol() +const INTERNAL_flushStoreHook = Symbol.for('JOTAI.EXPERIMENTAL.FLUSHSTOREHOOK') const buildStore = (...storeArgs: StoreArgs): Store => { const [ diff --git a/tests/vanilla/effect.test.ts b/tests/vanilla/effect.test.ts index 9ce918c314..3fe43e8d5a 100644 --- a/tests/vanilla/effect.test.ts +++ b/tests/vanilla/effect.test.ts @@ -1,8 +1,6 @@ import { expect, it, vi } from 'vitest' import type { Atom, Getter, Setter } from 'jotai/vanilla' import { atom, createStore } from 'jotai/vanilla' -// For this test only -import { INTERNAL_flushStoreHook } from '../../src/vanilla/store.js' type Store = ReturnType type GetAtomState = Parameters[0]>[0] @@ -84,6 +82,7 @@ function syncEffect(effect: Effect): Atom { }) } +const INTERNAL_flushStoreHook = Symbol.for('JOTAI.EXPERIMENTAL.FLUSHSTOREHOOK') const syncEffectChannelSymbol = Symbol() function ensureSyncEffectChannel(store: any) { From 2cf0a88f4ed6adb55b8e09e4b815b8a59975d88d Mon Sep 17 00:00:00 2001 From: Daishi Kato Date: Sat, 11 Jan 2025 08:59:39 +0900 Subject: [PATCH 03/23] Update src/vanilla/store.ts Co-authored-by: David Maskasky --- src/vanilla/store.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/vanilla/store.ts b/src/vanilla/store.ts index 653e738a15..0ff73b6a82 100644 --- a/src/vanilla/store.ts +++ b/src/vanilla/store.ts @@ -468,7 +468,7 @@ const buildStore = (...storeArgs: StoreArgs): Store => { visiting.add(a) // Push unvisited dependents onto the stack for (const [d, s] of getMountedOrPendingDependents(aState)) { - if (a !== d && !visiting.has(d)) { + if (!visiting.has(d)) { stack.push([d, s]) } } From b1efd91c3b85fb96130823820abede87e2473bd1 Mon Sep 17 00:00:00 2001 From: daishi Date: Sat, 11 Jan 2025 11:55:23 +0900 Subject: [PATCH 04/23] add a failing test --- tests/vanilla/store.test.tsx | 16 +++++++++++++++- 1 file changed, 15 insertions(+), 1 deletion(-) diff --git a/tests/vanilla/store.test.tsx b/tests/vanilla/store.test.tsx index 425063617e..6d3cee69d3 100644 --- a/tests/vanilla/store.test.tsx +++ b/tests/vanilla/store.test.tsx @@ -984,6 +984,21 @@ it('mounted atom should be recomputed eagerly', () => { expect(result).toEqual(['bRead', 'aCallback', 'bCallback']) }) +it('reading atom in write should notify subscribers', () => { + const a = atom(1) + const b = atom((get) => get(a) * 2) + const c = atom((get) => get(b) + 1) + const d = atom(null, (get, set) => { + set(a, 2) + get(b) + }) + const store = createStore() + const callback = vi.fn() + store.sub(c, callback) + store.set(d) + expect(callback).toHaveBeenCalledTimes(1) +}) + it('should process all atom listeners even if some of them throw errors', () => { const store = createStore() const a = atom(0) @@ -1091,7 +1106,6 @@ it('recomputes dependents of unmounted atoms', () => { const a = atom(0) a.debugLabel = 'a' const bRead = vi.fn((get: Getter) => { - console.log('bRead') return get(a) }) const b = atom(bRead) From 037c27f06d0bf821f9f95a56bd1076c46a23192a Mon Sep 17 00:00:00 2001 From: daishi Date: Sat, 11 Jan 2025 12:34:25 +0900 Subject: [PATCH 05/23] fix it --- src/vanilla/store.ts | 25 +++++++++++++------------ 1 file changed, 13 insertions(+), 12 deletions(-) diff --git a/src/vanilla/store.ts b/src/vanilla/store.ts index 653e738a15..32ed00eace 100644 --- a/src/vanilla/store.ts +++ b/src/vanilla/store.ts @@ -7,6 +7,7 @@ type AnyWritableAtom = WritableAtom type OnUnmount = () => void type Getter = Parameters[0] type Setter = Parameters[1] +type EpochNumber = number const isSelfAtom = (atom: AnyAtom, a: AnyAtom): boolean => atom.unstable_is ? atom.unstable_is(a) : a === atom @@ -89,7 +90,7 @@ type AtomState = { * Map of atoms that the atom depends on. * The map value is the epoch number of the dependency. */ - readonly d: Map + readonly d: Map /** * Set of atoms with pending promise that depend on the atom. * @@ -97,7 +98,7 @@ type AtomState = { */ readonly p: Set /** The epoch number of the atom. */ - n: number + n: EpochNumber /** Object to store mounted state of the atom. */ m?: Mounted // only available if the atom is mounted /** @@ -114,8 +115,6 @@ type AtomState = { v?: Value /** Atom error */ e?: AnyError - /** Indicates that the atom value has to be recomputed */ - x?: true } const isAtomStateInitialized = (atomState: AtomState) => @@ -233,6 +232,7 @@ const buildStore = (...storeArgs: StoreArgs): Store => { return atomState } + const invalidatedAtoms = new WeakMap() const changedAtoms = new Map() const unmountCallbacks = new Set<() => void>() const mountCallbacks = new Set<() => void>() @@ -279,7 +279,6 @@ const buildStore = (...storeArgs: StoreArgs): Store => { atomState.v = valueOrPromise } delete atomState.e - delete atomState.x if (!hasPrevValue || !Object.is(prevValue, atomState.v)) { ++atomState.n if (pendingPromise) { @@ -295,7 +294,7 @@ const buildStore = (...storeArgs: StoreArgs): Store => { // If the atom is mounted, we can use cached atom state. // because it should have been updated by dependencies. // We can't use the cache if the atom is invalidated. - if (atomState.m && !atomState.x) { + if (atomState.m && invalidatedAtoms.get(atom) !== atomState.n) { return atomState } // Otherwise, check if the dependencies have changed. @@ -385,7 +384,6 @@ const buildStore = (...storeArgs: StoreArgs): Store => { } catch (error) { delete atomState.v atomState.e = error - delete atomState.x ++atomState.n return atomState } finally { @@ -422,8 +420,8 @@ const buildStore = (...storeArgs: StoreArgs): Store => { const aState = stack.pop()! if (!visited.has(aState)) { visited.add(aState) - for (const s of getMountedOrPendingDependents(aState).values()) { - s.x = true + for (const [d, s] of getMountedOrPendingDependents(aState)) { + invalidatedAtoms.set(d, s.n) stack.push(s) } } @@ -439,7 +437,7 @@ const buildStore = (...storeArgs: StoreArgs): Store => { const topSortedReversed: [ atom: AnyAtom, atomState: AtomState, - epochNumber: number, + epochNumber: EpochNumber, ][] = [] const visiting = new WeakSet() const visited = new WeakSet() @@ -457,8 +455,11 @@ const buildStore = (...storeArgs: StoreArgs): Store => { // The algorithm calls for pushing onto the front of the list. For // performance, we will simply push onto the end, and then will iterate in // reverse order later. - if (aState.x) { + if (invalidatedAtoms.get(a) === aState.n) { topSortedReversed.push([a, aState, aState.n]) + } else { + invalidatedAtoms.delete(a) + changedAtoms.set(a, aState) } // Atom has been visited but not yet processed visited.add(a) @@ -493,7 +494,7 @@ const buildStore = (...storeArgs: StoreArgs): Store => { aState.u?.() } } - delete aState.x + invalidatedAtoms.delete(a) } } From 0fd8aa7a35dc89a75ce5d97919da45afe387e7c8 Mon Sep 17 00:00:00 2001 From: daishi Date: Sat, 11 Jan 2025 13:02:03 +0900 Subject: [PATCH 06/23] add some comments --- src/vanilla/store.ts | 13 ++++++++++++- tests/vanilla/store.test.tsx | 2 +- 2 files changed, 13 insertions(+), 2 deletions(-) diff --git a/src/vanilla/store.ts b/src/vanilla/store.ts index 1891fa8902..b90e0c72b3 100644 --- a/src/vanilla/store.ts +++ b/src/vanilla/store.ts @@ -84,6 +84,9 @@ type Mounted = { /** * Mutable atom state, * tracked for both mounted and unmounted atoms in a store. + * + * This should be garbage collectable. + * We can mutate it during atom read. (except for fields with TODO) */ type AtomState = { /** @@ -95,20 +98,26 @@ type AtomState = { * Set of atoms with pending promise that depend on the atom. * * This may cause memory leaks, but it's for the capability to continue promises + * TODO(daishi): revisit how to handle this */ readonly p: Set /** The epoch number of the atom. */ n: EpochNumber - /** Object to store mounted state of the atom. */ + /** + * Object to store mounted state of the atom. + * TODO(daishi): move this out of AtomState + */ m?: Mounted // only available if the atom is mounted /** * Listener to notify when the atom value is updated. * This is an experimental API and will be changed in the next minor. + * TODO(daishi): move this store hooks */ u?: () => void /** * Listener to notify when the atom is mounted or unmounted. * This is an experimental API and will be changed in the next minor. + * TODO(daishi): move this store hooks */ h?: () => void /** Atom value */ @@ -232,6 +241,8 @@ const buildStore = (...storeArgs: StoreArgs): Store => { return atomState } + // These are store state. + // As they are not garbage collectable, they shouldn't be mutated during atom read. const invalidatedAtoms = new WeakMap() const changedAtoms = new Map() const unmountCallbacks = new Set<() => void>() diff --git a/tests/vanilla/store.test.tsx b/tests/vanilla/store.test.tsx index 6d3cee69d3..229cb8126b 100644 --- a/tests/vanilla/store.test.tsx +++ b/tests/vanilla/store.test.tsx @@ -984,7 +984,7 @@ it('mounted atom should be recomputed eagerly', () => { expect(result).toEqual(['bRead', 'aCallback', 'bCallback']) }) -it('reading atom in write should notify subscribers', () => { +it('should notify subscription even with reading atom in write', () => { const a = atom(1) const b = atom((get) => get(a) * 2) const c = atom((get) => get(b) + 1) From da78ad8dcbb5e6729453f75778447f323c6edfe6 Mon Sep 17 00:00:00 2001 From: David Maskasky Date: Sat, 11 Jan 2025 21:31:36 +0900 Subject: [PATCH 07/23] add failing test --- tests/vanilla/store.test.tsx | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/tests/vanilla/store.test.tsx b/tests/vanilla/store.test.tsx index 229cb8126b..ef2e979307 100644 --- a/tests/vanilla/store.test.tsx +++ b/tests/vanilla/store.test.tsx @@ -1122,3 +1122,18 @@ it('recomputes dependents of unmounted atoms', () => { store.set(w) expect(bRead).not.toHaveBeenCalled() }) + +it('should not inf on subscribe or unsubscribe', async () => { + const store = createStore() + const countAtom = atom(0) + const effectAtom = atom( + (get) => get(countAtom), + (_, set) => set, + ) + effectAtom.onMount = (setAtom) => { + const set = setAtom() + set(countAtom, 1) + } + store.sub(effectAtom, () => {}) + expect(store.get(countAtom)).toBe(1) +}) From 417532ed36b95de20ca42540bbaa925d395cb9cd Mon Sep 17 00:00:00 2001 From: daishi Date: Sat, 11 Jan 2025 21:31:50 +0900 Subject: [PATCH 08/23] naive workaround --- src/vanilla/store.ts | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/src/vanilla/store.ts b/src/vanilla/store.ts index b90e0c72b3..dd55afc0fe 100644 --- a/src/vanilla/store.ts +++ b/src/vanilla/store.ts @@ -248,7 +248,12 @@ const buildStore = (...storeArgs: StoreArgs): Store => { const unmountCallbacks = new Set<() => void>() const mountCallbacks = new Set<() => void>() + let isFlushing = false const flushCallbacks = () => { + if (isFlushing) { + return + } + isFlushing = true const errors: unknown[] = [] const call = (fn: () => void) => { try { @@ -267,6 +272,7 @@ const buildStore = (...storeArgs: StoreArgs): Store => { mountCallbacks.clear() recomputeInvalidatedAtoms() } while (changedAtoms.size) + isFlushing = false if (errors.length) { throw errors[0] } From 9827d96f919451d7e99fcc7688d942260a854431 Mon Sep 17 00:00:00 2001 From: daishi Date: Sat, 11 Jan 2025 21:39:44 +0900 Subject: [PATCH 09/23] refactor readAtomState --- src/vanilla/store.ts | 17 ++++++++--------- 1 file changed, 8 insertions(+), 9 deletions(-) diff --git a/src/vanilla/store.ts b/src/vanilla/store.ts index dd55afc0fe..e303299279 100644 --- a/src/vanilla/store.ts +++ b/src/vanilla/store.ts @@ -330,6 +330,12 @@ const buildStore = (...storeArgs: StoreArgs): Store => { // Compute a new state for this atom. atomState.d.clear() let isSync = true + const mountDependenciesIfAsync = () => { + if (atomState.m) { + mountDependencies(atom, atomState) + flushCallbacks() + } + } const getter: Getter = (a: Atom) => { if (isSelfAtom(atom, a)) { const aState = ensureAtomState(a) @@ -350,8 +356,7 @@ const buildStore = (...storeArgs: StoreArgs): Store => { } finally { addDependency(atom, atomState, a, aState) if (!isSync) { - mountDependencies(atom, atomState) - flushCallbacks() + mountDependenciesIfAsync() } } } @@ -389,13 +394,7 @@ const buildStore = (...storeArgs: StoreArgs): Store => { setAtomStateValueOrPromise(atom, atomState, valueOrPromise) if (isPromiseLike(valueOrPromise)) { valueOrPromise.onCancel?.(() => controller?.abort()) - const complete = () => { - if (atomState.m) { - mountDependencies(atom, atomState) - flushCallbacks() - } - } - valueOrPromise.then(complete, complete) + valueOrPromise.then(mountDependenciesIfAsync, mountDependenciesIfAsync) } return atomState } catch (error) { From 699bd701d19ea899388add9196b8c81d656b342f Mon Sep 17 00:00:00 2001 From: daishi Date: Sat, 11 Jan 2025 23:52:06 +0900 Subject: [PATCH 10/23] is this better? --- src/vanilla/store.ts | 114 ++++++++++++++++--------------------------- 1 file changed, 43 insertions(+), 71 deletions(-) diff --git a/src/vanilla/store.ts b/src/vanilla/store.ts index e303299279..30fa88018f 100644 --- a/src/vanilla/store.ts +++ b/src/vanilla/store.ts @@ -248,12 +248,8 @@ const buildStore = (...storeArgs: StoreArgs): Store => { const unmountCallbacks = new Set<() => void>() const mountCallbacks = new Set<() => void>() - let isFlushing = false - const flushCallbacks = () => { - if (isFlushing) { - return - } - isFlushing = true + let inTransaction = 0 + const runWithTransaction = (fn: () => T): T => { const errors: unknown[] = [] const call = (fn: () => void) => { try { @@ -262,20 +258,31 @@ const buildStore = (...storeArgs: StoreArgs): Store => { errors.push(e) } } - do { - ;(store as any)[INTERNAL_flushStoreHook]?.() - changedAtoms.forEach((atomState) => atomState.m?.l.forEach(call)) - changedAtoms.clear() - unmountCallbacks.forEach(call) - unmountCallbacks.clear() - mountCallbacks.forEach(call) - mountCallbacks.clear() - recomputeInvalidatedAtoms() - } while (changedAtoms.size) - isFlushing = false + let result: T + ++inTransaction + try { + result = fn() + } finally { + if (inTransaction === 1) { + do { + if (changedAtoms.size) { + recomputeInvalidatedAtoms() + } + ;(store as any)[INTERNAL_flushStoreHook]?.() + changedAtoms.forEach((atomState) => atomState.m?.l.forEach(call)) + changedAtoms.clear() + unmountCallbacks.forEach(call) + unmountCallbacks.clear() + mountCallbacks.forEach(call) + mountCallbacks.clear() + } while (changedAtoms.size) + } + --inTransaction + } if (errors.length) { throw errors[0] } + return result } const setAtomStateValueOrPromise = ( @@ -332,8 +339,7 @@ const buildStore = (...storeArgs: StoreArgs): Store => { let isSync = true const mountDependenciesIfAsync = () => { if (atomState.m) { - mountDependencies(atom, atomState) - flushCallbacks() + runWithTransaction(() => mountDependencies(atom, atomState)) } } const getter: Getter = (a: Atom) => { @@ -518,14 +524,13 @@ 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) - try { + return runWithTransaction(() => { if (isSelfAtom(atom, a)) { if (!hasInitialValue(a)) { // NOTE technically possible but restricted as it may cause bugs @@ -544,31 +549,15 @@ const buildStore = (...storeArgs: StoreArgs): Store => { } else { return writeAtomState(a, ...args) } - } finally { - if (!isSync) { - recomputeInvalidatedAtoms() - flushCallbacks() - } - } - } - try { - return atomWrite(atom, getter, setter, ...args) - } finally { - isSync = false + }) } + return atomWrite(atom, getter, setter, ...args) } const writeAtom = ( atom: WritableAtom, ...args: Args - ): Result => { - try { - return writeAtomState(atom, ...args) - } finally { - recomputeInvalidatedAtoms() - flushCallbacks() - } - } + ): Result => runWithTransaction(() => writeAtomState(atom, ...args)) const mountDependencies = (atom: AnyAtom, atomState: AtomState) => { if (atomState.m && !isPendingPromise(atomState.v)) { @@ -610,31 +599,12 @@ 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) { - recomputeInvalidatedAtoms() - flushCallbacks() - } - } - } - try { - return fn() - } finally { - isSync = false - } - } const processOnMount = () => { - const onUnmount = createInvocationContext(() => - atomOnMount(atom, (...args) => setAtom(...args)), + const onUnmount = atomOnMount(atom, (...args) => + runWithTransaction(() => writeAtomState(atom, ...args)), ) if (onUnmount) { - mounted.u = () => createInvocationContext(onUnmount) + mounted.u = onUnmount } } mountCallbacks.add(processOnMount) @@ -671,15 +641,17 @@ const buildStore = (...storeArgs: StoreArgs): Store => { const subscribeAtom = (atom: AnyAtom, listener: () => void) => { const atomState = ensureAtomState(atom) - const mounted = mountAtom(atom, atomState) - const listeners = mounted.l - listeners.add(listener) - flushCallbacks() - return () => { - listeners.delete(listener) - unmountAtom(atom, atomState) - flushCallbacks() - } + return runWithTransaction(() => { + const mounted = mountAtom(atom, atomState) + const listeners = mounted.l + listeners.add(listener) + return () => { + runWithTransaction(() => { + listeners.delete(listener) + unmountAtom(atom, atomState) + }) + } + }) } const unstable_derive: Store['unstable_derive'] = (fn) => From 6c66cc8a1aad8689f4f95f7af549209ec1a15df3 Mon Sep 17 00:00:00 2001 From: David Maskasky Date: Sun, 12 Jan 2025 07:50:26 +0900 Subject: [PATCH 11/23] update test --- tests/vanilla/store.test.tsx | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/tests/vanilla/store.test.tsx b/tests/vanilla/store.test.tsx index ef2e979307..b6013f8c4a 100644 --- a/tests/vanilla/store.test.tsx +++ b/tests/vanilla/store.test.tsx @@ -1133,7 +1133,12 @@ it('should not inf on subscribe or unsubscribe', async () => { effectAtom.onMount = (setAtom) => { const set = setAtom() set(countAtom, 1) + return () => { + set(countAtom, 2) + } } - store.sub(effectAtom, () => {}) + const unsub = store.sub(effectAtom, () => {}) expect(store.get(countAtom)).toBe(1) + unsub() + expect(store.get(countAtom)).toBe(2) }) From 4a3039040664c1890611f84d7e5ab14972abe547 Mon Sep 17 00:00:00 2001 From: David Maskasky Date: Mon, 13 Jan 2025 18:46:41 +0900 Subject: [PATCH 12/23] add a test --- tests/vanilla/store.test.tsx | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/tests/vanilla/store.test.tsx b/tests/vanilla/store.test.tsx index b6013f8c4a..301af00df4 100644 --- a/tests/vanilla/store.test.tsx +++ b/tests/vanilla/store.test.tsx @@ -1142,3 +1142,15 @@ it('should not inf on subscribe or unsubscribe', async () => { unsub() expect(store.get(countAtom)).toBe(2) }) + +it('supports recursion in an atom subscriber', () => { + const a = atom(0) + const store = createStore() + store.sub(a, () => { + if (store.get(a) < 3) { + store.set(a, (v) => v + 1) + } + }) + store.set(a, 1) + expect(store.get(a)).toBe(3) +}) From 43f7e0062dd0a00af698d1bd9e340502a4804bb2 Mon Sep 17 00:00:00 2001 From: daishi Date: Mon, 13 Jan 2025 18:59:33 +0900 Subject: [PATCH 13/23] fix it --- src/vanilla/store.ts | 20 ++++++++++++++------ 1 file changed, 14 insertions(+), 6 deletions(-) diff --git a/src/vanilla/store.ts b/src/vanilla/store.ts index 30fa88018f..9b611bc6f0 100644 --- a/src/vanilla/store.ts +++ b/src/vanilla/store.ts @@ -173,6 +173,15 @@ const addDependency = ( aState.m?.t.add(atom) } +const clearThenForEach = ( + items: { values: () => Iterable; clear: () => void }, + callback: (item: T) => void, +) => { + const values = new Set(items.values()) + items.clear() + values.forEach(callback) +} + // internal & unstable type type StoreArgs = readonly [ getAtomState: (atom: Atom) => AtomState | undefined, @@ -269,12 +278,11 @@ const buildStore = (...storeArgs: StoreArgs): Store => { recomputeInvalidatedAtoms() } ;(store as any)[INTERNAL_flushStoreHook]?.() - changedAtoms.forEach((atomState) => atomState.m?.l.forEach(call)) - changedAtoms.clear() - unmountCallbacks.forEach(call) - unmountCallbacks.clear() - mountCallbacks.forEach(call) - mountCallbacks.clear() + clearThenForEach(changedAtoms, (atomState) => + atomState.m?.l.forEach(call), + ) + clearThenForEach(unmountCallbacks, call) + clearThenForEach(mountCallbacks, call) } while (changedAtoms.size) } --inTransaction From 304e9e6f7708d4dcf4a95ae2047d6f07a283c206 Mon Sep 17 00:00:00 2001 From: Daishi Kato Date: Tue, 14 Jan 2025 07:57:32 +0900 Subject: [PATCH 14/23] Update src/vanilla/store.ts Co-authored-by: David Maskasky --- src/vanilla/store.ts | 14 +++++++++----- 1 file changed, 9 insertions(+), 5 deletions(-) diff --git a/src/vanilla/store.ts b/src/vanilla/store.ts index 9b611bc6f0..54e61f022d 100644 --- a/src/vanilla/store.ts +++ b/src/vanilla/store.ts @@ -278,11 +278,15 @@ const buildStore = (...storeArgs: StoreArgs): Store => { recomputeInvalidatedAtoms() } ;(store as any)[INTERNAL_flushStoreHook]?.() - clearThenForEach(changedAtoms, (atomState) => - atomState.m?.l.forEach(call), - ) - clearThenForEach(unmountCallbacks, call) - clearThenForEach(mountCallbacks, call) + const callbacks = [ + ...changedAtoms.map((atomState) => atomState.m?.l ?? []), + ...unmountCallbacks, + ...mountCallbacks, + ] + callbacks.forEach(call) + changedAtoms.clear() + unmountCallbacks.clear() + mountCallbacks.clear() } while (changedAtoms.size) } --inTransaction From f53094af4c93348c46d6d5abcdd3e2c368b120a3 Mon Sep 17 00:00:00 2001 From: daishi Date: Tue, 14 Jan 2025 08:27:35 +0900 Subject: [PATCH 15/23] refactor --- src/vanilla/store.ts | 21 ++++++--------------- 1 file changed, 6 insertions(+), 15 deletions(-) diff --git a/src/vanilla/store.ts b/src/vanilla/store.ts index 54e61f022d..2667764aa5 100644 --- a/src/vanilla/store.ts +++ b/src/vanilla/store.ts @@ -173,15 +173,6 @@ const addDependency = ( aState.m?.t.add(atom) } -const clearThenForEach = ( - items: { values: () => Iterable; clear: () => void }, - callback: (item: T) => void, -) => { - const values = new Set(items.values()) - items.clear() - values.forEach(callback) -} - // internal & unstable type type StoreArgs = readonly [ getAtomState: (atom: Atom) => AtomState | undefined, @@ -278,15 +269,15 @@ const buildStore = (...storeArgs: StoreArgs): Store => { recomputeInvalidatedAtoms() } ;(store as any)[INTERNAL_flushStoreHook]?.() - const callbacks = [ - ...changedAtoms.map((atomState) => atomState.m?.l ?? []), - ...unmountCallbacks, - ...mountCallbacks, - ] - callbacks.forEach(call) + const callbacks = new Set<() => void>() + const add = callbacks.add.bind(callbacks) + changedAtoms.forEach((atomState) => atomState.m?.l.forEach(add)) + unmountCallbacks.forEach(add) + mountCallbacks.forEach(add) changedAtoms.clear() unmountCallbacks.clear() mountCallbacks.clear() + callbacks.forEach(call) } while (changedAtoms.size) } --inTransaction From 1696a2da7f360cd403d87d5b4e85e82f421d0cc6 Mon Sep 17 00:00:00 2001 From: daishi Date: Tue, 14 Jan 2025 09:03:50 +0900 Subject: [PATCH 16/23] add another failing test --- tests/vanilla/store.test.tsx | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/tests/vanilla/store.test.tsx b/tests/vanilla/store.test.tsx index 301af00df4..7202851f1d 100644 --- a/tests/vanilla/store.test.tsx +++ b/tests/vanilla/store.test.tsx @@ -1154,3 +1154,18 @@ it('supports recursion in an atom subscriber', () => { store.set(a, 1) expect(store.get(a)).toBe(3) }) + +it('allows subcribing to atoms during mount', () => { + const store = createStore() + const a = atom(0) + a.onMount = () => { + store.sub(b, () => {}) + } + const b = atom(0) + let bMounted = false + b.onMount = () => { + bMounted = true + } + store.sub(a, () => {}) + expect(bMounted).toBe(true) +}) From d7ffe0ae71558d03f238b4e23d4a6033e2815f87 Mon Sep 17 00:00:00 2001 From: daishi Date: Tue, 14 Jan 2025 09:04:40 +0900 Subject: [PATCH 17/23] fix it --- src/vanilla/store.ts | 16 +++++++++------- 1 file changed, 9 insertions(+), 7 deletions(-) diff --git a/src/vanilla/store.ts b/src/vanilla/store.ts index 2667764aa5..53f9e40156 100644 --- a/src/vanilla/store.ts +++ b/src/vanilla/store.ts @@ -264,21 +264,23 @@ const buildStore = (...storeArgs: StoreArgs): Store => { result = fn() } finally { if (inTransaction === 1) { - do { - if (changedAtoms.size) { - recomputeInvalidatedAtoms() - } + 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)) - unmountCallbacks.forEach(add) - mountCallbacks.forEach(add) changedAtoms.clear() + unmountCallbacks.forEach(add) unmountCallbacks.clear() + mountCallbacks.forEach(add) mountCallbacks.clear() callbacks.forEach(call) - } while (changedAtoms.size) + } } --inTransaction } From ab61bbd6f9aae34c7c9a55805a1b42abc3282a5d Mon Sep 17 00:00:00 2001 From: daishi Date: Tue, 14 Jan 2025 09:05:43 +0900 Subject: [PATCH 18/23] a dirty hack for syncEffect --- src/vanilla/store.ts | 1 + 1 file changed, 1 insertion(+) diff --git a/src/vanilla/store.ts b/src/vanilla/store.ts index 53f9e40156..b08a6d209c 100644 --- a/src/vanilla/store.ts +++ b/src/vanilla/store.ts @@ -264,6 +264,7 @@ const buildStore = (...storeArgs: StoreArgs): Store => { result = fn() } finally { if (inTransaction === 1) { + ;(store as any)[INTERNAL_flushStoreHook]?.() // FIXME we would like to avoid this while ( changedAtoms.size || unmountCallbacks.size || From bacc9c1dd9cb1f437ce582cc2b87308986fc396b Mon Sep 17 00:00:00 2001 From: David Maskasky Date: Tue, 14 Jan 2025 10:51:56 +0900 Subject: [PATCH 19/23] add effect test --- tests/vanilla/effect.test.ts | 48 ++++++++++++++++++++++++++++++++++++ 1 file changed, 48 insertions(+) diff --git a/tests/vanilla/effect.test.ts b/tests/vanilla/effect.test.ts index 3fe43e8d5a..f5dfb0d0e8 100644 --- a/tests/vanilla/effect.test.ts +++ b/tests/vanilla/effect.test.ts @@ -233,3 +233,51 @@ it('sets values to atoms without causing infinite loop', () => { unsub() expect(effect).toBeCalledTimes(2) }) + +// TODO: consider removing this after we provide a new syncEffect implementation +it('supports recursive setting synchronous in read', async () => { + const store = createStore() + const a = atom(0) + const refreshAtom = atom(0) + type Ref = { + isMounted?: true + isRecursing?: true + set: Setter + } + const refAtom = atom( + () => ({}) as Ref, + (get, set) => { + const ref = get(refAtom) + ref.isMounted = true + ref.set = set + set(refreshAtom, (v) => v + 1) + }, + ) + refAtom.onMount = (mount) => mount() + const effectAtom = atom((get) => { + get(refreshAtom) + const ref = get(refAtom) + if (!ref.isMounted) { + return + } + const recurse = ((...args) => { + ref.isRecursing = true + const value = ref.set(...args) + delete ref.isRecursing + return value + }) as Setter + function runEffect() { + const v = get(a) + if (v < 5) { + recurse(a, (v) => v + 1) + } + } + if (ref.isRecursing) { + return runEffect() + } + return Promise.resolve().then(runEffect) + }) + store.sub(effectAtom, () => {}) + await Promise.resolve() + expect(store.get(a)).toBe(5) +}) From 22cd11da926c1ee803c4d2fd23d66fdb5f6d6104 Mon Sep 17 00:00:00 2001 From: daishi Date: Tue, 14 Jan 2025 10:53:07 +0900 Subject: [PATCH 20/23] would it be acceptable? --- tests/vanilla/effect.test.ts | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/tests/vanilla/effect.test.ts b/tests/vanilla/effect.test.ts index f5dfb0d0e8..13b5e89921 100644 --- a/tests/vanilla/effect.test.ts +++ b/tests/vanilla/effect.test.ts @@ -272,12 +272,9 @@ it('supports recursive setting synchronous in read', async () => { recurse(a, (v) => v + 1) } } - if (ref.isRecursing) { - return runEffect() - } return Promise.resolve().then(runEffect) }) store.sub(effectAtom, () => {}) - await Promise.resolve() + await new Promise((r) => setTimeout(r)) expect(store.get(a)).toBe(5) }) From 9043fb8f7a35b1f811b4ec64e61919eccd9d547a Mon Sep 17 00:00:00 2001 From: daishi Date: Tue, 14 Jan 2025 10:59:55 +0900 Subject: [PATCH 21/23] possible fix for typing --- tests/vanilla/effect.test.ts | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/tests/vanilla/effect.test.ts b/tests/vanilla/effect.test.ts index 13b5e89921..b771fdc59b 100644 --- a/tests/vanilla/effect.test.ts +++ b/tests/vanilla/effect.test.ts @@ -1,5 +1,5 @@ import { expect, it, vi } from 'vitest' -import type { Atom, Getter, Setter } from 'jotai/vanilla' +import type { Atom, Getter, Setter, WritableAtom } from 'jotai/vanilla' import { atom, createStore } from 'jotai/vanilla' type Store = ReturnType @@ -260,12 +260,15 @@ it('supports recursive setting synchronous in read', async () => { if (!ref.isMounted) { return } - const recurse = ((...args) => { + const recurse = ( + a: WritableAtom, + ...args: Args + ): Result => { ref.isRecursing = true - const value = ref.set(...args) + const value = ref.set(a, ...args) delete ref.isRecursing - return value - }) as Setter + return value as Result + } function runEffect() { const v = get(a) if (v < 5) { From fb0510f2a0e137cf9dfeb3123c8977d30b0d664f Mon Sep 17 00:00:00 2001 From: daishi Date: Tue, 14 Jan 2025 11:16:57 +0900 Subject: [PATCH 22/23] enjoying the puzzle --- tests/vanilla/effect.test.ts | 18 +++++++++++++----- 1 file changed, 13 insertions(+), 5 deletions(-) diff --git a/tests/vanilla/effect.test.ts b/tests/vanilla/effect.test.ts index b771fdc59b..f64688d0ba 100644 --- a/tests/vanilla/effect.test.ts +++ b/tests/vanilla/effect.test.ts @@ -241,11 +241,11 @@ it('supports recursive setting synchronous in read', async () => { const refreshAtom = atom(0) type Ref = { isMounted?: true - isRecursing?: true + recursing: number set: Setter } const refAtom = atom( - () => ({}) as Ref, + () => ({ recursing: 0 }) as Ref, (get, set) => { const ref = get(refAtom) ref.isMounted = true @@ -264,9 +264,8 @@ it('supports recursive setting synchronous in read', async () => { a: WritableAtom, ...args: Args ): Result => { - ref.isRecursing = true + ++ref.recursing const value = ref.set(a, ...args) - delete ref.isRecursing return value as Result } function runEffect() { @@ -275,9 +274,18 @@ it('supports recursive setting synchronous in read', async () => { recurse(a, (v) => v + 1) } } + if (ref.recursing) { + let prevRecursing = ref.recursing + do { + prevRecursing = ref.recursing + runEffect() + } while (prevRecursing !== ref.recursing) + ref.recursing = 0 + return Promise.resolve() + } return Promise.resolve().then(runEffect) }) store.sub(effectAtom, () => {}) - await new Promise((r) => setTimeout(r)) + await Promise.resolve() expect(store.get(a)).toBe(5) }) From c601bd9c536dc000b9efb5bff456a8650b23e2fb Mon Sep 17 00:00:00 2001 From: daishi Date: Tue, 14 Jan 2025 11:38:03 +0900 Subject: [PATCH 23/23] fix the dirty hack ab61bbd6f9aae34c7c9a55805a1b42abc3282a5d --- src/vanilla/store.ts | 1 - tests/vanilla/effect.test.ts | 24 +++++++++++++++--------- 2 files changed, 15 insertions(+), 10 deletions(-) diff --git a/src/vanilla/store.ts b/src/vanilla/store.ts index b08a6d209c..53f9e40156 100644 --- a/src/vanilla/store.ts +++ b/src/vanilla/store.ts @@ -264,7 +264,6 @@ const buildStore = (...storeArgs: StoreArgs): Store => { result = fn() } finally { if (inTransaction === 1) { - ;(store as any)[INTERNAL_flushStoreHook]?.() // FIXME we would like to avoid this while ( changedAtoms.size || unmountCallbacks.size || diff --git a/tests/vanilla/effect.test.ts b/tests/vanilla/effect.test.ts index f64688d0ba..7a1dfb9906 100644 --- a/tests/vanilla/effect.test.ts +++ b/tests/vanilla/effect.test.ts @@ -19,15 +19,21 @@ type Ref = { function syncEffect(effect: Effect): Atom { const refAtom = atom(() => ({ inProgress: 0, epoch: 0 })) const refreshAtom = atom(0) - const internalAtom = atom((get) => { - get(refreshAtom) - const ref = get(refAtom) - if (ref.inProgress) { - return ref.epoch - } - ref.get = get - return ++ref.epoch - }) + const internalAtom = atom( + (get) => { + get(refreshAtom) + const ref = get(refAtom) + if (ref.inProgress) { + return ref.epoch + } + ref.get = get + return ++ref.epoch + }, + () => {}, + ) + internalAtom.onMount = () => { + return () => {} + } internalAtom.unstable_onInit = (store) => { const ref = store.get(refAtom) const runEffect = () => {