fix(utils): refactor loadable (#857)

* fix(utils): refactor loadable

* add more test

Co-authored-by: Christian Silver <Pinpickle@users.noreply.github.com>

Author: dai-shi

src/utils/loadable.ts

@@ -11,32 +11,57 @@ type Loadable<Value> =
   | { state: 'hasError'; error: unknown }
   | { state: 'hasData'; data: ResolveType<Value> }
 
+const LOADING: Loadable<unknown> = { state: 'loading' }
+
 export function loadable<Value>(anAtom: Atom<Value>): Atom<Loadable<Value>> {
   return memoizeAtom(() => {
-    // TODO we should revisit this for a better solution than refAtom
-    const refAtom = atom(() => ({} as { prev?: Loadable<Value> }))
+    const loadableAtomCache = new WeakMap<
+      Promise<void>,
+      Atom<Loadable<Value>>
+    >()
 
-    const derivedAtom = atom((get): Loadable<Value> => {
-      const ref = get(refAtom)
-      let curr = ref.prev
+    const catchAtom = atom((get) => {
+      let promise: Promise<void>
       try {
-        const value = get(anAtom) as ResolveType<Value>
-        if (curr?.state !== 'hasData' || !Object.is(curr.data, value)) {
-          curr = { state: 'hasData', data: value }
-        }
+        const data = get(anAtom) as ResolveType<Value>
+        const loadableAtom = atom({ state: 'hasData', data } as Loadable<Value>)
+        return loadableAtom
       } catch (error) {
         if (error instanceof Promise) {
-          if (curr?.state !== 'loading') {
-            curr = { state: 'loading' }
-          }
+          promise = error
         } else {
-          if (curr?.state !== 'hasError' || !Object.is(curr.error, error)) {
-            curr = { state: 'hasError', error }
+          const loadableAtom = atom({
+            state: 'hasError',
+            error,
+          } as Loadable<Value>)
+          return loadableAtom
+        }
+      }
+      const cached = loadableAtomCache.get(promise)
+      if (cached) {
+        return cached
+      }
+      const loadableAtom = atom(
+        LOADING as Loadable<Value>,
+        async (get, set) => {
+          try {
+            const data: Value = await get(anAtom, { unstable_promise: true })
+            set(loadableAtom, { state: 'hasData', data })
+          } catch (error) {
+            set(loadableAtom, { state: 'hasError', error })
           }
         }
+      )
+      loadableAtom.onMount = (init) => {
+        init()
       }
-      ref.prev = curr
-      return curr as Loadable<Value>
+      loadableAtomCache.set(promise, loadableAtom)
+      return loadableAtom
+    })
+
+    const derivedAtom = atom((get) => {
+      const loadableAtom = get(catchAtom)
+      return get(loadableAtom)
     })
 
     return derivedAtom

tests/utils/loadable.test.tsx

@@ -1,3 +1,4 @@
+import { Suspense, useEffect } from 'react'
 import { fireEvent, render } from '@testing-library/react'
 import { Atom, atom } from 'jotai'
 import { loadable, useAtomValue, useUpdateAtom } from 'jotai/utils'
@@ -130,13 +131,76 @@ it('loadable can recover from error', async () => {
   await findByText('Data: 6')
 })
 
+it('loadable immediately resolves sync values', async () => {
+  const syncAtom = atom(5)
+  const effectCallback = jest.fn()
+
+  const { getByText } = render(
+    <Provider>
+      <LoadableComponent effectCallback={effectCallback} asyncAtom={syncAtom} />
+    </Provider>
+  )
+
+  getByText('Data: 5')
+  expect(effectCallback.mock.calls).not.toContain(
+    expect.objectContaining({ state: 'loading' })
+  )
+  expect(effectCallback).toHaveBeenLastCalledWith({ state: 'hasData', data: 5 })
+})
+
+it('loadable can use resolved promises syncronously', async () => {
+  const asyncAtom = atom(Promise.resolve(5))
+  const effectCallback = jest.fn()
+
+  const ResolveAtomComponent = () => {
+    useAtomValue(asyncAtom)
+
+    return <div>Ready</div>
+  }
+
+  const { getByText, findByText, rerender } = render(
+    <Provider>
+      <Suspense fallback={null}>
+        <ResolveAtomComponent />
+      </Suspense>
+    </Provider>
+  )
+
+  await findByText('Ready')
+
+  rerender(
+    <Provider>
+      <LoadableComponent
+        effectCallback={effectCallback}
+        asyncAtom={asyncAtom}
+      />
+    </Provider>
+  )
+  getByText('Data: 5')
+
+  expect(effectCallback.mock.calls).not.toContain(
+    expect.objectContaining({ state: 'loading' })
+  )
+  expect(effectCallback).toHaveBeenLastCalledWith({ state: 'hasData', data: 5 })
+})
+
 interface LoadableComponentProps {
-  asyncAtom: Atom<Promise<number> | Promise<string>>
+  asyncAtom: Atom<Promise<number> | Promise<string> | string | number>
+  effectCallback?: (loadableValue: any) => void
 }
 
-const LoadableComponent = ({ asyncAtom }: LoadableComponentProps) => {
+const LoadableComponent = ({
+  asyncAtom,
+  effectCallback,
+}: LoadableComponentProps) => {
   const value = useAtomValue(loadable(asyncAtom))
 
+  useEffect(() => {
+    if (effectCallback) {
+      effectCallback(value)
+    }
+  }, [value, effectCallback])
+
   if (value.state === 'loading') {
     return <>Loading...</>
   }