support recurse

src/vanilla/store.ts

@@ -256,28 +256,6 @@ 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()
@@ -523,8 +501,6 @@ 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) {

tests/vanilla/effect.test.ts

@@ -3,46 +3,45 @@ 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 GetterWithPeak = Getter & { peak: Getter }
+type SetterWithRecurse = Setter & { recurse: Setter }
 type Cleanup = () => void
-type Effect = (get: Getter, set: Setter) => void | Cleanup
+type Effect = (get: GetterWithPeak, set: SetterWithRecurse) => void | Cleanup
 type Ref = {
-  getter?: Getter
-  setter?: Setter
+  get: GetterWithPeak
+  set?: SetterWithRecurse
   cleanup?: Cleanup | void
+  fromCleanup: boolean
+  inProgress: number
   isPending: boolean
   deps: Set<AnyAtom>
 }
 
 function atomSyncEffect(effect: Effect) {
   const refAtom = atom(
-    () => ({ deps: new Set() }) as Ref,
+    () => ({ deps: new Set(), inProgress: 0 }) as Ref,
     (get, set) => {
       const ref = get(refAtom)
-      ref.setter = set
-      ref.isPending = true
+      ref.get.peak ??= get
+      const setter: Setter = (a, ...args) => {
+        try {
+          ++ref.inProgress
+          return set(a, ...args)
+        } finally {
+          --ref.inProgress
+        }
+      }
+      const recurse: Setter = (a, ...args) => {
+        if (ref.fromCleanup) {
+          if (process.env.NODE_ENV !== 'production') {
+            console.warn('cannot recurse inside cleanup')
+          }
+          return undefined as any
+        }
+        return set(a, ...args)
+      }
+      ref.set ??= Object.assign(setter, { recurse })
+      ref.isPending = ref.inProgress === 0
       return () => {
         ref.cleanup?.()
         ref.cleanup = undefined
@@ -51,38 +50,49 @@ function atomSyncEffect(effect: Effect) {
       }
     },
   )
-  refAtom.onMount = (setSelf) => setSelf()
-  refAtom.debugPrivate = true
-  function onAfterFlushPending(get: Getter) {
+  refAtom.onMount = (mount) => mount()
+  const internalAtom = atom((get) => {
+    const ref = get(refAtom)
+    ref.get ??= ((a) => {
+      ref.deps.add(a)
+      return get(a)
+    }) as Getter & { peak: Getter }
+    ref.deps.forEach(get)
+    ref.isPending = true
+  })
+  internalAtom.onAfterFlushPending = (get: Getter) => {
     const ref = get(refAtom)
-    if (!ref.isPending) {
+    if (!ref.isPending || ref.inProgress > 0) {
       return
     }
     ref.isPending = false
     ref.cleanup?.()
-    ref.cleanup = effectAtom.effect(ref.getter!, ref.setter!)
+    const cleanup = effectAtom.effect(ref.get!, ref.set!)
+    ref.cleanup =
+      typeof cleanup === 'function'
+        ? () => {
+            try {
+              ref.fromCleanup = true
+              cleanup()
+            } finally {
+              ref.fromCleanup = false
+            }
+          }
+        : undefined
+  }
+  if (process.env.NODE_ENV !== 'production') {
+    refAtom.debugPrivate = true
+    internalAtom.debugPrivate = true
   }
   const effectAtom = Object.assign(
-    atom((get) => {
-      const ref = get(refAtom)
-      ref.getter = <Value>(a: Atom<Value>): Value => {
-        ref.deps.add(a)
-        return get(a)
-      }
-      ref.deps.forEach(get)
-      ref.isPending = true
-      return
-    }),
+    atom((get) => get(internalAtom)),
     { effect },
   )
-  effectAtom.onAfterFlushPending = onAfterFlushPending
   return effectAtom
 }
 
 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)
@@ -100,31 +110,29 @@ it('responds to changes to atoms when subscribed', () => {
   })
   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)
+  expect(results).toStrictEqual([11, 21])
   store.set(b, 2)
-  expect(results).toStrictEqual([11, 21, 22]) // store.set(b, 2)
+  expect(results).toStrictEqual([11, 21, 22])
   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(results).toStrictEqual([11, 21, 22, 33])
   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)
+  store.set(a, 4)
+  // the effect is unmounted so no more updates
+  expect(results).toStrictEqual([11, 21, 22, 33])
+  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)
@@ -144,23 +152,71 @@ it('responds to changes to atoms when mounted with get', () => {
   e.debugLabel = 'e'
   const d = atom((get) => get(e))
   d.debugLabel = 'd'
-  expect(results).toStrictEqual([])
   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)
+  expect(results).toStrictEqual([11, 21])
   store.set(b, 2)
-  expect(results).toStrictEqual([11, 21, 22]) // store.set(b, 2)
+  expect(results).toStrictEqual([11, 21, 22])
   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(results).toStrictEqual([11, 21, 22, 33])
   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('sets values to atoms without causing infinite loop', () => {
+  const store = createStore()
+  const a = atom(1)
+  a.debugLabel = 'a'
+  const effect = vi.fn((get: Getter, set: Setter) => {
+    set(a, get(a) + 1)
+  })
+  const e = atomSyncEffect(effect)
+  e.debugLabel = 'e'
+  const unsub = store.sub(e, () => {}) // mount syncEffect
+  expect(effect).toBeCalledTimes(1)
+  expect(store.get(a)).toBe(2) // initial values at time of effect mount
+  store.set(a, (v) => ++v)
+  expect(store.get(a)).toBe(4)
+  expect(effect).toBeCalledTimes(2)
+  unsub()
+  expect(effect).toBeCalledTimes(2)
+})
+
+it('reads the value with peak without subscribing to updates', () => {
+  const store = createStore()
+  const a = atom(1)
+  a.debugLabel = 'a'
+  let result = 0
+  const effect = vi.fn((get: GetterWithPeak) => {
+    result = get.peak(a)
+  })
+  const e = atomSyncEffect(effect)
+  e.debugLabel = 'e'
+  store.sub(e, () => {}) // mount syncEffect
+  expect(effect).toBeCalledTimes(1)
+  expect(result).toBe(1) // initial values at time of effect mount
+  store.set(a, 2)
+  expect(effect).toBeCalledTimes(1)
+})
+
+it('supports recursion', () => {
+  const store = createStore()
+  const a = atom(1)
+  a.debugLabel = 'a'
+  const effect = vi.fn((get: Getter, set: SetterWithRecurse) => {
+    if (get(a) < 3) {
+      set.recurse(a, (v) => ++v)
+    }
+  })
+  const e = atomSyncEffect(effect)
+  e.debugLabel = 'e'
+  const unsub = store.sub(e, () => {}) // mount syncEffect
+  expect(effect).toBeCalledTimes(3)
+  expect(store.get(a)).toBe(3)
+})