propose: synchronous effects

src/vanilla/store.ts

@@ -19,6 +19,15 @@ const hasInitialValue = <T extends Atom<AnyValue>>(
 const isActuallyWritableAtom = (atom: AnyAtom): atom is AnyWritableAtom =>
   !!(atom as AnyWritableAtom).write
 
+type Cleanup = () => void
+type Effect = (get: Getter, set: Setter) => void | Cleanup
+
+type SyncEffectAtom = Atom<undefined> & { effect: Effect }
+
+const isSyncEffectAtom = (
+  atom: AnyAtom | Atom<undefined>,
+): atom is SyncEffectAtom => !!(atom as SyncEffectAtom).effect
+
 //
 // Cancelable Promise
 //
@@ -165,9 +174,15 @@ type Pending = readonly [
   dependents: Map<AnyAtom, Set<AnyAtom>>,
   atomStates: Map<AnyAtom, AtomState>,
   functions: Set<() => void>,
+  effects: Set<SyncEffectAtom>,
 ]
 
-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 +204,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 +220,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 = <Value>(
   atom: Atom<Value>,
   originAtomState?: AtomState,
@@ -253,6 +263,67 @@ 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 = <V>(a: Atom<V>) => {
+      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 {
+      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 +361,10 @@ const buildStore = (getAtomState: StoreArgs[0]): Store => {
     atomState: AtomState<Value>,
     dirtyAtoms?: Set<AnyAtom>,
   ): AtomState<Value> => {
+    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 +546,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 +566,7 @@ const buildStore = (getAtomState: StoreArgs[0]): Store => {
     atomState: AtomState<Value>,
     ...args: Args
   ): Result => {
+    let isSync = true
     const getter: Getter = <V>(a: Atom<V>) =>
       returnAtomValue(readAtomState(pending, a, getAtomState(a, atomState)))
     const setter: Setter = <V, As extends unknown[], R>(
@@ -513,11 +592,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 = <Value, Args extends unknown[], Result>(
@@ -587,6 +669,9 @@ const buildStore = (getAtomState: StoreArgs[0]): Store => {
           }
         })
       }
+      if (isSyncEffectAtom(atom)) {
+        addPendingEffect(pending, atom)
+      }
     }
     return atomState.m
   }

tests/vanilla/effect.test.ts

@@ -0,0 +1,71 @@
+import { createStore, atom, type Getter, type Setter } from 'jotai/vanilla'
+import { vi, expect, it } from 'vitest'
+
+type Cleanup = () => void
+type Effect = (get: Getter, set: Setter) => void | Cleanup
+
+function atomSyncEffect(effect: Effect) {
+  const refAtom = atom<{
+    cleanup?: Cleanup | void
+  }>(() => ({}))
+  return Object.assign(
+    atom(() => {}),
+    {
+      effect: (get: Getter, set: Setter) => {
+        const ref = get(refAtom)
+        ref.cleanup?.()
+        ref.cleanup = effect(get, set)
+      },
+    },
+  )
+}
+
+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)
+})