(working - get only) onAfterFlushPending based synchronous effects

src/vanilla/store.ts

@@ -256,6 +256,28 @@ const buildStore = (getAtomState: StoreArgs[0]): Store => {
   }
 
   const flushPending = (pending: Pending, shouldProcessFinalizers = true) => {
+    // TODO: remove this after debugging is complete
+    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 +288,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 = <Value>(a: Atom<Value>) => getAtomState(a).v!
+            const get = function get<Value>(a: Atom<Value>) {
+              return getAtomState(a).v!
+            }
             finalizerAtom.onAfterFlushPending!(get)
             pending[3].delete(finalizerAtom)
           }
@@ -499,6 +523,8 @@ const buildStore = (getAtomState: StoreArgs[0]): Store => {
         }
       }
       if (hasChangedDeps) {
+        // TODO: remove this after debugging is complete
+        console.log('recompute', a, aState, prevEpochNumber, hasChangedDeps)
         readAtomState(pending, a, aState, markedAtoms)
         mountDependencies(pending, a, aState)
         if (prevEpochNumber !== aState.n) {
@@ -619,6 +645,7 @@ const buildStore = (getAtomState: StoreArgs[0]): Store => {
           }
         })
       }
+      addPendingFinalizer(pending, atom)
     }
     return atomState.m
   }

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<unknown>
+type Store = ReturnType<typeof createStore>
+type PrdStore = Exclude<Store, { dev4_get_internal_weak_map: any }>
+type DevStoreRev4 = Omit<
+  Extract<Store, { dev4_get_internal_weak_map: any }>,
+  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<undefined> & { effect: Effect }
 type Effect = (get: Getter, set: Setter) => void | Cleanup
 type Ref = {
   getter?: Getter
   setter?: Setter
   cleanup?: Cleanup | void
   isPending: boolean
+  deps: Set<AnyAtom>
 }
 
 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 = <Value>(a: Atom<Value>): 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('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)
 })