refactor: prefer onInit over setSelf (#3172)

Author: dai-shi

src/vanilla/utils/loadable.ts

@@ -21,51 +21,52 @@ export function loadable<Value>(anAtom: Atom<Value>): Atom<Loadable<Value>> {
       PromiseLike<Awaited<Value>>,
       Loadable<Value>
     >()
-    const refreshAtom = atom(0)
+    const refreshAtom = atom([() => {}, 0] as [() => void, number])
+    refreshAtom.unstable_onInit = (store) => {
+      store.set(refreshAtom, ([, c]) => [
+        () => store.set(refreshAtom, ([f, c]) => [f, c + 1]),
+        c,
+      ])
+    }
 
     if (import.meta.env?.MODE !== 'production') {
       refreshAtom.debugPrivate = true
     }
 
-    const derivedAtom = atom(
-      (get, { setSelf }) => {
-        get(refreshAtom)
-        let value: Value
-        try {
-          value = get(anAtom)
-        } catch (error) {
-          return { state: 'hasError', error } as Loadable<Value>
-        }
-        if (!isPromiseLike<Value>(value)) {
-          return { state: 'hasData', data: value } as Loadable<Value>
-        }
-        const promise = value
-        const cached1 = loadableCache.get(promise)
-        if (cached1) {
-          return cached1
-        }
-        promise.then(
-          (data) => {
-            loadableCache.set(promise, { state: 'hasData', data })
-            setSelf()
-          },
-          (error) => {
-            loadableCache.set(promise, { state: 'hasError', error })
-            setSelf()
-          },
-        )
+    const derivedAtom = atom((get) => {
+      const [triggerRefresh] = get(refreshAtom)
+      let value: Value
+      try {
+        value = get(anAtom)
+      } catch (error) {
+        return { state: 'hasError', error } as Loadable<Value>
+      }
+      if (!isPromiseLike<Value>(value)) {
+        return { state: 'hasData', data: value } as Loadable<Value>
+      }
+      const promise = value
+      const cached1 = loadableCache.get(promise)
+      if (cached1) {
+        return cached1
+      }
+      promise.then(
+        (data) => {
+          loadableCache.set(promise, { state: 'hasData', data })
+          triggerRefresh()
+        },
+        (error) => {
+          loadableCache.set(promise, { state: 'hasError', error })
+          triggerRefresh()
+        },
+      )
 
-        const cached2 = loadableCache.get(promise)
-        if (cached2) {
-          return cached2
-        }
-        loadableCache.set(promise, LOADING as Loadable<Value>)
-        return LOADING as Loadable<Value>
-      },
-      (_get, set) => {
-        set(refreshAtom, (c) => c + 1)
-      },
-    )
+      const cached2 = loadableCache.get(promise)
+      if (cached2) {
+        return cached2
+      }
+      loadableCache.set(promise, LOADING as Loadable<Value>)
+      return LOADING as Loadable<Value>
+    })
 
     if (import.meta.env?.MODE !== 'production') {
       derivedAtom.debugPrivate = true

src/vanilla/utils/unwrap.ts

@@ -47,57 +47,58 @@ export function unwrap<Value, Args extends unknown[], Result, PendingValue>(
         PromiseLike<unknown>,
         Awaited<Value>
       >()
-      const refreshAtom = atom(0)
+      const refreshAtom = atom([() => {}, 0] as [() => void, number])
+      refreshAtom.unstable_onInit = (store) => {
+        store.set(refreshAtom, ([, c]) => [
+          () => store.set(refreshAtom, ([f, c]) => [f, c + 1]),
+          c,
+        ])
+      }
 
       if (import.meta.env?.MODE !== 'production') {
         refreshAtom.debugPrivate = true
       }
 
-      const promiseAndValueAtom: WritableAtom<PromiseAndValue, [], void> & {
+      const promiseAndValueAtom: Atom<PromiseAndValue> & {
         init?: undefined
-      } = atom(
-        (get, { setSelf }) => {
-          get(refreshAtom)
-          let prev: PromiseAndValue | undefined
-          try {
-            prev = get(promiseAndValueAtom) as PromiseAndValue | undefined
-          } catch {
-            // ignore previous errors to avoid getting stuck in error state
-          }
-          const promise = get(anAtom)
-          if (!isPromiseLike(promise)) {
-            return { v: promise as Awaited<Value> }
-          }
-          if (promise !== prev?.p) {
-            promise.then(
-              (v) => {
-                promiseResultCache.set(promise, v as Awaited<Value>)
-                setSelf()
-              },
-              (e) => {
-                promiseErrorCache.set(promise, e)
-                setSelf()
-              },
-            )
-          }
-          if (promiseErrorCache.has(promise)) {
-            throw promiseErrorCache.get(promise)
-          }
-          if (promiseResultCache.has(promise)) {
-            return {
-              p: promise,
-              v: promiseResultCache.get(promise) as Awaited<Value>,
-            }
+      } = atom((get) => {
+        const [triggerRefresh] = get(refreshAtom)
+        let prev: PromiseAndValue | undefined
+        try {
+          prev = get(promiseAndValueAtom) as PromiseAndValue | undefined
+        } catch {
+          // ignore previous errors to avoid getting stuck in error state
+        }
+        const promise = get(anAtom)
+        if (!isPromiseLike(promise)) {
+          return { v: promise as Awaited<Value> }
+        }
+        if (promise !== prev?.p) {
+          promise.then(
+            (v) => {
+              promiseResultCache.set(promise, v as Awaited<Value>)
+              triggerRefresh()
+            },
+            (e) => {
+              promiseErrorCache.set(promise, e)
+              triggerRefresh()
+            },
+          )
+        }
+        if (promiseErrorCache.has(promise)) {
+          throw promiseErrorCache.get(promise)
+        }
+        if (promiseResultCache.has(promise)) {
+          return {
+            p: promise,
+            v: promiseResultCache.get(promise) as Awaited<Value>,
           }
-          if (prev && 'v' in prev) {
-            return { p: promise, f: fallback(prev.v), v: prev.v }
-          }
-          return { p: promise, f: fallback() }
-        },
-        (_get, set) => {
-          set(refreshAtom, (c) => c + 1)
-        },
-      )
+        }
+        if (prev && 'v' in prev) {
+          return { p: promise, f: fallback(prev.v), v: prev.v }
+        }
+        return { p: promise, f: fallback() }
+      })
       // HACK to read PromiseAndValue atom before initialization
       promiseAndValueAtom.init = undefined