diff --git a/src/vanilla/store.ts b/src/vanilla/store.ts index d0d9a49453..b4ce043368 100644 --- a/src/vanilla/store.ts +++ b/src/vanilla/store.ts @@ -256,6 +256,27 @@ const buildStore = (getAtomState: StoreArgs[0]): Store => { } const flushPending = (pending: Pending, shouldProcessFinalizers = true) => { + console.log( + pending[3].size + + '.'.repeat(30) + + '\n' + + new Error().stack + ?.split('\n') + .filter( + (l) => + l.includes('/Users/dmaskasky/Code/jotai/src') || + l.includes('/Users/dmaskasky/Code/jotai/test'), + ) + .map((l) => + ' - ' + l.includes('/Users/dmaskasky/Code/jotai/src') + ? l.trim().split(' ')[1] + : l.trim().split(' '), + ) + .join('\n') + + '\n' + + ','.repeat(30) + + '\n', + ) do { while (pending[1].size || pending[2].size) { pending[0].clear() @@ -266,13 +287,15 @@ const buildStore = (getAtomState: StoreArgs[0]): Store => { atomStates.forEach((atomState) => atomState.m?.l.forEach((l) => l())) functions.forEach((fn) => fn()) } - // Process syncEffects + // Process finalizers after all atoms are updated if (shouldProcessFinalizers) { const pendingFinalizers = Array.from(pending[3]) for (const finalizerAtom of pendingFinalizers) { const atomState = getAtomState(finalizerAtom) if (atomState.m) { - const get = (a: Atom) => getAtomState(a).v! + const get = function get(a: Atom) { + return getAtomState(a).v! + } finalizerAtom.onAfterFlushPending!(get) pending[3].delete(finalizerAtom) } @@ -499,6 +522,7 @@ const buildStore = (getAtomState: StoreArgs[0]): Store => { } } if (hasChangedDeps) { + console.log('recompute', a, aState, prevEpochNumber, hasChangedDeps) readAtomState(pending, a, aState, markedAtoms) mountDependencies(pending, a, aState) if (prevEpochNumber !== aState.n) { @@ -619,6 +643,7 @@ const buildStore = (getAtomState: StoreArgs[0]): Store => { } }) } + addPendingFinalizer(pending, atom) } return atomState.m } diff --git a/tests/vanilla/effect.test.ts b/tests/vanilla/effect.test.ts index 762bf255c1..ba8cd1250b 100644 --- a/tests/vanilla/effect.test.ts +++ b/tests/vanilla/effect.test.ts @@ -1,25 +1,44 @@ -import { - createStore, - atom, - type Getter, - type Setter, - type Atom, -} from 'jotai/vanilla' +import type { Getter, Setter, Atom } from 'jotai/vanilla' +import { createStore, atom } from 'jotai/vanilla' import { vi, expect, it } from 'vitest' +type AnyAtom = Atom +type Store = ReturnType +type PrdStore = Exclude +type DevStoreRev4 = Omit< + Extract, + keyof PrdStore +> + +function isDevStoreRev4(store: Store): store is PrdStore & DevStoreRev4 { + return ( + typeof (store as DevStoreRev4).dev4_get_internal_weak_map === 'function' && + typeof (store as DevStoreRev4).dev4_get_mounted_atoms === 'function' && + typeof (store as DevStoreRev4).dev4_restore_atoms === 'function' + ) +} + +function assertIsDevStore( + store: Store, +): asserts store is PrdStore & DevStoreRev4 { + if (!isDevStoreRev4(store)) { + throw new Error('Store is not a dev store') + } +} + type Cleanup = () => void -type AtomEffect = Atom & { effect: Effect } type Effect = (get: Getter, set: Setter) => void | Cleanup type Ref = { getter?: Getter setter?: Setter cleanup?: Cleanup | void isPending: boolean + deps: Set } function atomSyncEffect(effect: Effect) { const refAtom = atom( - () => ({}) as Ref, + () => ({ deps: new Set() }) as Ref, (get, set) => { const ref = get(refAtom) ref.setter = set @@ -28,6 +47,7 @@ function atomSyncEffect(effect: Effect) { ref.cleanup?.() ref.cleanup = undefined ref.isPending = false + ref.deps.clear() } }, ) @@ -45,32 +65,66 @@ function atomSyncEffect(effect: Effect) { const effectAtom = Object.assign( atom((get) => { const ref = get(refAtom) - ref.getter = get + ref.getter = (a: Atom): Value => { + ref.deps.add(a) + return get(a) + } + ref.deps.forEach(get) ref.isPending = true return }), - { effect, onAfterFlushPending }, + { effect }, ) + effectAtom.onAfterFlushPending = onAfterFlushPending return effectAtom } -it('responds to changes to atoms', () => { - const atomState = new Map() - const store = createStore().unstable_derive(() => { - return [ - (atom) => { - if (!atomState.has(atom)) { - atomState.set(atom, { - name: atom.debugLabel, - d: new Map(), - p: new Set(), - n: 0, - }) - } - return atomState.get(atom) - }, - ] +it('responds to changes to atoms when subscribed', () => { + const store = createStore() + assertIsDevStore(store) + const weakMap = store.dev4_get_internal_weak_map() + const a = atom(1) + a.debugLabel = 'a' + const b = atom(1) + b.debugLabel = 'b' + const w = atom(null, (_get, set, value: number) => { + set(a, value) + set(b, value) }) + w.debugLabel = 'w' + const results: number[] = [] + const cleanup = vi.fn() + const effect = vi.fn((get: Getter) => { + results.push(get(a) * 10 + get(b)) + return cleanup + }) + const e = atomSyncEffect(effect) + e.debugLabel = 'e' + expect(results).toStrictEqual([]) + const unsub = store.sub(e, () => {}) // mount syncEffect + expect(effect).toBeCalledTimes(1) + expect(results).toStrictEqual([11]) // initial values at time of effect mount + store.set(a, 2) + expect(results).toStrictEqual([11, 21]) // store.set(a, 2) + store.set(b, 2) + expect(results).toStrictEqual([11, 21, 22]) // store.set(b, 2) + store.set(w, 3) + // intermediate state of '32' should not be recorded since the effect runs _after_ graph has been computed + expect(results).toStrictEqual([11, 21, 22, 33]) // store.set(w, 3) + expect(cleanup).toBeCalledTimes(3) + expect(effect).toBeCalledTimes(4) + expect(Array.from(weakMap.get(e)!.d.keys())).toEqual( + expect.arrayContaining([a, b]), + ) + unsub() + expect(cleanup).toBeCalledTimes(4) + expect(effect).toBeCalledTimes(4) +}) + +it.only('responds to changes to atoms when mounted with get', () => { + const store = createStore() + assertIsDevStore(store) + const weakMap = store.dev4_get_internal_weak_map() const a = atom(1) a.debugLabel = 'a' const b = atom(1) @@ -82,24 +136,31 @@ it('responds to changes to atoms', () => { w.debugLabel = 'w' const results: number[] = [] const cleanup = vi.fn() - const effectFn = vi.fn((get: Getter) => { + const effect = vi.fn((get: Getter) => { results.push(get(a) * 10 + get(b)) return cleanup }) - const e = atomSyncEffect(effectFn) + const e = atomSyncEffect(effect) e.debugLabel = 'e' + const d = atom((get) => get(e)) + d.debugLabel = 'd' expect(results).toStrictEqual([]) - const subscriber = vi.fn() - store.sub(e, subscriber) // mount syncEffect + const unsub = store.sub(d, () => {}) // mount syncEffect + expect(effect).toBeCalledTimes(1) expect(results).toStrictEqual([11]) // initial values at time of effect mount store.set(a, 2) expect(results).toStrictEqual([11, 21]) // store.set(a, 2) store.set(b, 2) expect(results).toStrictEqual([11, 21, 22]) // store.set(b, 2) store.set(w, 3) - // intermediate state of '32' should not be recorded as effect runs _after_ graph has been computed + // intermediate state of '32' should not be recorded since the effect runs _after_ graph has been computed expect(results).toStrictEqual([11, 21, 22, 33]) // store.set(w, 3) - expect(subscriber).toBeCalledTimes(0) - expect(effectFn).toBeCalledTimes(4) expect(cleanup).toBeCalledTimes(3) + expect(effect).toBeCalledTimes(4) + expect(Array.from(weakMap.get(e)!.d.keys())).toEqual( + expect.arrayContaining([a, b]), + ) + unsub() + expect(cleanup).toBeCalledTimes(4) + expect(effect).toBeCalledTimes(4) })