fix: unwrap synchronously returns awaited values of resolved promise …

…atoms

src/vanilla/store.ts

@@ -102,6 +102,8 @@ type AtomState<Value = AnyValue> = {
   m?: Mounted // only available if the atom is mounted
   /** Atom value */
   v?: Value
+  /** Awaited atom value if the promise value has been resolved */
+  w?: Awaited<Value>
   /** Atom error */
   e?: AnyError
   /** Indicates that the atom value has been changed */
@@ -306,14 +308,14 @@ const buildStore = (
     const pendingPromise = isPendingPromise(atomState.v) ? atomState.v : null
     if (isPromiseLike(valueOrPromise)) {
       patchPromiseForCancelability(valueOrPromise)
+      valueOrPromise.then((v) => (atomState.w = v))
       for (const a of atomState.d.keys()) {
         addPendingPromiseToDependency(atom, valueOrPromise, getAtomState(a))
       }
-      atomState.v = valueOrPromise
-    } else {
-      atomState.v = valueOrPromise
     }
+    atomState.v = valueOrPromise
     delete atomState.e
+    delete atomState.w
     delete atomState.x
     if (!hasPrevValue || !Object.is(prevValue, atomState.v)) {
       ++atomState.n
@@ -426,6 +428,7 @@ const buildStore = (
       return atomState
     } catch (error) {
       delete atomState.v
+      delete atomState.w
       atomState.e = error
       delete atomState.x
       ++atomState.n

src/vanilla/utils/unwrap.ts

@@ -9,10 +9,27 @@ const memo2 = <T>(create: () => T, dep1: object, dep2: object): T => {
   return getCached(create, cache2, dep2)
 }
 
-const isPromise = (x: unknown): x is Promise<unknown> => x instanceof Promise
+const isPromise = <Value>(x: Value): x is Value & Promise<Awaited<Value>> =>
+  x instanceof Promise
 
 const defaultFallback = () => undefined
 
+// HACK to access atom state
+type AtomState<Value> = { w?: Awaited<Value> }
+type GetAtomState = <Value>(atom: Atom<Value>) => AtomState<Value>
+type GetAtomStateRef = { value: GetAtomState }
+const getAtomStateAtom = atom(() => ({}) as GetAtomStateRef)
+getAtomStateAtom.unstable_onInit = (store) => {
+  store.unstable_derive((...args) => {
+    const ref = store.get(getAtomStateAtom)
+    ref.value = args[0]
+    return args
+  })
+}
+if (import.meta.env?.MODE !== 'production') {
+  getAtomStateAtom.debugPrivate = true
+}
+
 export function unwrap<Value, Args extends unknown[], Result>(
   anAtom: WritableAtom<Value, Args, Result>,
 ): WritableAtom<Awaited<Value> | undefined, Args, Result>
@@ -31,85 +48,86 @@ export function unwrap<Value, PendingValue>(
   fallback: (prev?: Awaited<Value>) => PendingValue,
 ): Atom<Awaited<Value> | PendingValue>
 
-export function unwrap<Value, Args extends unknown[], Result, PendingValue>(
-  anAtom: WritableAtom<Value, Args, Result> | Atom<Value>,
+export function unwrap<Value, T extends Atom<Value>, PendingValue>(
+  targetAtom: T,
   fallback: (prev?: Awaited<Value>) => PendingValue = defaultFallback as never,
 ) {
   return memo2(
     () => {
-      type PromiseAndValue = { readonly p?: Promise<unknown> } & (
-        | { readonly v: Awaited<Value> }
-        | { readonly f: PendingValue; readonly v?: Awaited<Value> }
-      )
       const promiseErrorCache = new WeakMap<Promise<unknown>, unknown>()
       const promiseResultCache = new WeakMap<Promise<unknown>, Awaited<Value>>()
       const refreshAtom = atom(0)
-
-      if (import.meta.env?.MODE !== 'production') {
-        refreshAtom.debugPrivate = true
-      }
-
-      const promiseAndValueAtom: WritableAtom<PromiseAndValue, [], void> & {
-        init?: undefined
-      } = atom(
+      type PrevRef = { value: Awaited<Value> | undefined }
+      const prevAtom = atom(() => ({}) as PrevRef)
+      const valueOrPromiseAtom = atom(
         (get, { setSelf }) => {
           get(refreshAtom)
-          const prev = get(promiseAndValueAtom) as PromiseAndValue | undefined
-          const promise = get(anAtom)
-          if (!isPromise(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()
-              },
-            )
+          const prev = get(prevAtom)
+          const atomState = get(getAtomStateAtom).value(targetAtom)
+          const valueOrPromise = get(targetAtom)
+          if (!isPromise(valueOrPromise)) {
+            return { v: valueOrPromise as Awaited<Value> }
           }
-          if (promiseErrorCache.has(promise)) {
-            throw promiseErrorCache.get(promise)
-          }
-          if (promiseResultCache.has(promise)) {
-            return {
-              p: promise,
-              v: promiseResultCache.get(promise) as Awaited<Value>,
+          if (
+            !promiseResultCache.has(valueOrPromise) &&
+            !promiseErrorCache.has(valueOrPromise)
+          ) {
+            if ('w' in atomState) {
+              // resolved
+              promiseResultCache.set(valueOrPromise, atomState.w)
+            } else {
+              // not settled
+              valueOrPromise.then(
+                (v) => {
+                  promiseResultCache.set(valueOrPromise, v)
+                  setSelf()
+                },
+                (e) => {
+                  promiseErrorCache.set(valueOrPromise, e)
+                  setSelf()
+                },
+              )
             }
           }
-          if (prev && 'v' in prev) {
-            return { p: promise, f: fallback(prev.v), v: prev.v }
+          if (promiseErrorCache.has(valueOrPromise)) {
+            throw promiseErrorCache.get(valueOrPromise)
           }
-          return { p: promise, f: fallback() }
-        },
-        (_get, set) => {
-          set(refreshAtom, (c) => c + 1)
+          if (!promiseResultCache.has(valueOrPromise)) {
+            // not resolved
+            return { f: fallback(prev.value) }
+          }
+          return { v: promiseResultCache.get(valueOrPromise)! }
         },
+        (_, set) => set(refreshAtom, (v) => v + 1),
       )
-      // HACK to read PromiseAndValue atom before initialization
-      promiseAndValueAtom.init = undefined
 
       if (import.meta.env?.MODE !== 'production') {
-        promiseAndValueAtom.debugPrivate = true
+        refreshAtom.debugPrivate = true
+        prevAtom.debugPrivate = true
+        valueOrPromiseAtom.debugPrivate = true
       }
 
-      return atom(
-        (get) => {
-          const state = get(promiseAndValueAtom)
-          if ('f' in state) {
-            // is pending
-            return state.f
-          }
-          return state.v
-        },
-        (_get, set, ...args) =>
-          set(anAtom as WritableAtom<Value, unknown[], unknown>, ...args),
+      const descriptors = Object.getOwnPropertyDescriptors(
+        targetAtom as Atom<Awaited<Value> | PendingValue>,
       )
+      descriptors.read.value = (get) => {
+        const state = get(valueOrPromiseAtom)
+        if ('v' in state) {
+          get(prevAtom).value = state.v
+          return state.v
+        }
+        // is pending
+        return state.f
+      }
+      if ('write' in targetAtom) {
+        descriptors.write!.value = (
+          targetAtom as T & WritableAtom<unknown, unknown[], unknown>
+        ).write.bind(targetAtom)
+      }
+      // avoid reading `init` to preserve lazy initialization
+      return Object.create(Object.getPrototypeOf(targetAtom), descriptors)
     },
-    anAtom,
+    targetAtom,
     fallback,
   )
 }

tests/vanilla/utils/unwrap.test.ts

@@ -17,19 +17,19 @@ describe('unwrap', () => {
 
     expect(store.get(syncAtom)).toBe(undefined)
     resolve()
-    await new Promise((r) => setTimeout(r)) // wait for a tick
+    await store.get(asyncAtom)
     expect(store.get(syncAtom)).toBe(2)
 
     store.set(countAtom, 2)
     expect(store.get(syncAtom)).toBe(undefined)
     resolve()
-    await new Promise((r) => setTimeout(r)) // wait for a tick
+    await store.get(asyncAtom)
     expect(store.get(syncAtom)).toBe(4)
 
     store.set(countAtom, 3)
     expect(store.get(syncAtom)).toBe(undefined)
     resolve()
-    await new Promise((r) => setTimeout(r)) // wait for a tick
+    await store.get(asyncAtom)
     expect(store.get(syncAtom)).toBe(6)
   })
 
@@ -45,17 +45,17 @@ describe('unwrap', () => {
     const syncAtom = unwrap(asyncAtom, () => -1)
     expect(store.get(syncAtom)).toBe(-1)
     resolve()
-    await new Promise((r) => setTimeout(r)) // wait for a tick
+    await store.get(asyncAtom)
     expect(store.get(syncAtom)).toBe(2)
     store.set(countAtom, 2)
     expect(store.get(syncAtom)).toBe(-1)
     resolve()
-    await new Promise((r) => setTimeout(r)) // wait for a tick
+    await store.get(asyncAtom)
     expect(store.get(syncAtom)).toBe(4)
     store.set(countAtom, 3)
     expect(store.get(syncAtom)).toBe(-1)
     resolve()
-    await new Promise((r) => setTimeout(r)) // wait for a tick
+    await store.get(asyncAtom)
     expect(store.get(syncAtom)).toBe(6)
   })
 
@@ -72,19 +72,19 @@ describe('unwrap', () => {
 
     expect(store.get(syncAtom)).toBe(0)
     resolve()
-    await new Promise((r) => setTimeout(r)) // wait for a tick
+    await store.get(asyncAtom)
     expect(store.get(syncAtom)).toBe(2)
 
     store.set(countAtom, 2)
     expect(store.get(syncAtom)).toBe(2)
     resolve()
-    await new Promise((r) => setTimeout(r)) // wait for a tick
+    await store.get(asyncAtom)
     expect(store.get(syncAtom)).toBe(4)
 
     store.set(countAtom, 3)
     expect(store.get(syncAtom)).toBe(4)
     resolve()
-    await new Promise((r) => setTimeout(r)) // wait for a tick
+    await store.get(asyncAtom)
     expect(store.get(syncAtom)).toBe(6)
 
     store.set(countAtom, 4)
@@ -93,7 +93,7 @@ describe('unwrap', () => {
     store.set(countAtom, 5)
     expect(store.get(syncAtom)).not.toBe(0) // expect 6 or 8
     resolve()
-    await new Promise((r) => setTimeout(r)) // wait for a tick
+    await store.get(asyncAtom)
     expect(store.get(syncAtom)).toBe(10)
   })
 
@@ -114,39 +114,33 @@ describe('unwrap', () => {
     const syncAtom = unwrap(asyncAtom, (prev?: number) => prev ?? 0)
 
     expect(store.get(syncAtom)).toBe(0)
-    await new Promise((r) => setTimeout(r)) // wait for a tick
+    await store.get(asyncAtom)
     expect(store.get(syncAtom)).toBe(1)
 
     store.set(syncAtom, Promise.resolve(2))
     expect(store.get(syncAtom)).toBe(1)
-    await new Promise((r) => setTimeout(r)) // wait for a tick
+    await store.get(asyncAtom)
     expect(store.get(syncAtom)).toBe(2)
 
     store.set(syncAtom, Promise.resolve(3))
     expect(store.get(syncAtom)).toBe(2)
-    await new Promise((r) => setTimeout(r)) // wait for a tick
+    await store.get(asyncAtom)
     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(undefined)
-    await new Promise((r) => setTimeout(r)) // wait for a tick
+    await store.get(asyncAtom)
     expect(store.get(unwrap(asyncAtom))).toEqual('concrete')
   })
 
   it('should get a fulfilled value after the promise resolves', async () => {
     const store = createStore()
     const asyncAtom = atom(Promise.resolve('concrete'))
     const syncAtom = unwrap(asyncAtom)
-
     expect(store.get(syncAtom)).toEqual(undefined)
-
     await store.get(asyncAtom)
-
     expect(store.get(syncAtom)).toEqual('concrete')
   })
 })