fix(utils): make 'unwrap' and 'loadable' able to access async atom's resolved value immediately

src/vanilla/utils/loadable.ts

@@ -5,6 +5,15 @@ const cache1 = new WeakMap()
 const memo1 = <T>(create: () => T, dep1: object): T =>
   (cache1.has(dep1) ? cache1 : cache1.set(dep1, create())).get(dep1)
 
+type PromiseMeta<Value> =
+  | { status?: 'pending' }
+  | { status: 'fulfilled'; value: Awaited<Value> }
+  | { status: 'rejected'; reason: unknown }
+
+const isPromise = <Value>(
+  x: unknown,
+): x is Promise<Awaited<Value>> & PromiseMeta<Value> => x instanceof Promise
+
 export type Loadable<Value> =
   | { state: 'loading' }
   | { state: 'hasError'; error: unknown }
@@ -14,7 +23,10 @@ const LOADING: Loadable<unknown> = { state: 'loading' }
 
 export function loadable<Value>(anAtom: Atom<Value>): Atom<Loadable<Value>> {
   return memo1(() => {
-    const loadableCache = new WeakMap<Promise<void>, Loadable<Value>>()
+    const loadableCache = new WeakMap<
+      Promise<Awaited<Value>>,
+      Loadable<Value>
+    >()
     const refreshAtom = atom(0)
 
     if (import.meta.env?.MODE !== 'production') {
@@ -30,25 +42,38 @@ export function loadable<Value>(anAtom: Atom<Value>): Atom<Loadable<Value>> {
         } catch (error) {
           return { state: 'hasError', error } as Loadable<Value>
         }
-        if (!(value instanceof Promise)) {
+        if (!isPromise<Value>(value)) {
           return { state: 'hasData', data: value } as Loadable<Value>
         }
         const promise = value
-        const cached = loadableCache.get(promise)
-        if (cached) {
-          return cached
+        const cached1 = loadableCache.get(promise)
+        if (cached1) {
+          return cached1
+        }
+        if (promise.status === 'fulfilled') {
+          loadableCache.set(promise, { state: 'hasData', data: promise.value })
+        } else if (promise.status === 'rejected') {
+          loadableCache.set(promise, {
+            state: 'hasError',
+            error: promise.reason,
+          })
+        } else {
+          promise
+            .then(
+              (data) => {
+                loadableCache.set(promise, { state: 'hasData', data })
+              },
+              (error) => {
+                loadableCache.set(promise, { state: 'hasError', error })
+              },
+            )
+            .finally(setSelf)
+        }
+        const cached2 = loadableCache.get(promise)
+        if (cached2) {
+          return cached2
         }
         loadableCache.set(promise, LOADING as Loadable<Value>)
-        promise
-          .then(
-            (data) => {
-              loadableCache.set(promise, { state: 'hasData', data })
-            },
-            (error) => {
-              loadableCache.set(promise, { state: 'hasError', error })
-            },
-          )
-          .finally(setSelf)
         return LOADING as Loadable<Value>
       },
       (_get, set) => {

src/vanilla/utils/unwrap.ts

@@ -9,6 +9,14 @@ const memo2 = <T>(create: () => T, dep1: object, dep2: object): T => {
   return getCached(create, cache2, dep2)
 }
 
+type PromiseMeta =
+  | { status?: 'pending' }
+  | { status: 'fulfilled'; value: unknown }
+  | { status: 'rejected'; reason: unknown }
+
+const isPromise = (x: unknown): x is Promise<unknown> & PromiseMeta =>
+  x instanceof Promise
+
 const defaultFallback = () => undefined
 
 export function unwrap<Value, Args extends unknown[], Result>(
@@ -54,27 +62,31 @@ export function unwrap<Value, Args extends unknown[], Result, PendingValue>(
           get(refreshAtom)
           const prev = get(promiseAndValueAtom) as PromiseAndValue | undefined
           const promise = get(anAtom)
-          if (!(promise instanceof Promise)) {
+          if (!isPromise(promise)) {
             return { v: promise as Awaited<Value> }
           }
-          if (promise === prev?.p) {
-            if (promiseErrorCache.has(promise)) {
-              throw promiseErrorCache.get(promise)
-            }
-            if (promiseResultCache.has(promise)) {
-              return {
-                p: promise,
-                v: promiseResultCache.get(promise) as Awaited<Value>,
-              }
+          if (promise !== prev?.p) {
+            if (promise.status === 'fulfilled') {
+              promiseResultCache.set(promise, promise.value as Awaited<Value>)
+            } else if (promise.status === 'rejected') {
+              promiseErrorCache.set(promise, promise.reason)
+            } else {
+              promise
+                .then(
+                  (v) => promiseResultCache.set(promise, v as Awaited<Value>),
+                  (e) => promiseErrorCache.set(promise, e),
+                )
+                .finally(setSelf)
             }
           }
-          if (promise !== prev?.p) {
-            promise
-              .then(
-                (v) => promiseResultCache.set(promise, v as Awaited<Value>),
-                (e) => promiseErrorCache.set(promise, e),
-              )
-              .finally(setSelf)
+          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 }

tests/vanilla/utils/loadable.test.ts

@@ -0,0 +1,16 @@
+import { describe, expect, it } from 'vitest'
+import { atom, createStore } from 'jotai/vanilla'
+import { loadable } from 'jotai/vanilla/utils'
+
+describe('loadable', () => {
+  it('should return fulfilled value of an already resolved async atom', async () => {
+    const store = createStore()
+    const asyncAtom = atom(Promise.resolve('concrete'))
+
+    expect(await store.get(asyncAtom)).toEqual('concrete')
+    expect(store.get(loadable(asyncAtom))).toEqual({
+      state: 'hasData',
+      data: 'concrete',
+    })
+  })
+})

tests/vanilla/utils/unwrap.test.ts

@@ -127,4 +127,12 @@ describe('unwrap', () => {
     await new Promise((r) => setTimeout(r)) // wait for a tick
     expect(store.get(syncAtom)).toBe(3)
   })
+
+  it('should unwrap to a fulfilled value of an already resolved async atom', async () => {
+    const store = createStore()
+    const asyncAtom = atom(Promise.resolve('concrete'))
+
+    expect(await store.get(asyncAtom)).toEqual('concrete')
+    expect(store.get(unwrap(asyncAtom))).toEqual('concrete')
+  })
 })