[BREAKING] selectAtom does not resolve promises internally (#2435)

* test: selectAtom should not return async value when the base atom curr and prev values are synchronous

* selectAtom replaces promise with fulfilled value when promise resolves

* replace refAtom with weakMap implementation

* partition state by store

* BREAKING: selectAtom does not internally resolve promises

---------

Co-authored-by: David Maskasky <dmaskasky@otter.ai>
Co-authored-by: Daishi Kato <dai-shi@users.noreply.github.com>

src/vanilla/utils/selectAtom.ts

@@ -17,20 +17,20 @@ const memo3 = <T>(
 
 export function selectAtom<Value, Slice>(
   anAtom: Atom<Value>,
-  selector: (v: Awaited<Value>, prevSlice?: Slice) => Slice,
+  selector: (v: Value, prevSlice?: Slice) => Slice,
   equalityFn?: (a: Slice, b: Slice) => boolean,
-): Atom<Value extends Promise<unknown> ? Promise<Slice> : Slice>
+): Atom<Slice>
 
 export function selectAtom<Value, Slice>(
   anAtom: Atom<Value>,
-  selector: (v: Awaited<Value>, prevSlice?: Slice) => Slice,
-  equalityFn: (a: Slice, b: Slice) => boolean = Object.is,
+  selector: (v: Value, prevSlice?: Slice) => Slice,
+  equalityFn: (prevSlice: Slice, slice: Slice) => boolean = Object.is,
 ) {
   return memo3(
     () => {
       const EMPTY = Symbol()
       const selectValue = ([value, prevSlice]: readonly [
-        Awaited<Value>,
+        Value,
         Slice | typeof EMPTY,
       ]) => {
         if (prevSlice === EMPTY) {
@@ -39,15 +39,12 @@ export function selectAtom<Value, Slice>(
         const slice = selector(value, prevSlice)
         return equalityFn(prevSlice, slice) ? prevSlice : slice
       }
-      const derivedAtom: Atom<Slice | Promise<Slice> | typeof EMPTY> & {
+      const derivedAtom: Atom<Slice | typeof EMPTY> & {
         init?: typeof EMPTY
       } = atom((get) => {
         const prev = get(derivedAtom)
         const value = get(anAtom)
-        if (value instanceof Promise || prev instanceof Promise) {
-          return Promise.all([value, prev] as const).then(selectValue)
-        }
-        return selectValue([value as Awaited<Value>, prev] as const)
+        return selectValue([value, prev] as const)
       })
       // HACK to read derived atom before initialization
       derivedAtom.init = EMPTY

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

@@ -1,7 +1,7 @@
-import { StrictMode, Suspense, useEffect, useRef } from 'react'
+import { StrictMode, useEffect, useRef } from 'react'
 import { fireEvent, render } from '@testing-library/react'
 import { it } from 'vitest'
-import { useAtom, useAtomValue, useSetAtom } from 'jotai/react'
+import { useAtomValue, useSetAtom } from 'jotai/react'
 import { atom } from 'jotai/vanilla'
 import { selectAtom } from 'jotai/vanilla/utils'
 
@@ -58,54 +58,6 @@ it('selectAtom works as expected', async () => {
   await findByText('a: 3')
 })
 
-it('selectAtom works with async atom', async () => {
-  const bigAtom = atom({ a: 0, b: 'othervalue' })
-  const bigAtomAsync = atom((get) => Promise.resolve(get(bigAtom)))
-  const littleAtom = selectAtom(bigAtomAsync, (v) => v.a)
-
-  const Parent = () => {
-    const setValue = useSetAtom(bigAtom)
-    return (
-      <>
-        <button
-          onClick={() =>
-            setValue((oldValue) => ({ ...oldValue, a: oldValue.a + 1 }))
-          }
-        >
-          increment
-        </button>
-      </>
-    )
-  }
-
-  const Selector = () => {
-    const a = useAtomValue(littleAtom)
-    return (
-      <>
-        <div>a: {a}</div>
-      </>
-    )
-  }
-
-  const { findByText, getByText } = render(
-    <StrictMode>
-      <Suspense fallback={null}>
-        <Parent />
-        <Selector />
-      </Suspense>
-    </StrictMode>,
-  )
-
-  await findByText('a: 0')
-
-  fireEvent.click(getByText('increment'))
-  await findByText('a: 1')
-  fireEvent.click(getByText('increment'))
-  await findByText('a: 2')
-  fireEvent.click(getByText('increment'))
-  await findByText('a: 3')
-})
-
 it('do not update unless equality function says value has changed', async () => {
   const bigAtom = atom({ a: 0 })
   const littleAtom = selectAtom(
@@ -177,57 +129,3 @@ it('do not update unless equality function says value has changed', async () =>
   await findByText('value: {"a":3}')
   await findByText('commits: 4')
 })
-
-it('equality function works even if suspend', async () => {
-  const bigAtom = atom({ a: 0 })
-  const bigAtomAsync = atom((get) => Promise.resolve(get(bigAtom)))
-  const littleAtom = selectAtom(
-    bigAtomAsync,
-    (value) => value,
-    (left, right) => left.a === right.a,
-  )
-
-  const Controls = () => {
-    const [value, setValue] = useAtom(bigAtom)
-    return (
-      <>
-        <div>bigValue: {JSON.stringify(value)}</div>
-        <button
-          onClick={() =>
-            setValue((oldValue) => ({ ...oldValue, a: oldValue.a + 1 }))
-          }
-        >
-          increment
-        </button>
-        <button onClick={() => setValue((oldValue) => ({ ...oldValue, b: 2 }))}>
-          other
-        </button>
-      </>
-    )
-  }
-
-  const Selector = () => {
-    const value = useAtomValue(littleAtom)
-    return <div>littleValue: {JSON.stringify(value)}</div>
-  }
-
-  const { findByText, getByText } = render(
-    <StrictMode>
-      <Suspense fallback={null}>
-        <Controls />
-        <Selector />
-      </Suspense>
-    </StrictMode>,
-  )
-
-  await findByText('bigValue: {"a":0}')
-  await findByText('littleValue: {"a":0}')
-
-  fireEvent.click(getByText('increment'))
-  await findByText('bigValue: {"a":1}')
-  await findByText('littleValue: {"a":1}')
-
-  fireEvent.click(getByText('other'))
-  await findByText('bigValue: {"a":1,"b":2}')
-  await findByText('littleValue: {"a":1}')
-})

tests/vanilla/utils/types.test.tsx

@@ -10,16 +10,6 @@ it('selectAtom() should return the correct types', () => {
   const syncAtom = atom(0)
   const syncSelectedAtom = selectAtom(syncAtom, doubleCount)
   expectType<TypeEqual<Atom<number>, typeof syncSelectedAtom>>(true)
-
-  const asyncAtom = atom(Promise.resolve(0))
-  const asyncSelectedAtom = selectAtom(asyncAtom, doubleCount)
-  expectType<TypeEqual<Atom<Promise<number>>, typeof asyncSelectedAtom>>(true)
-
-  const maybeAsyncAtom = atom(Promise.resolve(0) as number | Promise<number>)
-  const maybeAsyncSelectedAtom = selectAtom(maybeAsyncAtom, doubleCount)
-  expectType<
-    TypeEqual<Atom<number | Promise<number>>, typeof maybeAsyncSelectedAtom>
-  >(true)
 })
 
 it('unwrap() should return the correct types', () => {