diff --git a/src/vanilla/store.ts b/src/vanilla/store.ts index 0a7d7e368a..8a8ffe9068 100644 --- a/src/vanilla/store.ts +++ b/src/vanilla/store.ts @@ -102,6 +102,8 @@ type AtomState = { 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 /** 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 diff --git a/src/vanilla/utils/unwrap.ts b/src/vanilla/utils/unwrap.ts index 2504f68d51..ec0b7b7cfa 100644 --- a/src/vanilla/utils/unwrap.ts +++ b/src/vanilla/utils/unwrap.ts @@ -9,10 +9,27 @@ const memo2 = (create: () => T, dep1: object, dep2: object): T => { return getCached(create, cache2, dep2) } -const isPromise = (x: unknown): x is Promise => x instanceof Promise +const isPromise = (x: Value): x is Value & Promise> => + x instanceof Promise const defaultFallback = () => undefined +// HACK to access atom state +type AtomState = { w?: Awaited } +type GetAtomState = (atom: Atom) => AtomState +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( anAtom: WritableAtom, ): WritableAtom | undefined, Args, Result> @@ -31,85 +48,86 @@ export function unwrap( fallback: (prev?: Awaited) => PendingValue, ): Atom | PendingValue> -export function unwrap( - anAtom: WritableAtom | Atom, +export function unwrap, PendingValue>( + targetAtom: T, fallback: (prev?: Awaited) => PendingValue = defaultFallback as never, ) { return memo2( () => { - type PromiseAndValue = { readonly p?: Promise } & ( - | { readonly v: Awaited } - | { readonly f: PendingValue; readonly v?: Awaited } - ) const promiseErrorCache = new WeakMap, unknown>() const promiseResultCache = new WeakMap, Awaited>() const refreshAtom = atom(0) - - if (import.meta.env?.MODE !== 'production') { - refreshAtom.debugPrivate = true - } - - const promiseAndValueAtom: WritableAtom & { - init?: undefined - } = atom( + type PrevRef = { value: Awaited | 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 } - } - if (promise !== prev?.p) { - promise.then( - (v) => { - promiseResultCache.set(promise, v as Awaited) - 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 } } - if (promiseErrorCache.has(promise)) { - throw promiseErrorCache.get(promise) - } - if (promiseResultCache.has(promise)) { - return { - p: promise, - v: promiseResultCache.get(promise) as Awaited, + 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, ...args), + const descriptors = Object.getOwnPropertyDescriptors( + targetAtom as Atom | 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 + ).write.bind(targetAtom) + } + // avoid reading `init` to preserve lazy initialization + return Object.create(Object.getPrototypeOf(targetAtom), descriptors) }, - anAtom, + targetAtom, fallback, ) } diff --git a/tests/vanilla/utils/unwrap.test.ts b/tests/vanilla/utils/unwrap.test.ts index be294dde54..651631dc16 100644 --- a/tests/vanilla/utils/unwrap.test.ts +++ b/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,27 +114,24 @@ 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') }) @@ -142,11 +139,8 @@ describe('unwrap', () => { 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') }) })