Merge branch 'main' into fix/resolved-async-unavailable-with-loadable…

…-or-unwrap

docs/recipes/atom-with-refresh.mdx

@@ -4,59 +4,5 @@ nav: 8.06
 keywords: creators,refresh
 ---
 
-> `atomWithRefresh` creates a derived atom that can be force-refreshed, by using
-> the update function.
-
-This is helpful when you need to refresh asynchronous data after performing a
-side effect.
-
-It can also be used to implement "pull to refresh" functionality.
-
-```ts
-import { atom, Getter } from 'jotai'
-
-export function atomWithRefresh<T>(fn: (get: Getter) => T) {
-  const refreshCounter = atom(0)
-
-  return atom(
-    (get) => {
-      get(refreshCounter)
-      return fn(get)
-    },
-    (_, set) => set(refreshCounter, (i) => i + 1),
-  )
-}
-```
-
-Here's how you'd use it to implement an refresh-able source of data:
-
-```js
-import { atomWithRefresh } from 'XXX'
-
-const postsAtom = atomWithRefresh((get) =>
-  fetch('https://jsonplaceholder.typicode.com/posts').then((r) => r.json()),
-)
-```
-
-In a component:
-
-```jsx
-const PostsList = () => {
-  const [posts, refreshPosts] = useAtom(postsAtom)
-
-  return (
-    <div>
-      <ul>
-        {posts.map((post) => (
-          <li key={post.id}>{post.title}</li>
-        ))}
-      </ul>
-
-      {/* Clicking this button will re-fetch the posts */}
-      <button type="button" onClick={refreshPosts}>
-        Refresh posts
-      </button>
-    </div>
-  )
-}
-```
+`atomWithRefresh` has been provided by `jotai/utils` since v2.7.0.
+[Jump to the doc](../utilities/resettable.mdx#atomWithRefresh)

docs/utilities/resettable.mdx

@@ -146,3 +146,62 @@ is no longer used and any change in dependencies atoms will not trigger an updat
 
 Resetting the value allows us to restore its original default value,
 discarding changes made previously via the `set` function.
+
+## atomWithRefresh
+
+```ts
+function atomWithRefresh<Value>(
+  read: Read<Value, [], void>,
+): WritableAtom<Value, [], void>
+```
+
+Creates an atom that we can refresh,
+which is to force reevaluating the read function.
+
+This is helpful when you need to refresh asynchronous data.
+It can also be used to implement "pull to refresh" functionality.
+
+```ts
+function atomWithRefresh<Value, Args extends unknown[], Result>(
+  read: Read<Value, Args, Result>,
+  write: Write<Value, Args, Result>,
+): WritableAtom<Value, Args | [], Result | void>
+```
+
+Passing zero arguments to `set` will refresh.
+Passing one or more arguments to `set` will call "write" function.
+
+### Example
+
+Here's how you'd use it to implement an refresh-able source of data:
+
+```js
+import { atomWithRefresh } from 'XXX'
+
+const postsAtom = atomWithRefresh((get) =>
+  fetch('https://jsonplaceholder.typicode.com/posts').then((r) => r.json()),
+)
+```
+
+In a component:
+
+```jsx
+const PostsList = () => {
+  const [posts, refreshPosts] = useAtom(postsAtom)
+
+  return (
+    <div>
+      <ul>
+        {posts.map((post) => (
+          <li key={post.id}>{post.title}</li>
+        ))}
+      </ul>
+
+      {/* Clicking this button will re-fetch the posts */}
+      <button type="button" onClick={refreshPosts}>
+        Refresh posts
+      </button>
+    </div>
+  )
+}
+```

src/vanilla/store.ts

@@ -155,9 +155,10 @@ type MountedAtoms = Set<AnyAtom>
 export const createStore = () => {
   const atomStateMap = new WeakMap<AnyAtom, AtomState>()
   const mountedMap = new WeakMap<AnyAtom, Mounted>()
-  const pendingMap = new Map<
+  const pendingStack: Set<AnyAtom>[] = []
+  const pendingMap = new WeakMap<
     AnyAtom,
-    AtomState /* prevAtomState */ | undefined
+    [prevAtomState: AtomState | undefined, dependents: Dependents]
   >()
   let storeListenersRev2: Set<StoreListenerRev2>
   let mountedAtoms: MountedAtoms
@@ -169,6 +170,20 @@ export const createStore = () => {
   const getAtomState = <Value>(atom: Atom<Value>) =>
     atomStateMap.get(atom) as AtomState<Value> | undefined
 
+  const addPendingDependent = (atom: AnyAtom, atomState: AtomState) => {
+    atomState.d.forEach((_, a) => {
+      if (!pendingMap.has(a)) {
+        const aState = getAtomState(a)
+        pendingStack[pendingStack.length - 1]?.add(a)
+        pendingMap.set(a, [aState, new Set()])
+        if (aState) {
+          addPendingDependent(a, aState)
+        }
+      }
+      pendingMap.get(a)![1].add(atom)
+    })
+  }
+
   const setAtomState = <Value>(
     atom: Atom<Value>,
     atomState: AtomState<Value>,
@@ -179,7 +194,9 @@ export const createStore = () => {
     const prevAtomState = getAtomState(atom)
     atomStateMap.set(atom, atomState)
     if (!pendingMap.has(atom)) {
-      pendingMap.set(atom, prevAtomState)
+      pendingStack[pendingStack.length - 1]?.add(atom)
+      pendingMap.set(atom, [prevAtomState, new Set()])
+      addPendingDependent(atom, atomState)
     }
     if (hasPromiseAtomValue(prevAtomState)) {
       const next =
@@ -479,10 +496,8 @@ export const createStore = () => {
   const recomputeDependents = (atom: AnyAtom): void => {
     const getDependents = (a: AnyAtom): Dependents => {
       const dependents = new Set(mountedMap.get(a)?.t)
-      pendingMap.forEach((_, pendingAtom) => {
-        if (getAtomState(pendingAtom)?.d.has(a)) {
-          dependents.add(pendingAtom)
-        }
+      pendingMap.get(a)?.[1].forEach((dependent) => {
+        dependents.add(dependent)
       })
       return dependents
     }
@@ -566,10 +581,10 @@ export const createStore = () => {
         r = writeAtomState(a as AnyWritableAtom, ...args) as R
       }
       if (!isSync) {
-        const flushed = flushPending()
+        const flushed = flushPending([a])
         if (import.meta.env?.MODE !== 'production') {
           storeListenersRev2.forEach((l) =>
-            l({ type: 'async-write', flushed: flushed as Set<AnyAtom> }),
+            l({ type: 'async-write', flushed: flushed! }),
           )
         }
       }
@@ -584,12 +599,11 @@ export const createStore = () => {
     atom: WritableAtom<Value, Args, Result>,
     ...args: Args
   ): Result => {
+    pendingStack.push(new Set([atom]))
     const result = writeAtomState(atom, ...args)
-    const flushed = flushPending()
+    const flushed = flushPending(pendingStack.pop()!)
     if (import.meta.env?.MODE !== 'production') {
-      storeListenersRev2.forEach((l) =>
-        l({ type: 'write', flushed: flushed as Set<AnyAtom> }),
-      )
+      storeListenersRev2.forEach((l) => l({ type: 'write', flushed: flushed! }))
     }
     return result
   }
@@ -709,43 +723,58 @@ export const createStore = () => {
     })
   }
 
-  const flushPending = (): void | Set<AnyAtom> => {
+  const flushPending = (
+    pendingAtoms: AnyAtom[] | Set<AnyAtom>,
+  ): void | Set<AnyAtom> => {
     let flushed: Set<AnyAtom>
     if (import.meta.env?.MODE !== 'production') {
       flushed = new Set()
     }
-    while (pendingMap.size) {
-      const pending = Array.from(pendingMap)
-      pendingMap.clear()
-      pending.forEach(([atom, prevAtomState]) => {
-        const atomState = getAtomState(atom)
-        if (atomState) {
-          const mounted = mountedMap.get(atom)
-          if (mounted && atomState.d !== prevAtomState?.d) {
-            mountDependencies(atom, atomState, prevAtomState?.d)
-          }
-          if (
-            mounted &&
-            !(
-              // TODO This seems pretty hacky. Hope to fix it.
-              // Maybe we could `mountDependencies` in `setAtomState`?
-              (
-                !hasPromiseAtomValue(prevAtomState) &&
-                (isEqualAtomValue(prevAtomState, atomState) ||
-                  isEqualAtomError(prevAtomState, atomState))
-              )
+    const pending: [AnyAtom, AtomState | undefined][] = []
+    const collectPending = (pendingAtom: AnyAtom) => {
+      if (!pendingMap.has(pendingAtom)) {
+        return
+      }
+      const [prevAtomState, dependents] = pendingMap.get(pendingAtom)!
+      pendingMap.delete(pendingAtom)
+      pending.push([pendingAtom, prevAtomState])
+      dependents.forEach(collectPending)
+      // FIXME might be better if we can avoid collecting from dependencies
+      getAtomState(pendingAtom)?.d.forEach((_, a) => collectPending(a))
+    }
+    pendingAtoms.forEach(collectPending)
+    pending.forEach(([atom, prevAtomState]) => {
+      const atomState = getAtomState(atom)
+      if (!atomState) {
+        if (import.meta.env?.MODE !== 'production') {
+          console.warn('[Bug] no atom state to flush')
+        }
+        return
+      }
+      if (atomState !== prevAtomState) {
+        const mounted = mountedMap.get(atom)
+        if (mounted && atomState.d !== prevAtomState?.d) {
+          mountDependencies(atom, atomState, prevAtomState?.d)
+        }
+        if (
+          mounted &&
+          !(
+            // TODO This seems pretty hacky. Hope to fix it.
+            // Maybe we could `mountDependencies` in `setAtomState`?
+            (
+              !hasPromiseAtomValue(prevAtomState) &&
+              (isEqualAtomValue(prevAtomState, atomState) ||
+                isEqualAtomError(prevAtomState, atomState))
             )
-          ) {
-            mounted.l.forEach((listener) => listener())
-            if (import.meta.env?.MODE !== 'production') {
-              flushed.add(atom)
-            }
+          )
+        ) {
+          mounted.l.forEach((listener) => listener())
+          if (import.meta.env?.MODE !== 'production') {
+            flushed.add(atom)
           }
-        } else if (import.meta.env?.MODE !== 'production') {
-          console.warn('[Bug] no atom state to flush')
         }
-      })
-    }
+      }
+    })
     if (import.meta.env?.MODE !== 'production') {
       // @ts-expect-error Variable 'flushed' is used before being assigned.
       return flushed
@@ -754,7 +783,7 @@ export const createStore = () => {
 
   const subscribeAtom = (atom: AnyAtom, listener: () => void) => {
     const mounted = addAtom(atom)
-    const flushed = flushPending()
+    const flushed = flushPending([atom])
     const listeners = mounted.l
     listeners.add(listener)
     if (import.meta.env?.MODE !== 'production') {
@@ -791,15 +820,16 @@ export const createStore = () => {
       dev_get_atom_state: (a: AnyAtom) => atomStateMap.get(a),
       dev_get_mounted: (a: AnyAtom) => mountedMap.get(a),
       dev_restore_atoms: (values: Iterable<readonly [AnyAtom, AnyValue]>) => {
+        pendingStack.push(new Set())
         for (const [atom, valueOrPromise] of values) {
           if (hasInitialValue(atom)) {
             setAtomValueOrPromise(atom, valueOrPromise)
             recomputeDependents(atom)
           }
         }
-        const flushed = flushPending()
+        const flushed = flushPending(pendingStack.pop()!)
         storeListenersRev2.forEach((l) =>
-          l({ type: 'restore', flushed: flushed as Set<AnyAtom> }),
+          l({ type: 'restore', flushed: flushed! }),
         )
       },
     }

src/vanilla/utils.ts

@@ -14,3 +14,4 @@ export {
 export { atomWithObservable } from './utils/atomWithObservable.ts'
 export { loadable } from './utils/loadable.ts'
 export { unwrap } from './utils/unwrap.ts'
+export { atomWithRefresh } from './utils/atomWithRefresh.ts'

src/vanilla/utils/atomWithRefresh.ts

@@ -0,0 +1,45 @@
+import { atom } from '../../vanilla.ts'
+import type { WritableAtom } from '../../vanilla.ts'
+
+type Read<Value, Args extends unknown[], Result> = WritableAtom<
+  Value,
+  Args,
+  Result
+>['read']
+type Write<Value, Args extends unknown[], Result> = WritableAtom<
+  Value,
+  Args,
+  Result
+>['write']
+
+export function atomWithRefresh<Value, Args extends unknown[], Result>(
+  read: Read<Value, Args, Result>,
+  write: Write<Value, Args, Result>,
+): WritableAtom<Value, Args | [], Result | void>
+
+export function atomWithRefresh<Value>(
+  read: Read<Value, [], void>,
+): WritableAtom<Value, [], void>
+
+export function atomWithRefresh<Value, Args extends unknown[], Result>(
+  read: Read<Value, Args, Result>,
+  write?: Write<Value, Args, Result>,
+) {
+  const refreshAtom = atom(0)
+  if (import.meta.env?.MODE !== 'production') {
+    refreshAtom.debugPrivate = true
+  }
+  return atom(
+    (get, options) => {
+      get(refreshAtom)
+      return read(get, options as any)
+    },
+    (get, set, ...args: Args) => {
+      if (args.length === 0) {
+        set(refreshAtom, (c) => c + 1)
+      } else if (write) {
+        return write(get, set, ...args)
+      }
+    },
+  )
+}