selectAtom replaces promise with fulfilled value when promise resolves

src/vanilla/utils/selectAtom.ts

@@ -29,6 +29,11 @@ export function selectAtom<Value, Slice>(
   return memo3(
     () => {
       const EMPTY = Symbol()
+      const refAtom = atom(() => ({
+        version: 0,
+        prev: EMPTY as Slice | typeof EMPTY | Promise<Slice>,
+      }))
+
       const selectValue = ([value, prevSlice]: readonly [
         Awaited<Value>,
         Slice | typeof EMPTY,
@@ -42,12 +47,22 @@ export function selectAtom<Value, Slice>(
       const derivedAtom: Atom<Slice | Promise<Slice> | typeof EMPTY> & {
         init?: typeof EMPTY
       } = atom((get) => {
+        const ref = get(refAtom)
         const prev = get(derivedAtom)
+        const prevSlice = prev instanceof Promise ? ref.prev : prev
         const value = get(anAtom)
-        if (value instanceof Promise || prev instanceof Promise) {
-          return Promise.all([value, prev] as const).then(selectValue)
+        const version = ++ref.version
+        if (value instanceof Promise || prevSlice instanceof Promise) {
+          return (ref.prev = Promise.all([value, prev] as const)
+            .then(selectValue)
+            .then((slice) => {
+              if (version === ref.version) {
+                ref.prev = slice
+              }
+              return slice
+            }))
         }
-        return selectValue([value as Awaited<Value>, prev] as const)
+        return (ref.prev = selectValue([value as Awaited<Value>, prevSlice]))
       })
       // HACK to read derived atom before initialization
       derivedAtom.init = EMPTY

tests/react/vanilla-utils/selectAtom.test.tsx

@@ -232,12 +232,11 @@ it('equality function works even if suspend', async () => {
   await findByText('littleValue: {"a":1}')
 })
 
-it('should not return async value when the base atom values are synchronous', async () => {
+it.only('should not return async value when the base atom values are synchronous', async () => {
   expect.assertions(4)
   type Base = { id: number; value: number }
-  const baseAtom = atom<Base | Promise<Base>>(
-    Promise.resolve({ id: 0, value: 0 }),
-  )
+  const initialBase = Promise.resolve({ id: 0, value: 0 })
+  const baseAtom = atom<Base | Promise<Base>>(initialBase)
   const idAtom = selectAtom(
     baseAtom,
     ({ id }) => id,
@@ -255,7 +254,12 @@ it('should not return async value when the base atom values are synchronous', as
 
   expect(isPromiseLike(store.get(baseAtom))).toBe(true)
   expect(isPromiseLike(store.get(idAtom))).toBe(true)
+  await delay(0)
   await incrementValue()
   expect(isPromiseLike(store.get(baseAtom))).toBe(false)
   expect(isPromiseLike(store.get(idAtom))).toBe(false)
 })
+
+function delay(ms: number) {
+  return new Promise((resolve) => setTimeout(resolve, ms))
+}