From 9e458baf6d969d366a3af70fd168dd1dea89a3f7 Mon Sep 17 00:00:00 2001 From: David Maskasky Date: Thu, 7 Nov 2024 17:51:31 -0800 Subject: [PATCH] propose: synchronous effects --- src/vanilla/store.ts | 121 ++++++++++++++++++++++++++++++----- tests/vanilla/effect.test.ts | 60 +++++++++++++++++ 2 files changed, 165 insertions(+), 16 deletions(-) create mode 100644 tests/vanilla/effect.test.ts diff --git a/src/vanilla/store.ts b/src/vanilla/store.ts index 3f555f35f9..bd629c2191 100644 --- a/src/vanilla/store.ts +++ b/src/vanilla/store.ts @@ -19,6 +19,16 @@ const hasInitialValue = >( const isActuallyWritableAtom = (atom: AnyAtom): atom is AnyWritableAtom => !!(atom as AnyWritableAtom).write +type Cleanup = () => void + +export type Effect = (get: Getter, set: Setter) => void | Cleanup + +type SyncEffectAtom = Atom & { effect: Effect } + +const isSyncEffectAtom = ( + atom: AnyAtom | Atom, +): atom is SyncEffectAtom => !!(atom as SyncEffectAtom).effect + // // Cancelable Promise // @@ -104,6 +114,8 @@ type AtomState = { v?: Value /** Atom error */ e?: AnyError + /** Cleanup function for syncEffects */ + c?: void | Cleanup } const isAtomStateInitialized = (atomState: AtomState) => @@ -165,9 +177,15 @@ type Pending = readonly [ dependents: Map>, atomStates: Map, functions: Set<() => void>, + effects: Set, ] -const createPending = (): Pending => [new Map(), new Map(), new Set()] +const createPending = (): Pending => [ + new Map(), + new Map(), + new Set(), + new Set(), +] const addPendingAtom = ( pending: Pending, @@ -189,6 +207,13 @@ const addPendingDependent = ( if (dependents) { dependents.add(dependent) } + if (isSyncEffectAtom(dependent)) { + addPendingEffect(pending, dependent) + } +} + +const addPendingEffect = (pending: Pending, effectAtom: SyncEffectAtom) => { + pending[3].add(effectAtom) } const getPendingDependents = (pending: Pending, atom: AnyAtom) => @@ -198,18 +223,6 @@ const addPendingFunction = (pending: Pending, fn: () => void) => { pending[2].add(fn) } -const flushPending = (pending: Pending) => { - while (pending[1].size || pending[2].size) { - pending[0].clear() - const atomStates = new Set(pending[1].values()) - pending[1].clear() - const functions = new Set(pending[2]) - pending[2].clear() - atomStates.forEach((atomState) => atomState.m?.l.forEach((l) => l())) - functions.forEach((fn) => fn()) - } -} - type GetAtomState = ( atom: Atom, originAtomState?: AtomState, @@ -253,6 +266,68 @@ const buildStore = (getAtomState: StoreArgs[0]): Store => { debugMountedAtoms = new Set() } + const flushPending = (pending: Pending, shouldProcessEffects = true) => { + do { + while (pending[1].size || pending[2].size) { + pending[0].clear() + const atomStates = new Set(pending[1].values()) + pending[1].clear() + const functions = new Set(pending[2]) + pending[2].clear() + atomStates.forEach((atomState) => atomState.m?.l.forEach((l) => l())) + functions.forEach((fn) => fn()) + } + // Process syncEffects + if (shouldProcessEffects) { + for (const effectAtom of pending[3]) { + const atomState = getAtomState(effectAtom) + if (atomState.m) { + runEffect(pending, effectAtom, atomState) + } + } + pending[3].clear() + } + } while (pending[1].size || pending[2].size) + } + + const runEffect = ( + pending: Pending, + effectAtom: SyncEffectAtom, + atomState: AtomState, + ) => { + let isSync = true + const getter: Getter = (a: Atom) => { + if (isSyncEffectAtom(a)) { + return returnAtomValue(atomState) as V + } + // a !== atom + const aState = readAtomState(pending, a, getAtomState(a), undefined) + if (isSync) { + addDependency(pending, effectAtom, atomState, a, aState) + } else { + const pending = createPending() + addDependency(pending, effectAtom, atomState, a, aState) + mountDependencies(pending, effectAtom, atomState) + flushPending(pending) + } + return returnAtomValue(aState) + } + const setter: Setter = (a, ...args) => { + return writeAtomState(pending, a, getAtomState(a), ...args) + } + try { + atomState.c?.() + atomState.c = effectAtom.effect(getter, setter) + mountDependencies(pending, effectAtom, atomState) + } catch (error) { + delete atomState.v + atomState.e = error + } finally { + ++atomState.n + isSync = false + } + } + const setAtomStateValueOrPromise = ( atom: AnyAtom, atomState: AtomState, @@ -290,6 +365,10 @@ const buildStore = (getAtomState: StoreArgs[0]): Store => { atomState: AtomState, dirtyAtoms?: Set, ): AtomState => { + if (isSyncEffectAtom(atom)) { + atomState.v = undefined as Value + return atomState + } // See if we can skip recomputing this atom. if (isAtomStateInitialized(atomState)) { // If the atom is mounted, we can use cached atom state. @@ -471,6 +550,9 @@ const buildStore = (getAtomState: StoreArgs[0]): Store => { } } if (hasChangedDeps) { + if (isSyncEffectAtom(a)) { + addPendingEffect(pending, a) + } readAtomState(pending, a, aState, markedAtoms) mountDependencies(pending, a, aState) if (prevEpochNumber !== aState.n) { @@ -488,6 +570,7 @@ const buildStore = (getAtomState: StoreArgs[0]): Store => { atomState: AtomState, ...args: Args ): Result => { + let isSync = true const getter: Getter = (a: Atom) => returnAtomValue(readAtomState(pending, a, getAtomState(a, atomState))) const setter: Setter = ( @@ -513,11 +596,14 @@ const buildStore = (getAtomState: StoreArgs[0]): Store => { } else { r = writeAtomState(pending, a, aState, ...args) as R } - flushPending(pending) + flushPending(pending, !isSync) return r as R } - const result = atom.write(getter, setter, ...args) - return result + try { + return atom.write(getter, setter, ...args) + } finally { + isSync = false + } } const writeAtom = ( @@ -587,6 +673,9 @@ const buildStore = (getAtomState: StoreArgs[0]): Store => { } }) } + if (isSyncEffectAtom(atom)) { + addPendingEffect(pending, atom) + } } return atomState.m } diff --git a/tests/vanilla/effect.test.ts b/tests/vanilla/effect.test.ts new file mode 100644 index 0000000000..7181c8b5f9 --- /dev/null +++ b/tests/vanilla/effect.test.ts @@ -0,0 +1,60 @@ +import { createStore, atom, atomSyncEffect, type Getter } from 'jotai/vanilla' +import type { Effect } from 'jotai/vanilla/store' +import { vi, expect, it } from 'vitest' + +function atomSyncEffect(effect: Effect) { + return Object.assign( + atom(() => undefined), + { effect }, + ) +} + +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) + }, + ] + }) + 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 effectFn = vi.fn((get: Getter) => { + results.push(get(a) * 10 + get(b)) + return cleanup + }) + const e = atomSyncEffect(effectFn) + e.debugLabel = 'e' + expect(results).toStrictEqual([]) + const subscriber = vi.fn() + store.sub(e, subscriber) // mount syncEffect + 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 + expect(results).toStrictEqual([11, 21, 22, 33]) // store.set(w, 3) + expect(subscriber).toBeCalledTimes(0) + expect(effectFn).toBeCalledTimes(4) + expect(cleanup).toBeCalledTimes(3) +})