Skip to content

Commit

Permalink
BREAKING: selectAtom does not internally resolve promises
Browse files Browse the repository at this point in the history
  • Loading branch information
@dmaskasky
dmaskasky committed Mar 8, 2024
1 parent 48d89f4 commit 0b69a1c
Show file tree
Hide file tree
Showing 3 changed files with 13 additions and 183 deletions.
44 changes: 9 additions & 35 deletions src/vanilla/utils/selectAtom.ts
View file
Original file line number Diff line number Diff line change
@@ -1,11 +1,8 @@
import { atom } from '../../vanilla.ts'
import type { Atom } from '../../vanilla.ts'

const getCached = <T, K extends object>(
c: () => T,
m: WeakMap<K, T>,
k: K,
): T => (m.has(k) ? m : m.set(k, c())).get(k) as T
const getCached = <T>(c: () => T, m: WeakMap<object, T>, k: object): T =>
(m.has(k) ? m : m.set(k, c())).get(k) as T
const cache1 = new WeakMap()
const memo3 = <T>(
create: () => T,
Expand All @@ -18,32 +15,22 @@ const memo3 = <T>(
return getCached(create, cache3, dep3)
}

const stateCache = new WeakMap()

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 initState = () => ({
version: 0,
prev: EMPTY as Slice | typeof EMPTY | Promise<Slice>,
})
const refFamily = getCached(() => new WeakMap(), stateCache, anAtom)

const identityAtom = atom(() => ({}))

const selectValue = ([value, prevSlice]: readonly [
Awaited<Value>,
Value,
Slice | typeof EMPTY,
]) => {
if (prevSlice === EMPTY) {
Expand All @@ -52,25 +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 ref = getCached(initState, refFamily, get(identityAtom))
const prev = get(derivedAtom)
const prevSlice = prev instanceof Promise ? ref.prev : prev
const value = get(anAtom)
const version = ++ref.version
if (value instanceof Promise || prevSlice instanceof Promise) {
return (ref.prev = Promise.all([value, prevSlice] as const)
.then(selectValue)
.then((slice) => {
if (version === ref.version) {
ref.prev = slice
}
return slice
}))
}
return (ref.prev = selectValue([value as Awaited<Value>, prevSlice]))
return selectValue([value, prev] as const)
})
// HACK to read derived atom before initialization
derivedAtom.init = EMPTY
Expand Down
142 changes: 4 additions & 138 deletions tests/react/vanilla-utils/selectAtom.test.tsx
View file
Original file line number Diff line number Diff line change
@@ -1,8 +1,8 @@
import { StrictMode, Suspense, useEffect, useRef } from 'react'
import { StrictMode, useEffect, useRef } from 'react'
import { fireEvent, render } from '@testing-library/react'
import { expect, it } from 'vitest'
import { useAtom, useAtomValue, useSetAtom } from 'jotai/react'
import { atom, createStore } from 'jotai/vanilla'
import { it } from 'vitest'
import { useAtomValue, useSetAtom } from 'jotai/react'
import { atom } from 'jotai/vanilla'
import { selectAtom } from 'jotai/vanilla/utils'

const useCommitCount = () => {
Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -177,89 +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}')
})

it('should not return async value when the base atom values are synchronous', async () => {
expect.assertions(4)
type Base = { id: number; value: number }
const initialBase = Promise.resolve({ id: 0, value: 0 })
const baseAtom = atom<Base | Promise<Base>>(initialBase)
const idAtom = selectAtom(
baseAtom,
({ id }) => id,
(a, b) => a === b,
)

const isPromiseLike = (x: unknown): x is PromiseLike<unknown> =>
typeof (x as any)?.then === 'function'

const store = createStore()
async function incrementValue() {
const { id, value } = await store.get(baseAtom)
store.set(baseAtom, { id, value: value + 1 })
}

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))
}
10 changes: 0 additions & 10 deletions tests/vanilla/utils/types.test.tsx
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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', () => {
Expand Down

0 comments on commit 0b69a1c

Please sign in to comment.