From fd3102605d430ec2a33e1d68e597a5ae7e1a17af Mon Sep 17 00:00:00 2001 From: David Maskasky Date: Tue, 24 Sep 2024 16:31:16 -0700 Subject: [PATCH] chore(jotai): add jotai for local reference --- jotai/index.ts | 2 + jotai/mode.ts | 1 + jotai/react.ts | 4 + jotai/react/Provider.ts | 39 ++ jotai/react/useAtom.ts | 56 ++ jotai/react/useAtomValue.ts | 168 +++++ jotai/react/useSetAtom.ts | 42 ++ jotai/react/utils.ts | 4 + jotai/react/utils/useAtomCallback.ts | 17 + jotai/react/utils/useHydrateAtoms.ts | 58 ++ jotai/react/utils/useReducerAtom.ts | 46 ++ jotai/react/utils/useResetAtom.ts | 15 + jotai/types.d.ts | 0 jotai/utils.ts | 2 + jotai/vanilla.ts | 13 + jotai/vanilla/atom.ts | 136 ++++ jotai/vanilla/store.ts | 756 ++++++++++++++++++++++ jotai/vanilla/typeUtils.ts | 21 + jotai/vanilla/utils.ts | 18 + jotai/vanilla/utils/atomFamily.ts | 114 ++++ jotai/vanilla/utils/atomWithDefault.ts | 47 ++ jotai/vanilla/utils/atomWithLazy.ts | 15 + jotai/vanilla/utils/atomWithObservable.ts | 195 ++++++ jotai/vanilla/utils/atomWithReducer.ts | 21 + jotai/vanilla/utils/atomWithRefresh.ts | 46 ++ jotai/vanilla/utils/atomWithReset.ts | 33 + jotai/vanilla/utils/atomWithStorage.ts | 274 ++++++++ jotai/vanilla/utils/constants.ts | 5 + jotai/vanilla/utils/freezeAtom.ts | 65 ++ jotai/vanilla/utils/loadable.ts | 76 +++ jotai/vanilla/utils/selectAtom.ts | 57 ++ jotai/vanilla/utils/splitAtom.ts | 219 +++++++ jotai/vanilla/utils/unwrap.ts | 112 ++++ 33 files changed, 2677 insertions(+) create mode 100644 jotai/index.ts create mode 100644 jotai/mode.ts create mode 100644 jotai/react.ts create mode 100644 jotai/react/Provider.ts create mode 100644 jotai/react/useAtom.ts create mode 100644 jotai/react/useAtomValue.ts create mode 100644 jotai/react/useSetAtom.ts create mode 100644 jotai/react/utils.ts create mode 100644 jotai/react/utils/useAtomCallback.ts create mode 100644 jotai/react/utils/useHydrateAtoms.ts create mode 100644 jotai/react/utils/useReducerAtom.ts create mode 100644 jotai/react/utils/useResetAtom.ts create mode 100644 jotai/types.d.ts create mode 100644 jotai/utils.ts create mode 100644 jotai/vanilla.ts create mode 100644 jotai/vanilla/atom.ts create mode 100644 jotai/vanilla/store.ts create mode 100644 jotai/vanilla/typeUtils.ts create mode 100644 jotai/vanilla/utils.ts create mode 100644 jotai/vanilla/utils/atomFamily.ts create mode 100644 jotai/vanilla/utils/atomWithDefault.ts create mode 100644 jotai/vanilla/utils/atomWithLazy.ts create mode 100644 jotai/vanilla/utils/atomWithObservable.ts create mode 100644 jotai/vanilla/utils/atomWithReducer.ts create mode 100644 jotai/vanilla/utils/atomWithRefresh.ts create mode 100644 jotai/vanilla/utils/atomWithReset.ts create mode 100644 jotai/vanilla/utils/atomWithStorage.ts create mode 100644 jotai/vanilla/utils/constants.ts create mode 100644 jotai/vanilla/utils/freezeAtom.ts create mode 100644 jotai/vanilla/utils/loadable.ts create mode 100644 jotai/vanilla/utils/selectAtom.ts create mode 100644 jotai/vanilla/utils/splitAtom.ts create mode 100644 jotai/vanilla/utils/unwrap.ts diff --git a/jotai/index.ts b/jotai/index.ts new file mode 100644 index 0000000..f021512 --- /dev/null +++ b/jotai/index.ts @@ -0,0 +1,2 @@ +export * from './vanilla' +export * from './react' diff --git a/jotai/mode.ts b/jotai/mode.ts new file mode 100644 index 0000000..9dfdd56 --- /dev/null +++ b/jotai/mode.ts @@ -0,0 +1 @@ +export const MODE: 'development' | 'production' = 'development' diff --git a/jotai/react.ts b/jotai/react.ts new file mode 100644 index 0000000..e9a6c50 --- /dev/null +++ b/jotai/react.ts @@ -0,0 +1,4 @@ +export { Provider, useStore } from './react/Provider' +export { useAtomValue } from './react/useAtomValue' +export { useSetAtom } from './react/useSetAtom' +export { useAtom } from './react/useAtom' diff --git a/jotai/react/Provider.ts b/jotai/react/Provider.ts new file mode 100644 index 0000000..c350d72 --- /dev/null +++ b/jotai/react/Provider.ts @@ -0,0 +1,39 @@ +import { createContext, createElement, useContext, useRef } from 'react' +import type { FunctionComponentElement, ReactNode } from 'react' +import { createStore, getDefaultStore } from '../vanilla' + +type Store = ReturnType + +type StoreContextType = ReturnType> +const StoreContext: StoreContextType = createContext( + undefined, +) + +type Options = { + store?: Store +} + +export const useStore = (options?: Options): Store => { + const store = useContext(StoreContext) + return options?.store || store || getDefaultStore() +} + +export const Provider = ({ + children, + store, +}: { + children?: ReactNode + store?: Store +}): FunctionComponentElement<{ value: Store | undefined }> => { + const storeRef = useRef() + if (!store && !storeRef.current) { + storeRef.current = createStore() + } + return createElement( + StoreContext.Provider, + { + value: store || storeRef.current, + }, + children, + ) +} diff --git a/jotai/react/useAtom.ts b/jotai/react/useAtom.ts new file mode 100644 index 0000000..7cf76b0 --- /dev/null +++ b/jotai/react/useAtom.ts @@ -0,0 +1,56 @@ +import type { + Atom, + ExtractAtomArgs, + ExtractAtomResult, + ExtractAtomValue, + PrimitiveAtom, + SetStateAction, + WritableAtom, +} from '../vanilla' +import { useAtomValue } from './useAtomValue' +import { useSetAtom } from './useSetAtom' + +type SetAtom = (...args: Args) => Result + +type Options = Parameters[1] + +export function useAtom( + atom: WritableAtom, + options?: Options, +): [Awaited, SetAtom] + +export function useAtom( + atom: PrimitiveAtom, + options?: Options, +): [Awaited, SetAtom<[SetStateAction], void>] + +export function useAtom( + atom: Atom, + options?: Options, +): [Awaited, never] + +export function useAtom< + AtomType extends WritableAtom, +>( + atom: AtomType, + options?: Options, +): [ + Awaited>, + SetAtom, ExtractAtomResult>, +] + +export function useAtom>( + atom: AtomType, + options?: Options, +): [Awaited>, never] + +export function useAtom( + atom: Atom | WritableAtom, + options?: Options, +) { + return [ + useAtomValue(atom, options), + // We do wrong type assertion here, which results in throwing an error. + useSetAtom(atom as WritableAtom, options), + ] +} diff --git a/jotai/react/useAtomValue.ts b/jotai/react/useAtomValue.ts new file mode 100644 index 0000000..0011618 --- /dev/null +++ b/jotai/react/useAtomValue.ts @@ -0,0 +1,168 @@ +/// +import ReactExports, { useDebugValue, useEffect, useReducer } from 'react' +import type { ReducerWithoutAction } from 'react' +import type { Atom, ExtractAtomValue } from '../vanilla' +import { useStore } from './Provider' +import { MODE } from '../mode' + +type Store = ReturnType + +const isPromiseLike = (x: unknown): x is PromiseLike => + typeof (x as any)?.then === 'function' + +const attachPromiseMeta = ( + promise: PromiseLike & { + status?: 'pending' | 'fulfilled' | 'rejected' + value?: T + reason?: unknown + }, +) => { + promise.status = 'pending' + promise.then( + (v) => { + promise.status = 'fulfilled' + promise.value = v + }, + (e) => { + promise.status = 'rejected' + promise.reason = e + }, + ) +} + +const use = + ReactExports.use || + (( + promise: PromiseLike & { + status?: 'pending' | 'fulfilled' | 'rejected' + value?: T + reason?: unknown + }, + ): T => { + if (promise.status === 'pending') { + throw promise + } else if (promise.status === 'fulfilled') { + return promise.value as T + } else if (promise.status === 'rejected') { + throw promise.reason + } else { + attachPromiseMeta(promise) + throw promise + } + }) + +const continuablePromiseMap = new WeakMap< + PromiseLike, + Promise +>() + +const createContinuablePromise = (promise: PromiseLike) => { + let continuablePromise = continuablePromiseMap.get(promise) + if (!continuablePromise) { + continuablePromise = new Promise((resolve, reject) => { + let curr = promise + const onFulfilled = (me: PromiseLike) => (v: T) => { + if (curr === me) { + resolve(v) + } + } + const onRejected = (me: PromiseLike) => (e: unknown) => { + if (curr === me) { + reject(e) + } + } + const registerCancelHandler = (p: PromiseLike) => { + if ('onCancel' in p && typeof p.onCancel === 'function') { + p.onCancel((nextValue: PromiseLike | T) => { + if (MODE !== 'production' && nextValue === p) { + throw new Error('[Bug] p is not updated even after cancelation') + } + if (isPromiseLike(nextValue)) { + continuablePromiseMap.set(nextValue, continuablePromise!) + curr = nextValue + nextValue.then(onFulfilled(nextValue), onRejected(nextValue)) + registerCancelHandler(nextValue) + } else { + resolve(nextValue) + } + }) + } + } + promise.then(onFulfilled(promise), onRejected(promise)) + registerCancelHandler(promise) + }) + continuablePromiseMap.set(promise, continuablePromise) + } + return continuablePromise +} + +type Options = Parameters[0] & { + delay?: number +} + +export function useAtomValue( + atom: Atom, + options?: Options, +): Awaited + +export function useAtomValue>( + atom: AtomType, + options?: Options, +): Awaited> + +export function useAtomValue(atom: Atom, options?: Options) { + const store = useStore(options) + + const [[valueFromReducer, storeFromReducer, atomFromReducer], rerender] = + useReducer< + ReducerWithoutAction, + undefined + >( + (prev) => { + const nextValue = store.get(atom) + if ( + Object.is(prev[0], nextValue) && + prev[1] === store && + prev[2] === atom + ) { + return prev + } + return [nextValue, store, atom] + }, + undefined, + () => [store.get(atom), store, atom], + ) + + let value = valueFromReducer + if (storeFromReducer !== store || atomFromReducer !== atom) { + rerender() + value = store.get(atom) + } + + const delay = options?.delay + useEffect(() => { + const unsub = store.sub(atom, () => { + if (typeof delay === 'number') { + const value = store.get(atom) + if (isPromiseLike(value)) { + attachPromiseMeta(createContinuablePromise(value)) + } + // delay rerendering to wait a promise possibly to resolve + setTimeout(rerender, delay) + return + } + rerender() + }) + rerender() + return unsub + }, [store, atom, delay]) + + useDebugValue(value) + // The use of isPromiseLike is to be consistent with `use` type. + // `instanceof Promise` actually works fine in this case. + if (isPromiseLike(value)) { + const promise = createContinuablePromise(value) + return use(promise) + } + return value as Awaited +} diff --git a/jotai/react/useSetAtom.ts b/jotai/react/useSetAtom.ts new file mode 100644 index 0000000..200b5c0 --- /dev/null +++ b/jotai/react/useSetAtom.ts @@ -0,0 +1,42 @@ +import { useCallback } from 'react' +import type { + ExtractAtomArgs, + ExtractAtomResult, + WritableAtom, +} from '../vanilla' +import { useStore } from './Provider' +import { MODE } from '../mode' + +type SetAtom = (...args: Args) => Result +type Options = Parameters[0] + +export function useSetAtom( + atom: WritableAtom, + options?: Options, +): SetAtom + +export function useSetAtom< + AtomType extends WritableAtom, +>( + atom: AtomType, + options?: Options, +): SetAtom, ExtractAtomResult> + +export function useSetAtom( + atom: WritableAtom, + options?: Options, +) { + const store = useStore(options) + const setAtom = useCallback( + (...args: Args) => { + if (MODE !== 'production' && !('write' in atom)) { + // useAtom can pass non writable atom with wrong type assertion, + // so we should check here. + throw new Error('not writable atom') + } + return store.set(atom, ...args) + }, + [store, atom], + ) + return setAtom +} diff --git a/jotai/react/utils.ts b/jotai/react/utils.ts new file mode 100644 index 0000000..9124088 --- /dev/null +++ b/jotai/react/utils.ts @@ -0,0 +1,4 @@ +export { useResetAtom } from './utils/useResetAtom' +export { useReducerAtom } from './utils/useReducerAtom' +export { useAtomCallback } from './utils/useAtomCallback' +export { useHydrateAtoms } from './utils/useHydrateAtoms' diff --git a/jotai/react/utils/useAtomCallback.ts b/jotai/react/utils/useAtomCallback.ts new file mode 100644 index 0000000..3604371 --- /dev/null +++ b/jotai/react/utils/useAtomCallback.ts @@ -0,0 +1,17 @@ +import { useMemo } from 'react' +import { useSetAtom } from '../useSetAtom' +import { atom } from '../../vanilla' +import type { Getter, Setter } from '../../vanilla' + +type Options = Parameters[1] + +export function useAtomCallback( + callback: (get: Getter, set: Setter, ...arg: Args) => Result, + options?: Options, +): (...args: Args) => Result { + const anAtom = useMemo( + () => atom(null, (get, set, ...args: Args) => callback(get, set, ...args)), + [callback], + ) + return useSetAtom(anAtom, options) +} diff --git a/jotai/react/utils/useHydrateAtoms.ts b/jotai/react/utils/useHydrateAtoms.ts new file mode 100644 index 0000000..0bd293b --- /dev/null +++ b/jotai/react/utils/useHydrateAtoms.ts @@ -0,0 +1,58 @@ +import { useStore } from '../Provider' +import type { WritableAtom } from '../../vanilla' + +type Store = ReturnType +type Options = Parameters[0] & { + dangerouslyForceHydrate?: boolean +} +type AnyWritableAtom = WritableAtom + +type InferAtomTuples = { + [K in keyof T]: T[K] extends readonly [infer A, unknown] + ? A extends WritableAtom + ? readonly [A, Args[0]] + : T[K] + : never +} + +// For internal use only +// This can be changed without notice. +export type INTERNAL_InferAtomTuples = InferAtomTuples + +const hydratedMap: WeakMap> = new WeakMap() + +export function useHydrateAtoms< + T extends (readonly [AnyWritableAtom, unknown])[], +>(values: InferAtomTuples, options?: Options): void + +export function useHydrateAtoms>( + values: T, + options?: Options, +): void + +export function useHydrateAtoms< + T extends Iterable, +>(values: InferAtomTuples, options?: Options): void + +export function useHydrateAtoms< + T extends Iterable, +>(values: T, options?: Options) { + const store = useStore(options) + + const hydratedSet = getHydratedSet(store) + for (const [atom, value] of values) { + if (!hydratedSet.has(atom) || options?.dangerouslyForceHydrate) { + hydratedSet.add(atom) + store.set(atom, value as never) + } + } +} + +const getHydratedSet = (store: Store) => { + let hydratedSet = hydratedMap.get(store) + if (!hydratedSet) { + hydratedSet = new WeakSet() + hydratedMap.set(store, hydratedSet) + } + return hydratedSet +} diff --git a/jotai/react/utils/useReducerAtom.ts b/jotai/react/utils/useReducerAtom.ts new file mode 100644 index 0000000..bc30a87 --- /dev/null +++ b/jotai/react/utils/useReducerAtom.ts @@ -0,0 +1,46 @@ +import { useCallback } from 'react' +import { useAtom } from '../useAtom' +import type { PrimitiveAtom } from '../../vanilla' +import { MODE } from '../../mode' + +type Options = Parameters[1] + +/** + * @deprecated please use a recipe instead + * https://github.com/pmndrs/jotai/pull/2467 + */ +export function useReducerAtom( + anAtom: PrimitiveAtom, + reducer: (v: Value, a?: Action) => Value, + options?: Options, +): [Value, (action?: Action) => void] + +/** + * @deprecated please use a recipe instead + * https://github.com/pmndrs/jotai/pull/2467 + */ +export function useReducerAtom( + anAtom: PrimitiveAtom, + reducer: (v: Value, a: Action) => Value, + options?: Options, +): [Value, (action: Action) => void] + +export function useReducerAtom( + anAtom: PrimitiveAtom, + reducer: (v: Value, a: Action) => Value, + options?: Options, +) { + if (MODE !== 'production') { + console.warn( + '[DEPRECATED] useReducerAtom is deprecated and will be removed in the future. Please create your own version using the recipe. https://github.com/pmndrs/jotai/pull/2467', + ) + } + const [state, setState] = useAtom(anAtom, options) + const dispatch = useCallback( + (action: Action) => { + setState((prev) => reducer(prev, action)) + }, + [setState, reducer], + ) + return [state, dispatch] +} diff --git a/jotai/react/utils/useResetAtom.ts b/jotai/react/utils/useResetAtom.ts new file mode 100644 index 0000000..b860f72 --- /dev/null +++ b/jotai/react/utils/useResetAtom.ts @@ -0,0 +1,15 @@ +import { useCallback } from 'react' +import { useSetAtom } from '../useSetAtom' +import { RESET } from '../../vanilla/utils' +import type { WritableAtom } from '../../vanilla' + +type Options = Parameters[1] + +export function useResetAtom( + anAtom: WritableAtom, + options?: Options, +): () => T { + const setAtom = useSetAtom(anAtom, options) + const resetAtom = useCallback(() => setAtom(RESET), [setAtom]) + return resetAtom +} diff --git a/jotai/types.d.ts b/jotai/types.d.ts new file mode 100644 index 0000000..e69de29 diff --git a/jotai/utils.ts b/jotai/utils.ts new file mode 100644 index 0000000..d273f5a --- /dev/null +++ b/jotai/utils.ts @@ -0,0 +1,2 @@ +export * from './vanilla/utils' +export * from './react/utils' diff --git a/jotai/vanilla.ts b/jotai/vanilla.ts new file mode 100644 index 0000000..55e3581 --- /dev/null +++ b/jotai/vanilla.ts @@ -0,0 +1,13 @@ +export { atom } from './vanilla/atom' +export type { Atom, WritableAtom, PrimitiveAtom } from './vanilla/atom' + +export { createStore, getDefaultStore } from './vanilla/store' + +export type { + Getter, + Setter, + ExtractAtomValue, + ExtractAtomArgs, + ExtractAtomResult, + SetStateAction, +} from './vanilla/typeUtils' diff --git a/jotai/vanilla/atom.ts b/jotai/vanilla/atom.ts new file mode 100644 index 0000000..2c5158a --- /dev/null +++ b/jotai/vanilla/atom.ts @@ -0,0 +1,136 @@ +import { MODE } from '../mode' + +type Getter = (atom: Atom) => Value + +type Setter = ( + atom: WritableAtom, + ...args: Args +) => Result + +type SetAtom = ( + ...args: A +) => Result + +/** + * setSelf is for internal use only and subject to change without notice. + */ +type Read = ( + get: Getter, + options: { readonly signal: AbortSignal; readonly setSelf: SetSelf }, +) => Value + +type Write = ( + get: Getter, + set: Setter, + ...args: Args +) => Result + +// This is an internal type and not part of public API. +// Do not depend on it as it can change without notice. +type WithInitialValue = { + init: Value +} + +type OnUnmount = () => void + +type OnMount = < + S extends SetAtom, +>( + setAtom: S, +) => OnUnmount | void + +export interface Atom { + toString: () => string + read: Read + unstable_is?(a: Atom): boolean + debugLabel?: string + /** + * To ONLY be used by Jotai libraries to mark atoms as private. Subject to change. + * @private + */ + debugPrivate?: boolean +} + +export interface WritableAtom + extends Atom { + read: Read> + write: Write + onMount?: OnMount +} + +type SetStateAction = Value | ((prev: Value) => Value) + +export type PrimitiveAtom = WritableAtom< + Value, + [SetStateAction], + void +> + +let keyCount = 0 // global key count for all atoms + +// writable derived atom +export function atom( + read: Read>, + write: Write, +): WritableAtom + +// read-only derived atom +export function atom(read: Read): Atom + +// write-only derived atom +export function atom( + initialValue: Value, + write: Write, +): WritableAtom & WithInitialValue + +// primitive atom without initial value +export function atom(): PrimitiveAtom & + WithInitialValue + +// primitive atom +export function atom( + initialValue: Value, +): PrimitiveAtom & WithInitialValue + +export function atom( + read?: Value | Read>, + write?: Write, +) { + const key = `atom${++keyCount}` + const config = { + toString() { + return MODE !== 'production' && this.debugLabel + ? key + ':' + this.debugLabel + : key + }, + } as WritableAtom & { init?: Value | undefined } + if (typeof read === 'function') { + config.read = read as Read> + } else { + config.init = read + config.read = defaultRead + config.write = defaultWrite as unknown as Write + } + if (write) { + config.write = write + } + return config +} + +function defaultRead(this: Atom, get: Getter) { + return get(this) +} + +function defaultWrite( + this: PrimitiveAtom, + get: Getter, + set: Setter, + arg: SetStateAction, +) { + return set( + this, + typeof arg === 'function' + ? (arg as (prev: Value) => Value)(get(this)) + : arg, + ) +} diff --git a/jotai/vanilla/store.ts b/jotai/vanilla/store.ts new file mode 100644 index 0000000..9e5a8d1 --- /dev/null +++ b/jotai/vanilla/store.ts @@ -0,0 +1,756 @@ +import type { Atom, WritableAtom } from './atom.js' +import { MODE } from '../mode' + +type AnyValue = unknown +type AnyError = unknown +type AnyAtom = Atom +type AnyWritableAtom = WritableAtom +type OnUnmount = () => void +type Getter = Parameters[0] +type Setter = Parameters[1] + +const isSelfAtom = (atom: AnyAtom, a: AnyAtom): boolean => + atom.unstable_is ? atom.unstable_is(a) : a === atom + +const hasInitialValue = >( + atom: T, +): atom is T & (T extends Atom ? { init: Value } : never) => + 'init' in atom + +const isActuallyWritableAtom = (atom: AnyAtom): atom is AnyWritableAtom => + !!(atom as AnyWritableAtom).write + +// +// Cancelable Promise +// + +type CancelHandler = (nextValue: unknown) => void +type PromiseState = [cancelHandlers: Set, settled: boolean] + +const cancelablePromiseMap = new WeakMap, PromiseState>() + +const isPendingPromise = (value: unknown): value is PromiseLike => + isPromiseLike(value) && !cancelablePromiseMap.get(value)?.[1] + +const cancelPromise = (promise: PromiseLike, nextValue: unknown) => { + const promiseState = cancelablePromiseMap.get(promise) + if (promiseState) { + promiseState[1] = true + promiseState[0].forEach((fn) => fn(nextValue)) + } else if (MODE !== 'production') { + throw new Error('[Bug] cancelable promise not found') + } +} + +const patchPromiseForCancelability = (promise: PromiseLike) => { + if (cancelablePromiseMap.has(promise)) { + // already patched + return + } + const promiseState: PromiseState = [new Set(), false] + cancelablePromiseMap.set(promise, promiseState) + const settle = () => { + promiseState[1] = true + } + promise.then(settle, settle) + ;(promise as { onCancel?: (fn: CancelHandler) => void }).onCancel = (fn) => { + promiseState[0].add(fn) + } +} + +const isPromiseLike = ( + x: unknown, +): x is PromiseLike & { onCancel?: (fn: CancelHandler) => void } => + typeof (x as any)?.then === 'function' + +/** + * State tracked for mounted atoms. An atom is considered "mounted" if it has a + * subscriber, or is a transitive dependency of another atom that has a + * subscriber. + * + * The mounted state of an atom is freed once it is no longer mounted. + */ +type Mounted = { + /** Set of listeners to notify when the atom value changes. */ + readonly l: Set<() => void> + /** Set of mounted atoms that the atom depends on. */ + readonly d: Set + /** Set of mounted atoms that depends on the atom. */ + readonly t: Set + /** Function to run when the atom is unmounted. */ + u?: OnUnmount +} + +/** + * Mutable atom state, + * tracked for both mounted and unmounted atoms in a store. + */ +type AtomState = { + /** + * Map of atoms that the atom depends on. + * The map value is the epoch number of the dependency. + */ + readonly d: Map + /** + * Set of atoms with pending promise that depend on the atom. + * + * This may cause memory leaks, but it's for the capability to continue promises + */ + readonly p: Set + /** The epoch number of the atom. */ + n: number + /** Object to store mounted state of the atom. */ + m?: Mounted // only available if the atom is mounted + /** Atom value */ + v?: Value + /** Atom error */ + e?: AnyError +} + +const isAtomStateInitialized = (atomState: AtomState) => + 'v' in atomState || 'e' in atomState + +const returnAtomValue = (atomState: AtomState): Value => { + if ('e' in atomState) { + throw atomState.e + } + if (MODE !== 'production' && !('v' in atomState)) { + throw new Error('[Bug] atom state is not initialized') + } + return atomState.v! +} + +const addPendingPromiseToDependency = ( + atom: AnyAtom, + promise: PromiseLike, + dependencyAtomState: AtomState, +) => { + if (!dependencyAtomState.p.has(atom)) { + dependencyAtomState.p.add(atom) + promise.then( + () => { + dependencyAtomState.p.delete(atom) + }, + () => { + dependencyAtomState.p.delete(atom) + }, + ) + } +} + +const addDependency = ( + pending: Pending | undefined, + atom: Atom, + atomState: AtomState, + a: AnyAtom, + aState: AtomState, +) => { + if (MODE !== 'production' && a === atom) { + throw new Error('[Bug] atom cannot depend on itself') + } + atomState.d.set(a, aState.n) + if (isPendingPromise(atomState.v)) { + addPendingPromiseToDependency(atom, atomState.v, aState) + } + aState.m?.t.add(atom) + if (pending) { + addPendingDependent(pending, a, atom) + } +} + +// +// Pending +// + +type Pending = readonly [ + dependents: Map>, + atomStates: Map, + functions: Set<() => void>, +] + +const createPending = (): Pending => [new Map(), new Map(), new Set()] + +const addPendingAtom = ( + pending: Pending, + atom: AnyAtom, + atomState: AtomState, +) => { + if (!pending[0].has(atom)) { + pending[0].set(atom, new Set()) + } + pending[1].set(atom, atomState) +} + +const addPendingDependent = ( + pending: Pending, + atom: AnyAtom, + dependent: AnyAtom, +) => { + const dependents = pending[0].get(atom) + if (dependents) { + dependents.add(dependent) + } +} + +const getPendingDependents = (pending: Pending, atom: AnyAtom) => + pending[0].get(atom) + +const addPendingFunction = (pending: Pending, fn: () => void) => { + pending[2].add(fn) +} + +const flushPending = (pending: Pending) => { + while (pending[1].size || pending[2].size) { + pending[0].clear() + const atomStates = new Set(pending[1].values()) + pending[1].clear() + const functions = new Set(pending[2]) + pending[2].clear() + atomStates.forEach((atomState) => atomState.m?.l.forEach((l) => l())) + functions.forEach((fn) => fn()) + } +} + +type GetAtomState = ( + atom: Atom, + originAtomState?: AtomState, +) => AtomState + +// internal & unstable type +type StoreArgs = readonly [ + getAtomState: GetAtomState, + atomRead: ( + atom: Atom, + ...params: Parameters['read']> + ) => Value, + atomWrite: ( + atom: WritableAtom, + ...params: Parameters['write']> + ) => Result, +] + +// for debugging purpose only +type DevStoreRev4 = { + dev4_get_internal_weak_map: () => { + get: (atom: AnyAtom) => AtomState | undefined + } + dev4_get_mounted_atoms: () => Set + dev4_restore_atoms: (values: Iterable) => void +} + +type PrdStore = { + get: (atom: Atom) => Value + set: ( + atom: WritableAtom, + ...args: Args + ) => Result + sub: (atom: AnyAtom, listener: () => void) => () => void + unstable_derive: (fn: (...args: StoreArgs) => StoreArgs) => Store +} + +type Store = PrdStore | (PrdStore & DevStoreRev4) + +export type INTERNAL_DevStoreRev4 = DevStoreRev4 +export type INTERNAL_PrdStore = PrdStore + +const buildStore = ( + getAtomState: StoreArgs[0], + atomRead: StoreArgs[1], + atomWrite: StoreArgs[2] +): Store => { + // for debugging purpose only + let debugMountedAtoms: Set + + if (MODE !== 'production') { + debugMountedAtoms = new Set() + } + + const setAtomStateValueOrPromise = ( + atom: AnyAtom, + /** the state of `atom` */ + atomState: AtomState, + valueOrPromise: unknown, + ) => { + const hasPrevValue = 'v' in atomState + const prevValue = atomState.v + const pendingPromise = isPendingPromise(atomState.v) ? atomState.v : null + if (isPromiseLike(valueOrPromise)) { + patchPromiseForCancelability(valueOrPromise) + for (const /** a dependency (`atomState.d`) of `atom` */ a of atomState.d.keys()) { + addPendingPromiseToDependency( + atom, + valueOrPromise, + getAtomState(a, atomState), + ) + } + atomState.v = valueOrPromise + delete atomState.e + } else { + atomState.v = valueOrPromise + delete atomState.e + } + if (!hasPrevValue || !Object.is(prevValue, atomState.v)) { + ++atomState.n + if (pendingPromise) { + cancelPromise(pendingPromise, valueOrPromise) + } + } + } + + const readAtomState = ( + pending: Pending | undefined, + atom: Atom, + /** the state of `atom` */ + atomState: AtomState, + force?: (a: AnyAtom) => boolean, + ): AtomState => { + // See if we can skip recomputing this atom. + if (!force?.(atom) && isAtomStateInitialized(atomState)) { + // If the atom is mounted, we can use the cache. + // because it should have been updated by dependencies. + if (atomState.m) { + return atomState + } + // Otherwise, check if the dependencies have changed. + // If all dependencies haven't changed, we can use the cache. + if ( + Array.from(atomState.d).every( + (keyValue) => { + /** a dependency (`atomState.d`) of `atom` */ + const a = keyValue[0] + const n = keyValue[1] + // Recursively, read the atom state of the dependency, and + // check if the atom epoch number is unchanged + return readAtomState( + pending, + a, + /** original atomState here is the first call to readAtomState */ + getAtomState(a, atomState), + force + ).n === n + }) + ) { + return atomState + } + } + // Compute a new state for this atom. + atomState.d.clear() + let isSync = true + const getter: Getter = (/** a dependency passed to `getter` */ a: Atom) => { + if (isSelfAtom(atom, a)) { + const aState = getAtomState(a, atomState) + if (!isAtomStateInitialized(aState)) { + if (hasInitialValue(a)) { + setAtomStateValueOrPromise(a, aState, a.init) + } else { + // NOTE invalid derived atoms can reach here + throw new Error('no atom init') + } + } + return returnAtomValue(aState) + } + // a !== atom + const aState = readAtomState( + pending, + a, + /** original atomState here is the first call to readAtomState */ + getAtomState(a, atomState), + force, + ) + if (isSync) { + addDependency(pending, atom, atomState, a, aState) + } else { + const pending = createPending() + addDependency(pending, atom, atomState, a, aState) + mountDependencies(pending, atom, atomState) + flushPending(pending) + } + return returnAtomValue(aState) + } + let controller: AbortController | undefined + let setSelf: ((...args: unknown[]) => unknown) | undefined + const options = { + get signal() { + if (!controller) { + controller = new AbortController() + } + return controller.signal + }, + get setSelf() { + if ( + MODE !== 'production' && + !isActuallyWritableAtom(atom) + ) { + console.warn('setSelf function cannot be used with read-only atom') + } + if (!setSelf && isActuallyWritableAtom(atom)) { + setSelf = (...args) => { + if (MODE !== 'production' && isSync) { + console.warn('setSelf function cannot be called in sync') + } + if (!isSync) { + return writeAtom(atom, ...args) + } + } + } + return setSelf + }, + } + try { + const valueOrPromise = atomRead(atom, getter, options as never) + setAtomStateValueOrPromise(atom, atomState, valueOrPromise) + if (isPromiseLike(valueOrPromise)) { + valueOrPromise.onCancel?.(() => controller?.abort()) + const complete = () => { + if (atomState.m) { + const pending = createPending() + mountDependencies(pending, atom, atomState) + flushPending(pending) + } + } + valueOrPromise.then(complete, complete) + } + return atomState + } catch (error) { + delete atomState.v + atomState.e = error + ++atomState.n + return atomState + } finally { + isSync = false + } + } + + const readAtom = ( + /** the atom read by `store.get` */ + atom: Atom + ): Value => + returnAtomValue(readAtomState(undefined, atom, getAtomState(atom))) + + const getDependents = ( + pending: Pending, + atom: Atom, + /** the state of `atom` */ + atomState: AtomState, + ): Map => { + const dependents = new Map() + for (const a of atomState.m?.t || []) { + /** + * original atomState is the atomState of the atom targeted by setter, + * or the of the atom dependent deeply on the atom targeted by setter + * passed from `recomputeDependents` + */ + dependents.set(a, getAtomState(a, atomState)) + } + for (const atomWithPendingPromise of atomState.p) { + dependents.set( + atomWithPendingPromise, + getAtomState(atomWithPendingPromise, atomState), + ) + } + getPendingDependents(pending, atom)?.forEach((dependent) => { + dependents.set(dependent, getAtomState(dependent, atomState)) + }) + return dependents + } + + const recomputeDependents = ( + pending: Pending, + atom: Atom, + atomState: AtomState, + ) => { + // This is a topological sort via depth-first search, slightly modified from + // what's described here for simplicity and performance reasons: + // https://en.wikipedia.org/wiki/Topological_sorting#Depth-first_search + + // Step 1: traverse the dependency graph to build the topsorted atom list + // We don't bother to check for cycles, which simplifies the algorithm. + const topsortedAtoms: (readonly [ + atom: AnyAtom, + atomState: AtomState, + epochNumber: number, + ])[] = [] + const markedAtoms = new Set() + const visit = (a: AnyAtom, aState: AtomState) => { + if (markedAtoms.has(a)) { + return + } + markedAtoms.add(a) + for (const [d, s] of getDependents(pending, a, aState)) { + if (a !== d) { + visit(d, s) + } + } + // The algorithm calls for pushing onto the front of the list. For + // performance, we will simply push onto the end, and then will iterate in + // reverse order later. + topsortedAtoms.push([a, aState, aState.n]) + } + // Visit the root atom. This is the only atom in the dependency graph + // without incoming edges, which is one reason we can simplify the algorithm + visit(atom, atomState) + // Step 2: use the topsorted atom list to recompute all affected atoms + // Track what's changed, so that we can short circuit when possible + const changedAtoms = new Set([atom]) + const isMarked = (a: AnyAtom) => markedAtoms.has(a) + for (let i = topsortedAtoms.length - 1; i >= 0; --i) { + const [a, aState, prevEpochNumber] = topsortedAtoms[i]! + let hasChangedDeps = false + for (const dep of aState.d.keys()) { + if (dep !== a && changedAtoms.has(dep)) { + hasChangedDeps = true + break + } + } + if (hasChangedDeps) { + readAtomState(pending, a, aState, isMarked) + mountDependencies(pending, a, aState) + if (prevEpochNumber !== aState.n) { + addPendingAtom(pending, a, aState) + changedAtoms.add(a) + } + } + markedAtoms.delete(a) + } + } + + const writeAtomState = ( + pending: Pending, + atom: WritableAtom, + /** the state of `atom` */ + atomState: AtomState, + ...args: Args + ): Result => { + const getter: Getter = (a: Atom) => { + /** original atomState is the atomState of the atom targeted by setter */ + const aState = getAtomState(a, atomState) + return returnAtomValue(readAtomState(pending, a, aState)) + } + const setter: Setter = ( + a: WritableAtom, + ...args: As + ) => { + /** original atomState is the atomState of the atom targeted by setter */ + const aState = getAtomState(a, atomState) + let r: R | undefined + if (isSelfAtom(atom, a)) { + if (!hasInitialValue(a)) { + // NOTE technically possible but restricted as it may cause bugs + throw new Error('atom not writable') + } + const hasPrevValue = 'v' in aState + const prevValue = aState.v + const v = args[0] as V + setAtomStateValueOrPromise(a, aState, v) + mountDependencies(pending, a, aState) + if (!hasPrevValue || !Object.is(prevValue, aState.v)) { + addPendingAtom(pending, a, aState) + recomputeDependents(pending, a, aState) + } + } else { + r = writeAtomState(pending, a, aState, ...args) as R + } + flushPending(pending) + return r as R + } + const result = atomWrite(atom, getter, setter, ...args) + return result + } + + const writeAtom = ( + /** the atom read by `store.set` */ + atom: WritableAtom, + ...args: Args + ): Result => { + const pending = createPending() + const result = writeAtomState(pending, atom, getAtomState(atom), ...args) + flushPending(pending) + return result + } + + const mountDependencies = ( + pending: Pending, + atom: AnyAtom, + atomState: AtomState, + ) => { + if (atomState.m && !isPendingPromise(atomState.v)) { + for (const a of atomState.d.keys()) { + if (!atomState.m.d.has(a)) { + const aMounted = mountAtom(pending, a, getAtomState(a, atomState)) + aMounted.t.add(atom) + atomState.m.d.add(a) + } + } + for (const a of atomState.m.d || []) { + if (!atomState.d.has(a)) { + atomState.m.d.delete(a) + const aMounted = unmountAtom(pending, a, getAtomState(a, atomState)) + aMounted?.t.delete(atom) + } + } + } + } + + const mountAtom = ( + pending: Pending, + atom: Atom, + atomState: AtomState, + ): Mounted => { + if (!atomState.m) { + // recompute atom state + readAtomState(pending, atom, atomState) + // mount dependencies first + for (const a of atomState.d.keys()) { + const aMounted = mountAtom(pending, a, getAtomState(a, atomState)) + aMounted.t.add(atom) + } + // mount self + atomState.m = { + l: new Set(), + d: new Set(atomState.d.keys()), + t: new Set(), + } + if (MODE !== 'production') { + debugMountedAtoms.add(atom) + } + if (isActuallyWritableAtom(atom) && atom.onMount) { + const mounted = atomState.m + const { onMount } = atom + addPendingFunction(pending, () => { + const onUnmount = onMount((...args) => + writeAtomState(pending, atom, atomState, ...args), + ) + if (onUnmount) { + mounted.u = onUnmount + } + }) + } + } + return atomState.m + } + + const unmountAtom = ( + pending: Pending, + atom: Atom, + atomState: AtomState, + ): Mounted | undefined => { + if ( + atomState.m && + !atomState.m.l.size && + !Array.from(atomState.m.t).some((a) => + getAtomState(a, atomState).m?.d.has(atom), + ) + ) { + // unmount self + const onUnmount = atomState.m.u + if (onUnmount) { + addPendingFunction(pending, onUnmount) + } + delete atomState.m + if (MODE !== 'production') { + debugMountedAtoms.delete(atom) + } + // unmount dependencies + for (const a of atomState.d.keys()) { + const aMounted = unmountAtom(pending, a, getAtomState(a, atomState)) + aMounted?.t.delete(atom) + } + return undefined + } + return atomState.m + } + + const subscribeAtom = (atom: AnyAtom, listener: () => void) => { + const pending = createPending() + const atomState = getAtomState(atom) + const mounted = mountAtom(pending, atom, atomState) + flushPending(pending) + const listeners = mounted.l + listeners.add(listener) + return () => { + listeners.delete(listener) + const pending = createPending() + unmountAtom(pending, atom, atomState) + flushPending(pending) + } + } + + const unstable_derive = (fn: (...args: StoreArgs) => StoreArgs) => + buildStore(...fn(getAtomState, atomRead, atomWrite)) + + const store: Store = { + get: readAtom, + set: writeAtom, + sub: subscribeAtom, + unstable_derive, + } + if (MODE !== 'production') { + const devStore: DevStoreRev4 = { + // store dev methods (these are tentative and subject to change without notice) + dev4_get_internal_weak_map: () => ({ + get: (atom) => { + const atomState = getAtomState(atom) + if (atomState.n === 0) { + // for backward compatibility + return undefined + } + return atomState + }, + }), + dev4_get_mounted_atoms: () => debugMountedAtoms, + dev4_restore_atoms: (values) => { + const pending = createPending() + for (const [atom, value] of values) { + if (hasInitialValue(atom)) { + const atomState = getAtomState(atom) + const hasPrevValue = 'v' in atomState + const prevValue = atomState.v + setAtomStateValueOrPromise(atom, atomState, value) + mountDependencies(pending, atom, atomState) + if (!hasPrevValue || !Object.is(prevValue, atomState.v)) { + addPendingAtom(pending, atom, atomState) + recomputeDependents(pending, atom, atomState) + } + } + } + flushPending(pending) + }, + } + Object.assign(store, devStore) + } + return store +} + +export const createStore = (): Store => { + const atomStateMap = new WeakMap() + const getAtomState = (atom: Atom) => { + let atomState = atomStateMap.get(atom) as AtomState | undefined + if (!atomState) { + atomState = { d: new Map(), p: new Set(), n: 0 } + atomStateMap.set(atom, atomState) + } + return atomState + } + return buildStore( + getAtomState, + (atom, ...params) => atom.read(...params), + (atom, ...params) => atom.write(...params), + ) +} + +let defaultStore: Store | undefined + +export const getDefaultStore = (): Store => { + if (!defaultStore) { + defaultStore = createStore() + if (MODE !== 'production') { + ;(globalThis as any).__JOTAI_DEFAULT_STORE__ ||= defaultStore + if ((globalThis as any).__JOTAI_DEFAULT_STORE__ !== defaultStore) { + console.warn( + 'Detected multiple Jotai instances. It may cause unexpected behavior with the default store. https://github.com/pmndrs/jotai/discussions/2044', + ) + } + } + } + return defaultStore +} diff --git a/jotai/vanilla/typeUtils.ts b/jotai/vanilla/typeUtils.ts new file mode 100644 index 0000000..2863451 --- /dev/null +++ b/jotai/vanilla/typeUtils.ts @@ -0,0 +1,21 @@ +import type { Atom, PrimitiveAtom, WritableAtom } from './atom' + +export type Getter = Parameters['read']>[0] +export type Setter = Parameters< + WritableAtom['write'] +>[1] + +export type ExtractAtomValue = + AtomType extends Atom ? Value : never + +export type ExtractAtomArgs = + AtomType extends WritableAtom + ? Args + : never + +export type ExtractAtomResult = + AtomType extends WritableAtom + ? Result + : never + +export type SetStateAction = ExtractAtomArgs>[0] diff --git a/jotai/vanilla/utils.ts b/jotai/vanilla/utils.ts new file mode 100644 index 0000000..af71aed --- /dev/null +++ b/jotai/vanilla/utils.ts @@ -0,0 +1,18 @@ +export { RESET } from './utils/constants' +export { atomWithReset } from './utils/atomWithReset' +export { atomWithReducer } from './utils/atomWithReducer' +export { atomFamily } from './utils/atomFamily' +export { selectAtom } from './utils/selectAtom' +export { freezeAtom, freezeAtomCreator } from './utils/freezeAtom' +export { splitAtom } from './utils/splitAtom' +export { atomWithDefault } from './utils/atomWithDefault' +export { + atomWithStorage, + createJSONStorage, + withStorageValidator as unstable_withStorageValidator, +} from './utils/atomWithStorage' +export { atomWithObservable } from './utils/atomWithObservable' +export { loadable } from './utils/loadable' +export { unwrap } from './utils/unwrap' +export { atomWithRefresh } from './utils/atomWithRefresh' +export { atomWithLazy } from './utils/atomWithLazy' diff --git a/jotai/vanilla/utils/atomFamily.ts b/jotai/vanilla/utils/atomFamily.ts new file mode 100644 index 0000000..7a9f9da --- /dev/null +++ b/jotai/vanilla/utils/atomFamily.ts @@ -0,0 +1,114 @@ +import { type Atom } from '../atom' + +/** + * in milliseconds + */ +type CreatedAt = number +type ShouldRemove = (createdAt: CreatedAt, param: Param) => boolean +type Cleanup = () => void +type Callback = (event: { + type: 'CREATE' | 'REMOVE' + param: Param + atom: AtomType +}) => void + +export interface AtomFamily { + (param: Param): AtomType + getParams(): Iterable + remove(param: Param): void + setShouldRemove(shouldRemove: ShouldRemove | null): void + /** + * fires when a atom is created or removed + * This API is for advanced use cases, and can change without notice. + */ + unstable_listen(callback: Callback): Cleanup +} + +export function atomFamily>( + initializeAtom: (param: Param) => AtomType, + areEqual?: (a: Param, b: Param) => boolean, +): AtomFamily + +export function atomFamily>( + initializeAtom: (param: Param) => AtomType, + areEqual?: (a: Param, b: Param) => boolean, +) { + let shouldRemove: ShouldRemove | null = null + const atoms: Map = new Map() + const listeners = new Set>() + const createAtom = (param: Param) => { + let item: [AtomType, CreatedAt] | undefined + if (areEqual === undefined) { + item = atoms.get(param) + } else { + // Custom comparator, iterate over all elements + for (const [key, value] of atoms) { + if (areEqual(key, param)) { + item = value + break + } + } + } + + if (item !== undefined) { + if (shouldRemove?.(item[1], param)) { + createAtom.remove(param) + } else { + return item[0] + } + } + + const newAtom = initializeAtom(param) + atoms.set(param, [newAtom, Date.now()]) + notifyListeners('CREATE', param, newAtom) + return newAtom + } + + function notifyListeners( + type: 'CREATE' | 'REMOVE', + param: Param, + atom: AtomType, + ) { + for (const listener of listeners) { + listener({ type, param, atom }) + } + } + + createAtom.unstable_listen = (callback: Callback) => { + listeners.add(callback) + return () => { + listeners.delete(callback) + } + } + + createAtom.getParams = () => atoms.keys() + + createAtom.remove = (param: Param) => { + if (areEqual === undefined) { + if (!atoms.has(param)) return + const [atom] = atoms.get(param)! + atoms.delete(param) + notifyListeners('REMOVE', param, atom) + } else { + for (const [key, [atom]] of atoms) { + if (areEqual(key, param)) { + atoms.delete(key) + notifyListeners('REMOVE', key, atom) + break + } + } + } + } + + createAtom.setShouldRemove = (fn: ShouldRemove | null) => { + shouldRemove = fn + if (!shouldRemove) return + for (const [key, [atom, createdAt]] of atoms) { + if (shouldRemove(createdAt, key)) { + atoms.delete(key) + notifyListeners('REMOVE', key, atom) + } + } + } + return createAtom +} diff --git a/jotai/vanilla/utils/atomWithDefault.ts b/jotai/vanilla/utils/atomWithDefault.ts new file mode 100644 index 0000000..43cf12f --- /dev/null +++ b/jotai/vanilla/utils/atomWithDefault.ts @@ -0,0 +1,47 @@ +import { atom } from '../atom' +import type { WritableAtom } from '../atom' +import type { SetStateAction } from '../typeUtils' +import { RESET } from './constants' +import { MODE } from '../../mode' + +type Read = WritableAtom< + Value, + Args, + Result +>['read'] + +export function atomWithDefault( + getDefault: Read | typeof RESET], void>, +): WritableAtom | typeof RESET], void> { + const EMPTY = Symbol() + const overwrittenAtom = atom(EMPTY) + + if (MODE !== 'production') { + overwrittenAtom.debugPrivate = true + } + + const anAtom: WritableAtom< + Value, + [SetStateAction | typeof RESET], + void + > = atom( + (get, options) => { + const overwritten = get(overwrittenAtom) + if (overwritten !== EMPTY) { + return overwritten + } + return getDefault(get, options) + }, + (get, set, update) => { + if (update === RESET) { + set(overwrittenAtom, EMPTY) + } else if (typeof update === 'function') { + const prevValue = get(anAtom) + set(overwrittenAtom, (update as (prev: Value) => Value)(prevValue)) + } else { + set(overwrittenAtom, update) + } + }, + ) + return anAtom +} diff --git a/jotai/vanilla/utils/atomWithLazy.ts b/jotai/vanilla/utils/atomWithLazy.ts new file mode 100644 index 0000000..50d219d --- /dev/null +++ b/jotai/vanilla/utils/atomWithLazy.ts @@ -0,0 +1,15 @@ +import { atom } from '../atom' +import type { PrimitiveAtom } from '../atom' + +export function atomWithLazy( + makeInitial: () => Value, +): PrimitiveAtom { + const a = atom(undefined as unknown as Value) + delete (a as { init?: Value }).init + Object.defineProperty(a, 'init', { + get() { + return makeInitial() + }, + }) + return a +} diff --git a/jotai/vanilla/utils/atomWithObservable.ts b/jotai/vanilla/utils/atomWithObservable.ts new file mode 100644 index 0000000..bd5bb6a --- /dev/null +++ b/jotai/vanilla/utils/atomWithObservable.ts @@ -0,0 +1,195 @@ +import { atom } from '../atom' +import type { Getter } from '../typeUtils' +import type { Atom, WritableAtom } from '../atom' +import { MODE } from '../../mode' + +type Timeout = ReturnType +type AnyError = unknown + +type Subscription = { + unsubscribe: () => void +} + +type Observer = { + next: (value: T) => void + error: (error: AnyError) => void + complete: () => void +} + +declare global { + interface SymbolConstructor { + readonly observable: symbol + } +} + +type ObservableLike = { + [Symbol.observable]?: () => ObservableLike | undefined +} & ( + | { + subscribe(observer: Observer): Subscription + } + | { + subscribe(observer: Partial>): Subscription + } + | { + subscribe(observer: Partial>): Subscription + // Overload function to make typing happy + subscribe(next: (value: T) => void): Subscription + } +) + +type SubjectLike = ObservableLike & Observer + +type Options = { + initialValue?: Data | (() => Data) + unstable_timeout?: number +} + +type OptionsWithInitialValue = { + initialValue: Data | (() => Data) + unstable_timeout?: number +} + +export function atomWithObservable( + getObservable: (get: Getter) => SubjectLike, + options: OptionsWithInitialValue, +): WritableAtom + +export function atomWithObservable( + getObservable: (get: Getter) => SubjectLike, + options?: Options, +): WritableAtom, [Data], void> + +export function atomWithObservable( + getObservable: (get: Getter) => ObservableLike, + options: OptionsWithInitialValue, +): Atom + +export function atomWithObservable( + getObservable: (get: Getter) => ObservableLike, + options?: Options, +): Atom> + +export function atomWithObservable( + getObservable: (get: Getter) => ObservableLike | SubjectLike, + options?: Options, +) { + type Result = { d: Data } | { e: AnyError } + const returnResultData = (result: Result) => { + if ('e' in result) { + throw result.e + } + return result.d + } + + const observableResultAtom = atom((get) => { + let observable = getObservable(get) + const itself = observable[Symbol.observable]?.() + if (itself) { + observable = itself + } + + let resolve: ((result: Result) => void) | undefined + const makePending = () => + new Promise((r) => { + resolve = r + }) + const initialResult: Result | Promise = + options && 'initialValue' in options + ? { + d: + typeof options.initialValue === 'function' + ? (options.initialValue as () => Data)() + : (options.initialValue as Data), + } + : makePending() + + let setResult: ((result: Result) => void) | undefined + let lastResult: Result | undefined + const listener = (result: Result) => { + lastResult = result + resolve?.(result) + setResult?.(result) + } + + let subscription: Subscription | undefined + let timer: Timeout | undefined + const isNotMounted = () => !setResult + const start = () => { + if (subscription) { + clearTimeout(timer) + subscription.unsubscribe() + } + subscription = observable.subscribe({ + next: (d) => listener({ d }), + error: (e) => listener({ e }), + complete: () => {}, + }) + if (isNotMounted() && options?.unstable_timeout) { + timer = setTimeout(() => { + if (subscription) { + subscription.unsubscribe() + subscription = undefined + } + }, options.unstable_timeout) + } + } + start() + + const resultAtom = atom(lastResult || initialResult) + + if (MODE !== 'production') { + resultAtom.debugPrivate = true + } + + resultAtom.onMount = (update) => { + setResult = update + if (lastResult) { + update(lastResult) + } + if (subscription) { + clearTimeout(timer) + } else { + start() + } + return () => { + setResult = undefined + if (subscription) { + subscription.unsubscribe() + subscription = undefined + } + } + } + return [resultAtom, observable, makePending, start, isNotMounted] as const + }) + + if (MODE !== 'production') { + observableResultAtom.debugPrivate = true + } + + const observableAtom = atom( + (get) => { + const [resultAtom] = get(observableResultAtom) + const result = get(resultAtom) + if (result instanceof Promise) { + return result.then(returnResultData) + } + return returnResultData(result) + }, + (get, set, data: Data) => { + const [resultAtom, observable, makePending, start, isNotMounted] = + get(observableResultAtom) + if ('next' in observable) { + if (isNotMounted()) { + set(resultAtom, makePending()) + start() + } + observable.next(data) + } else { + throw new Error('observable is not subject') + } + }, + ) + + return observableAtom +} diff --git a/jotai/vanilla/utils/atomWithReducer.ts b/jotai/vanilla/utils/atomWithReducer.ts new file mode 100644 index 0000000..07fa7bb --- /dev/null +++ b/jotai/vanilla/utils/atomWithReducer.ts @@ -0,0 +1,21 @@ +import { atom } from '../atom' +import type { WritableAtom } from '../atom' + +export function atomWithReducer( + initialValue: Value, + reducer: (value: Value, action?: Action) => Value, +): WritableAtom + +export function atomWithReducer( + initialValue: Value, + reducer: (value: Value, action: Action) => Value, +): WritableAtom + +export function atomWithReducer( + initialValue: Value, + reducer: (value: Value, action: Action) => Value, +) { + return atom(initialValue, function (this: never, get, set, action: Action) { + set(this, reducer(get(this), action)) + }) +} diff --git a/jotai/vanilla/utils/atomWithRefresh.ts b/jotai/vanilla/utils/atomWithRefresh.ts new file mode 100644 index 0000000..d30cf86 --- /dev/null +++ b/jotai/vanilla/utils/atomWithRefresh.ts @@ -0,0 +1,46 @@ +import { atom } from '../atom' +import type { WritableAtom } from '../atom' +import { MODE } from '../../mode' + +type Read = WritableAtom< + Value, + Args, + Result +>['read'] +type Write = WritableAtom< + Value, + Args, + Result +>['write'] + +export function atomWithRefresh( + read: Read, + write: Write, +): WritableAtom + +export function atomWithRefresh( + read: Read, +): WritableAtom + +export function atomWithRefresh( + read: Read, + write?: Write, +) { + const refreshAtom = atom(0) + if (MODE !== 'production') { + refreshAtom.debugPrivate = true + } + return atom( + (get, options) => { + get(refreshAtom) + return read(get, options as never) + }, + (get, set, ...args: Args) => { + if (args.length === 0) { + set(refreshAtom, (c) => c + 1) + } else if (write) { + return write(get, set, ...args) + } + }, + ) +} diff --git a/jotai/vanilla/utils/atomWithReset.ts b/jotai/vanilla/utils/atomWithReset.ts new file mode 100644 index 0000000..7d972b0 --- /dev/null +++ b/jotai/vanilla/utils/atomWithReset.ts @@ -0,0 +1,33 @@ +import { atom } from '../atom' +import type { WritableAtom } from '../atom' +import { RESET } from './constants' + +type SetStateActionWithReset = + | Value + | typeof RESET + | ((prev: Value) => Value | typeof RESET) + +// This is an internal type and not part of public API. +// Do not depend on it as it can change without notice. +type WithInitialValue = { + init: Value +} + +export function atomWithReset( + initialValue: Value, +): WritableAtom], void> & + WithInitialValue { + type Update = SetStateActionWithReset + const anAtom = atom( + initialValue, + (get, set, update) => { + const nextValue = + typeof update === 'function' + ? (update as (prev: Value) => Value | typeof RESET)(get(anAtom)) + : update + + set(anAtom, nextValue === RESET ? initialValue : nextValue) + }, + ) + return anAtom as WritableAtom & WithInitialValue +} diff --git a/jotai/vanilla/utils/atomWithStorage.ts b/jotai/vanilla/utils/atomWithStorage.ts new file mode 100644 index 0000000..ca73b2a --- /dev/null +++ b/jotai/vanilla/utils/atomWithStorage.ts @@ -0,0 +1,274 @@ +import { atom } from '../atom' +import type { WritableAtom } from '../atom' +import { RESET } from './constants' +import { MODE } from '../../mode' + +const isPromiseLike = (x: unknown): x is PromiseLike => + typeof (x as any)?.then === 'function' + +type Unsubscribe = () => void + +type Subscribe = ( + key: string, + callback: (value: Value) => void, + initialValue: Value, +) => Unsubscribe + +type StringSubscribe = ( + key: string, + callback: (value: string | null) => void, +) => Unsubscribe + +type SetStateActionWithReset = + | Value + | typeof RESET + | ((prev: Value) => Value | typeof RESET) + +export interface AsyncStorage { + getItem: (key: string, initialValue: Value) => PromiseLike + setItem: (key: string, newValue: Value) => PromiseLike + removeItem: (key: string) => PromiseLike + subscribe?: Subscribe +} + +export interface SyncStorage { + getItem: (key: string, initialValue: Value) => Value + setItem: (key: string, newValue: Value) => void + removeItem: (key: string) => void + subscribe?: Subscribe +} + +export interface AsyncStringStorage { + getItem: (key: string) => PromiseLike + setItem: (key: string, newValue: string) => PromiseLike + removeItem: (key: string) => PromiseLike + subscribe?: StringSubscribe +} + +export interface SyncStringStorage { + getItem: (key: string) => string | null + setItem: (key: string, newValue: string) => void + removeItem: (key: string) => void + subscribe?: StringSubscribe +} + +export function withStorageValidator( + validator: (value: unknown) => value is Value, +): { + (storage: AsyncStorage): AsyncStorage + (storage: SyncStorage): SyncStorage +} + +export function withStorageValidator( + validator: (value: unknown) => value is Value, +) { + return (unknownStorage: AsyncStorage | SyncStorage) => { + const storage = { + ...unknownStorage, + getItem: (key: string, initialValue: Value) => { + const validate = (value: unknown) => { + if (!validator(value)) { + return initialValue + } + return value + } + const value = unknownStorage.getItem(key, initialValue) + if (isPromiseLike(value)) { + return value.then(validate) + } + return validate(value) + }, + } + return storage + } +} + +type JsonStorageOptions = { + reviver?: (key: string, value: unknown) => unknown + replacer?: (key: string, value: unknown) => unknown +} + +export function createJSONStorage(): SyncStorage + +export function createJSONStorage( + getStringStorage: () => AsyncStringStorage, + options?: JsonStorageOptions, +): AsyncStorage + +export function createJSONStorage( + getStringStorage: () => SyncStringStorage, + options?: JsonStorageOptions, +): SyncStorage + +export function createJSONStorage( + getStringStorage: () => + | AsyncStringStorage + | SyncStringStorage + | undefined = () => { + try { + return window.localStorage + } catch (e) { + if (MODE !== 'production') { + if (typeof window !== 'undefined') { + console.warn(e) + } + } + return undefined + } + }, + options?: JsonStorageOptions, +): AsyncStorage | SyncStorage { + let lastStr: string | undefined + let lastValue: Value + + const storage: AsyncStorage | SyncStorage = { + getItem: (key, initialValue) => { + const parse = (str: string | null) => { + str = str || '' + if (lastStr !== str) { + try { + lastValue = JSON.parse(str, options?.reviver) + } catch { + return initialValue + } + lastStr = str + } + return lastValue + } + const str = getStringStorage()?.getItem(key) ?? null + if (isPromiseLike(str)) { + return str.then(parse) as never + } + return parse(str) as never + }, + setItem: (key, newValue) => + getStringStorage()?.setItem( + key, + JSON.stringify(newValue, options?.replacer), + ), + removeItem: (key) => getStringStorage()?.removeItem(key), + } + + const createHandleSubscribe = + (subscriber: StringSubscribe): Subscribe => + (key, callback, initialValue) => + subscriber(key, (v) => { + let newValue: Value + try { + newValue = JSON.parse(v || '') + } catch { + newValue = initialValue + } + callback(newValue) + }) + + let subscriber: StringSubscribe | undefined + try { + subscriber = getStringStorage()?.subscribe + } catch { + // ignore + } + if ( + !subscriber && + typeof window !== 'undefined' && + typeof window.addEventListener === 'function' && + window.Storage + ) { + subscriber = (key, callback) => { + if (!(getStringStorage() instanceof window.Storage)) { + return () => {} + } + const storageEventCallback = (e: StorageEvent) => { + if (e.storageArea === getStringStorage() && e.key === key) { + callback(e.newValue) + } + } + window.addEventListener('storage', storageEventCallback) + return () => { + window.removeEventListener('storage', storageEventCallback) + } + } + } + + if (subscriber) { + storage.subscribe = createHandleSubscribe(subscriber) + } + return storage +} + +const defaultStorage = createJSONStorage() + +export function atomWithStorage( + key: string, + initialValue: Value, + storage: AsyncStorage, + options?: { getOnInit?: boolean }, +): WritableAtom< + Value | Promise, + [SetStateActionWithReset>], + Promise +> + +export function atomWithStorage( + key: string, + initialValue: Value, + storage?: SyncStorage, + options?: { getOnInit?: boolean }, +): WritableAtom], void> + +export function atomWithStorage( + key: string, + initialValue: Value, + storage: + | SyncStorage + | AsyncStorage = defaultStorage as SyncStorage, + options?: { getOnInit?: boolean }, +) { + const getOnInit = options?.getOnInit + const baseAtom = atom( + getOnInit + ? (storage.getItem(key, initialValue) as Value | Promise) + : initialValue, + ) + + if (MODE !== 'production') { + baseAtom.debugPrivate = true + } + + baseAtom.onMount = (setAtom) => { + setAtom(storage.getItem(key, initialValue) as Value | Promise) + let unsub: Unsubscribe | undefined + if (storage.subscribe) { + unsub = storage.subscribe(key, setAtom, initialValue) + } + return unsub + } + + const anAtom = atom( + (get) => get(baseAtom), + (get, set, update: SetStateActionWithReset>) => { + const nextValue = + typeof update === 'function' + ? ( + update as ( + prev: Value | Promise, + ) => Value | Promise | typeof RESET + )(get(baseAtom)) + : update + if (nextValue === RESET) { + set(baseAtom, initialValue) + return storage.removeItem(key) + } + if (nextValue instanceof Promise) { + return nextValue.then((resolvedValue) => { + set(baseAtom, resolvedValue) + return storage.setItem(key, resolvedValue) + }) + } + set(baseAtom, nextValue) + return storage.setItem(key, nextValue) + }, + ) + + return anAtom as never +} diff --git a/jotai/vanilla/utils/constants.ts b/jotai/vanilla/utils/constants.ts new file mode 100644 index 0000000..0b9772f --- /dev/null +++ b/jotai/vanilla/utils/constants.ts @@ -0,0 +1,5 @@ +import { MODE } from '../../mode' + +export const RESET = Symbol( + MODE !== 'production' ? 'RESET' : '', +) diff --git a/jotai/vanilla/utils/freezeAtom.ts b/jotai/vanilla/utils/freezeAtom.ts new file mode 100644 index 0000000..1392215 --- /dev/null +++ b/jotai/vanilla/utils/freezeAtom.ts @@ -0,0 +1,65 @@ +import type { Atom, WritableAtom } from '../atom' +import { MODE } from '../../mode' + +const frozenAtoms = new WeakSet>() + +const deepFreeze = (obj: unknown) => { + if (typeof obj !== 'object' || obj === null) return + Object.freeze(obj) + const propNames = Object.getOwnPropertyNames(obj) + for (const name of propNames) { + const value = (obj as never)[name] + deepFreeze(value) + } + return obj +} + +export function freezeAtom>( + anAtom: AtomType, +): AtomType + +export function freezeAtom( + anAtom: WritableAtom, +): WritableAtom { + if (frozenAtoms.has(anAtom)) { + return anAtom + } + frozenAtoms.add(anAtom) + + const origRead = anAtom.read + anAtom.read = function (get, options) { + return deepFreeze(origRead.call(this, get, options)) + } + if ('write' in anAtom) { + const origWrite = anAtom.write + anAtom.write = function (get, set, ...args) { + return origWrite.call( + this, + get, + (...setArgs) => { + if (setArgs[0] === anAtom) { + setArgs[1] = deepFreeze(setArgs[1]) + } + + return set(...setArgs) + }, + ...args, + ) + } + } + return anAtom +} + +/** + * @deprecated Define it on users end + */ +export function freezeAtomCreator< + CreateAtom extends (...args: unknown[]) => Atom, +>(createAtom: CreateAtom): CreateAtom { + if (MODE !== 'production') { + console.warn( + '[DEPRECATED] freezeAtomCreator is deprecated, define it on users end', + ) + } + return ((...args: unknown[]) => freezeAtom(createAtom(...args))) as never +} diff --git a/jotai/vanilla/utils/loadable.ts b/jotai/vanilla/utils/loadable.ts new file mode 100644 index 0000000..895852b --- /dev/null +++ b/jotai/vanilla/utils/loadable.ts @@ -0,0 +1,76 @@ +import { atom } from '../atom' +import type { Atom } from '../atom' +import { MODE } from '../../mode' + +const cache1 = new WeakMap() +const memo1 = (create: () => T, dep1: object): T => + (cache1.has(dep1) ? cache1 : cache1.set(dep1, create())).get(dep1) + +const isPromise = (x: unknown): x is Promise> => + x instanceof Promise + +export type Loadable = + | { state: 'loading' } + | { state: 'hasError'; error: unknown } + | { state: 'hasData'; data: Awaited } + +const LOADING: Loadable = { state: 'loading' } + +export function loadable(anAtom: Atom): Atom> { + return memo1(() => { + const loadableCache = new WeakMap< + Promise>, + Loadable + >() + const refreshAtom = atom(0) + + if (MODE !== 'production') { + refreshAtom.debugPrivate = true + } + + const derivedAtom = atom( + (get, { setSelf }) => { + get(refreshAtom) + let value: Value + try { + value = get(anAtom) + } catch (error) { + return { state: 'hasError', error } as Loadable + } + if (!isPromise(value)) { + return { state: 'hasData', data: value } as Loadable + } + const promise = value + const cached1 = loadableCache.get(promise) + if (cached1) { + return cached1 + } + promise + .then( + (data) => { + loadableCache.set(promise, { state: 'hasData', data }) + }, + (error) => { + loadableCache.set(promise, { state: 'hasError', error }) + }, + ) + .finally(setSelf) + const cached2 = loadableCache.get(promise) + if (cached2) { + return cached2 + } + loadableCache.set(promise, LOADING as Loadable) + return LOADING as Loadable + }, + (_get, set) => { + set(refreshAtom, (c) => c + 1) + }, + ) + + if (MODE !== 'production') { + derivedAtom.debugPrivate = true + } + + return atom((get) => get(derivedAtom)) + }, anAtom) +} diff --git a/jotai/vanilla/utils/selectAtom.ts b/jotai/vanilla/utils/selectAtom.ts new file mode 100644 index 0000000..923d26b --- /dev/null +++ b/jotai/vanilla/utils/selectAtom.ts @@ -0,0 +1,57 @@ +import { atom } from '../atom' +import type { Atom } from '../atom' + +const getCached = (c: () => T, m: WeakMap, k: object): T => + (m.has(k) ? m : m.set(k, c())).get(k) as T +const cache1 = new WeakMap() +const memo3 = ( + create: () => T, + dep1: object, + dep2: object, + dep3: object, +): T => { + const cache2 = getCached(() => new WeakMap(), cache1, dep1) + const cache3 = getCached(() => new WeakMap(), cache2, dep2) + return getCached(create, cache3, dep3) +} + +export function selectAtom( + anAtom: Atom, + selector: (v: Value, prevSlice?: Slice) => Slice, + equalityFn?: (a: Slice, b: Slice) => boolean, +): Atom + +export function selectAtom( + anAtom: Atom, + selector: (v: Value, prevSlice?: Slice) => Slice, + equalityFn: (prevSlice: Slice, slice: Slice) => boolean = Object.is, +) { + return memo3( + () => { + const EMPTY = Symbol() + const selectValue = ([value, prevSlice]: readonly [ + Value, + Slice | typeof EMPTY, + ]) => { + if (prevSlice === EMPTY) { + return selector(value) + } + const slice = selector(value, prevSlice) + return equalityFn(prevSlice, slice) ? prevSlice : slice + } + const derivedAtom: Atom & { + init?: typeof EMPTY + } = atom((get) => { + const prev = get(derivedAtom) + const value = get(anAtom) + return selectValue([value, prev] as const) + }) + // HACK to read derived atom before initialization + derivedAtom.init = EMPTY + return derivedAtom + }, + anAtom, + selector, + equalityFn, + ) +} diff --git a/jotai/vanilla/utils/splitAtom.ts b/jotai/vanilla/utils/splitAtom.ts new file mode 100644 index 0000000..ac7893e --- /dev/null +++ b/jotai/vanilla/utils/splitAtom.ts @@ -0,0 +1,219 @@ +import { atom } from '../atom' +import type { + Atom, + PrimitiveAtom, + WritableAtom, +} from '../atom' +import type { + Getter, + SetStateAction, + Setter, +} from '../typeUtils' +import { MODE } from '../../mode' + +const getCached = (c: () => T, m: WeakMap, k: object): T => + (m.has(k) ? m : m.set(k, c())).get(k) as T +const cache1 = new WeakMap() +const memo2 = (create: () => T, dep1: object, dep2: object): T => { + const cache2 = getCached(() => new WeakMap(), cache1, dep1) + return getCached(create, cache2, dep2) +} +const cacheKeyForEmptyKeyExtractor = {} + +const isWritable = ( + atom: Atom | WritableAtom, +): atom is WritableAtom => + !!(atom as WritableAtom).write + +const isFunction = (x: T): x is T & ((...args: never[]) => unknown) => + typeof x === 'function' + +type SplitAtomAction = + | { type: 'remove'; atom: PrimitiveAtom } + | { + type: 'insert' + value: Item + before?: PrimitiveAtom + } + | { + type: 'move' + atom: PrimitiveAtom + before?: PrimitiveAtom + } + +export function splitAtom( + arrAtom: WritableAtom, + keyExtractor?: (item: Item) => Key, +): WritableAtom[], [SplitAtomAction], void> + +export function splitAtom( + arrAtom: Atom, + keyExtractor?: (item: Item) => Key, +): Atom[]> + +export function splitAtom( + arrAtom: WritableAtom | Atom, + keyExtractor?: (item: Item) => Key, +) { + return memo2( + () => { + type ItemAtom = PrimitiveAtom | Atom + type Mapping = { + arr: Item[] + atomList: ItemAtom[] + keyList: Key[] + } + const mappingCache = new WeakMap() + const getMapping = (arr: Item[], prev?: Item[]) => { + let mapping = mappingCache.get(arr) + if (mapping) { + return mapping + } + const prevMapping = prev && mappingCache.get(prev) + const atomList: Atom[] = [] + const keyList: Key[] = [] + arr.forEach((item, index) => { + const key = keyExtractor + ? keyExtractor(item) + : (index as unknown as Key) + keyList[index] = key + const cachedAtom = + prevMapping && + prevMapping.atomList[prevMapping.keyList.indexOf(key)] + if (cachedAtom) { + atomList[index] = cachedAtom + return + } + const read = (get: Getter) => { + const prev = get(mappingAtom) as Mapping | undefined + const currArr = get(arrAtom) + const mapping = getMapping(currArr, prev?.arr) + const index = mapping.keyList.indexOf(key) + if (index < 0 || index >= currArr.length) { + // returning a stale value to avoid errors for use cases such as react-spring + const prevItem = arr[getMapping(arr).keyList.indexOf(key)] + if (prevItem) { + return prevItem + } + throw new Error('splitAtom: index out of bounds for read') + } + return currArr[index]! + } + const write = ( + get: Getter, + set: Setter, + update: SetStateAction, + ) => { + const prev = get(mappingAtom) as Mapping | undefined + const arr = get(arrAtom) + const mapping = getMapping(arr, prev?.arr) + const index = mapping.keyList.indexOf(key) + if (index < 0 || index >= arr.length) { + throw new Error('splitAtom: index out of bounds for write') + } + const nextItem = isFunction(update) + ? (update as (prev: Item) => Item)(arr[index]!) + : update + if (!Object.is(arr[index], nextItem)) { + set(arrAtom as WritableAtom, [ + ...arr.slice(0, index), + nextItem, + ...arr.slice(index + 1), + ]) + } + } + atomList[index] = isWritable(arrAtom) ? atom(read, write) : atom(read) + }) + if ( + prevMapping && + prevMapping.keyList.length === keyList.length && + prevMapping.keyList.every((x, i) => x === keyList[i]) + ) { + // not changed + mapping = prevMapping + } else { + mapping = { arr, atomList, keyList } + } + mappingCache.set(arr, mapping) + return mapping + } + const mappingAtom: Atom & { + init?: undefined + } = atom((get) => { + const prev = get(mappingAtom) as Mapping | undefined + const arr = get(arrAtom) + const mapping = getMapping(arr, prev?.arr) + return mapping + }) + + if (MODE !== 'production') { + mappingAtom.debugPrivate = true + } + + // HACK to read mapping atom before initialization + mappingAtom.init = undefined + const splittedAtom = isWritable(arrAtom) + ? atom( + (get) => get(mappingAtom).atomList, + (get, set, action: SplitAtomAction) => { + switch (action.type) { + case 'remove': { + const index = get(splittedAtom).indexOf(action.atom) + if (index >= 0) { + const arr = get(arrAtom) + set(arrAtom as WritableAtom, [ + ...arr.slice(0, index), + ...arr.slice(index + 1), + ]) + } + break + } + case 'insert': { + const index = action.before + ? get(splittedAtom).indexOf(action.before) + : get(splittedAtom).length + if (index >= 0) { + const arr = get(arrAtom) + set(arrAtom as WritableAtom, [ + ...arr.slice(0, index), + action.value, + ...arr.slice(index), + ]) + } + break + } + case 'move': { + const index1 = get(splittedAtom).indexOf(action.atom) + const index2 = action.before + ? get(splittedAtom).indexOf(action.before) + : get(splittedAtom).length + if (index1 >= 0 && index2 >= 0) { + const arr = get(arrAtom) + if (index1 < index2) { + set(arrAtom as WritableAtom, [ + ...arr.slice(0, index1), + ...arr.slice(index1 + 1, index2), + arr[index1]!, + ...arr.slice(index2), + ]) + } else { + set(arrAtom as WritableAtom, [ + ...arr.slice(0, index2), + arr[index1]!, + ...arr.slice(index2, index1), + ...arr.slice(index1 + 1), + ]) + } + } + break + } + } + }, + ) + : atom((get) => get(mappingAtom).atomList) // read-only atom + return splittedAtom + }, + arrAtom, + keyExtractor || cacheKeyForEmptyKeyExtractor, + ) +} diff --git a/jotai/vanilla/utils/unwrap.ts b/jotai/vanilla/utils/unwrap.ts new file mode 100644 index 0000000..3a55766 --- /dev/null +++ b/jotai/vanilla/utils/unwrap.ts @@ -0,0 +1,112 @@ +import { atom } from '../atom' +import type { Atom, WritableAtom } from '../atom' +import { MODE } from '../../mode' + +const getCached = (c: () => T, m: WeakMap, k: object): T => + (m.has(k) ? m : m.set(k, c())).get(k) as T +const cache1 = new WeakMap() +const memo2 = (create: () => T, dep1: object, dep2: object): T => { + const cache2 = getCached(() => new WeakMap(), cache1, dep1) + return getCached(create, cache2, dep2) +} + +const isPromise = (x: unknown): x is Promise => x instanceof Promise + +const defaultFallback = () => undefined + +export function unwrap( + anAtom: WritableAtom, +): WritableAtom | undefined, Args, Result> + +export function unwrap( + anAtom: WritableAtom, + fallback: (prev?: Awaited) => PendingValue, +): WritableAtom | PendingValue, Args, Result> + +export function unwrap( + anAtom: Atom, +): Atom | undefined> + +export function unwrap( + anAtom: Atom, + fallback: (prev?: Awaited) => PendingValue, +): Atom | PendingValue> + +export function unwrap( + anAtom: WritableAtom | Atom, + 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 (MODE !== 'production') { + refreshAtom.debugPrivate = true + } + + const promiseAndValueAtom: WritableAtom & { + init?: undefined + } = 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), + (e) => promiseErrorCache.set(promise, e), + ) + .finally(setSelf) + } + if (promiseErrorCache.has(promise)) { + throw promiseErrorCache.get(promise) + } + if (promiseResultCache.has(promise)) { + return { + p: promise, + v: promiseResultCache.get(promise) as Awaited, + } + } + if (prev && 'v' in prev) { + return { p: promise, f: fallback(prev.v), v: prev.v } + } + return { p: promise, f: fallback() } + }, + (_get, set) => { + set(refreshAtom, (c) => c + 1) + }, + ) + // HACK to read PromiseAndValue atom before initialization + promiseAndValueAtom.init = undefined + + if (MODE !== 'production') { + promiseAndValueAtom.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), + ) + }, + anAtom, + fallback, + ) +}