diff --git a/src/state/internal/createQueryStore.ts b/src/state/internal/createQueryStore.ts new file mode 100644 index 00000000000..41ad4ad0980 --- /dev/null +++ b/src/state/internal/createQueryStore.ts @@ -0,0 +1,839 @@ +import { StateCreator, StoreApi, UseBoundStore, create } from 'zustand'; +import { subscribeWithSelector } from 'zustand/middleware'; +import { IS_DEV } from '@/env'; +import { RainbowError, logger } from '@/logger'; +import { RainbowPersistConfig, createRainbowStore, omitStoreMethods } from './createRainbowStore'; +import { $, AttachValue, SignalFunction, Unsubscribe, attachValueSubscriptionMap } from './signal'; + +const ENABLE_LOGS = false; + +/** + * A set of constants representing the various stages of a query's remote data fetching process. + */ +export const QueryStatuses = { + Error: 'error', + Idle: 'idle', + Loading: 'loading', + Success: 'success', +} as const; + +/** + * Represents the current status of the query's remote data fetching operation. + * + * Possible values: + * - **`'error'`** : The most recent request encountered an error. + * - **`'idle'`** : No request in progress, no error, no data yet. + * - **`'loading'`** : A request is currently in progress. + * - **`'success'`** : The most recent request has succeeded, and `data` is available. + */ +export type QueryStatus = (typeof QueryStatuses)[keyof typeof QueryStatuses]; + +/** + * Expanded status information for the currently specified query parameters. + */ +export type QueryStatusInfo = { + isError: boolean; + isFetching: boolean; + isIdle: boolean; + isInitialLoading: boolean; + isSuccess: boolean; +}; + +/** + * Defines additional options for a data fetch operation. + */ +interface FetchOptions { + /** + * Overrides the default cache duration for this fetch, in milliseconds. + * If data in the cache is older than this duration, it will be considered expired and + * will be pruned following a successful fetch. + */ + cacheTime?: number; + /** + * Forces a fetch request even if the current data is fresh and not stale. + * If `true`, the fetch operation bypasses existing cached data. + */ + force?: boolean; + /** + * Overrides the default stale duration for this fetch, in milliseconds. + * If the fetch is successful, the subsequently scheduled refetch will occur after + * the specified duration. + */ + staleTime?: number; +} + +/** + * Represents an entry in the query cache, which stores fetched data along with metadata, and error information + * in the event the most recent fetch failed. + */ +interface CacheEntry { + data: TData | null; + errorInfo: { + error: Error; + lastFailedAt: number; + retryCount: number; + } | null; + lastFetchedAt: number; +} + +/** + * A specialized store interface that combines Zustand's store capabilities with remote data fetching support. + * + * In addition to Zustand's store API (such as `getState()` and `subscribe()`), this interface provides: + * - **`enabled`**: A boolean indicating if the store is actively fetching data. + * - **`fetch(params, options)`**: Initiates a data fetch operation. + * - **`getData(params)`**: Returns the cached data, if available, for the current query parameters. + * - **`getStatus()`**: Returns expanded status information for the current query parameters. + * - **`isDataExpired(override?)`**: Checks if the current data has expired based on `cacheTime`. + * - **`isStale(override?)`**: Checks if the current data is stale based on `staleTime`. + * - **`reset()`**: Resets the store to its initial state, clearing data and errors. + */ +export interface QueryStore< + TData, + TParams extends Record, + S extends Omit, keyof PrivateStoreState>, +> extends UseBoundStore> { + /** + * Indicates whether the store should actively fetch data. + * When `false`, the store won't automatically refetch data. + */ + enabled: boolean; + /** + * The current query key, which is a string representation of the current query parameter values. + */ + queryKey: string; + /** + * Initiates a data fetch for the given parameters. If no parameters are provided, the store's + * current parameters are used. + * @param params - Optional parameters to pass to the fetcher function. + * @param options - Optional {@link FetchOptions} to customize the fetch behavior. + * @returns A promise that resolves when the fetch operation completes. + */ + fetch: (params?: TParams, options?: FetchOptions) => Promise; + /** + * Returns the cached data, if available, for the current query params. + * @returns The cached data, or `null` if no data is available. + */ + getData: (params?: TParams) => TData | null; + /** + * Returns expanded status information for the currently specified query parameters. The raw + * status can be obtained by directly reading the `status` property. + * @example + * ```ts + * const isInitialLoad = useQueryStore(state => state.getStatus().isInitialLoad); + * ``` + * @returns An object containing boolean flags for each status. + */ + getStatus: () => QueryStatusInfo; + /** + * Determines if the current data is expired based on whether `cacheTime` has been exceeded. + * @param override - An optional override for the default cache time, in milliseconds. + * @returns `true` if the data is expired, otherwise `false`. + */ + isDataExpired: (override?: number) => boolean; + /** + * Determines if the current data is stale based on whether `staleTime` has been exceeded. + * Stale data may be refreshed automatically in the background. + * @param override - An optional override for the default stale time, in milliseconds. + * @returns `true` if the data is stale, otherwise `false`. + */ + isStale: (override?: number) => boolean; + /** + * Resets the store to its initial state, clearing data, error, and any cached values. + */ + reset: () => void; +} + +/** + * The private state managed by the query store, omitted from the store's public interface. + */ +type PrivateStoreState = { + subscriptionCount: number; +}; + +/** + * The full state structure managed by the query store. This type is generally internal, + * though the state it defines can be accessed via the store's public interface. + */ +type StoreState> = Pick< + QueryStore>, + 'enabled' | 'queryKey' | 'fetch' | 'getData' | 'getStatus' | 'isDataExpired' | 'isStale' | 'reset' +> & { + error: Error | null; + lastFetchedAt: number | null; + queryCache: Record | undefined>; + status: QueryStatus; +}; + +/** + * Configuration options for creating a query-enabled Rainbow store. + */ +export type RainbowQueryStoreConfig, TData, S extends StoreState> = { + /** + * A function responsible for fetching data from a remote source. + * Receives parameters of type `TParams` and returns either a promise or a raw data value of type `TQueryFnData`. + */ + fetcher: (params: TParams) => TQueryFnData | Promise; + /** + * A callback invoked whenever a fetch operation fails. + * Receives the error and the current retry count. + */ + onError?: (error: Error, retryCount: number) => void; + /** + * A callback invoked whenever fresh data is successfully fetched. + * Receives the transformed data and the store's set function, which can optionally be used to update store state. + */ + onFetched?: (data: TData, set: (partial: S | Partial | ((state: S) => S | Partial)) => void) => void; + /** + * A function that overrides the default behavior of setting the fetched data in the store's query cache. + * Receives an object containing the transformed data, the query parameters, the query key, and the store's set function. + * + * When using `setData`, it’s important to note that you are taking full responsibility for managing query data. If your + * query supports variable parameters (and thus multiple query keys) and you want to cache data for each key, you’ll need + * to manually handle storing data based on the provided `params` or `queryKey`. Naturally, you will also bear + * responsibility for pruning this data in the event you do not want it persisted indefinitely. + * + * Automatic refetching per your specified `staleTime` is still managed internally by the store. While no query *data* + * will be cached internally if `setData` is provided, metadata such as the last fetch time for each query key is still + * cached and tracked by the store, unless caching is fully disabled via `disableCache: true`. + */ + setData?: (info: { + data: TData; + params: TParams; + queryKey: string; + set: (partial: S | Partial | ((state: S) => S | Partial)) => void; + }) => void; + /** + * A function to transform the raw fetched data (`TQueryFnData`) into another form (`TData`). + * If not provided, the raw data returned by `fetcher` is used. + */ + transform?: (data: TQueryFnData) => TData; + /** + * The maximum duration, in milliseconds, that fetched data is considered fresh. + * After this time, data is considered expired and will be refetched when requested. + * @default time.days(7) + */ + cacheTime?: number; + /** + * If `true`, the store's caching mechanisms will be fully disabled, meaning that the store will + * always refetch data on every call to `fetch()`, and the fetched data will not be stored unless + * a `setData` function is provided. + * + * Disable caching if you always want fresh data on refetch. + * @default false + */ + disableCache?: boolean; + /** + * When `true`, the store actively fetches and refetches data as needed. + * When `false`, the store will not automatically fetch data until explicitly enabled. + * @default true + */ + enabled?: boolean; + /** + * The maximum number of times to retry a failed fetch operation. + * @default 3 + */ + maxRetries?: number; + /** + * Parameters to be passed to the fetcher, defined as either direct values or `ParamResolvable` functions. + * Dynamic parameters using `AttachValue` will cause the store to refetch when their values change. + */ + params?: { + [K in keyof TParams]: ParamResolvable; + }; + /** + * The delay between retries after a fetch error occurs, in milliseconds, defined as a number or a function that + * receives the error and current retry count and returns a number. + * @default time.seconds(5) + */ + retryDelay?: number | ((retryCount: number, error: Error) => number); + /** + * The duration, in milliseconds, that data is considered fresh after fetching. + * After becoming stale, the store may automatically refetch data in the background if there are active subscribers. + * + * **Note:** Stale times under 5 seconds are strongly discouraged. + * @default time.minutes(2) + */ + staleTime?: number; + /** + * Suppresses warnings in the event a `staleTime` under the minimum is desired. + * @default false + */ + suppressStaleTimeWarning?: boolean; +}; + +/** + * Represents a parameter that can be provided directly or defined via a reactive `AttachValue`. + * A parameter can be: + * - A static value (e.g. `string`, `number`). + * - A function that returns an `AttachValue` when given a `SignalFunction`. + */ +type ParamResolvable, S extends StoreState, TData> = + | T + | (($: SignalFunction, store: QueryStore) => AttachValue); + +interface ResolvedParamsResult { + /** + * Direct, non-reactive values resolved from the initial configuration. + */ + directValues: Partial; + /** + * Reactive parameter values wrapped in `AttachValue`, which trigger refetches when they change. + */ + paramAttachVals: Partial>>; + /** + * Fully resolved parameters, merging both direct and reactive values. + */ + resolvedParams: TParams; +} + +/** + * The keys that make up the internal state of the store. + */ +type InternalStateKeys = keyof (StoreState> & PrivateStoreState); + +const [persist, discard] = [true, false]; + +const SHOULD_PERSIST_INTERNAL_STATE_MAP: Record = { + /* Internal state to persist if the store is persisted */ + enabled: persist, + error: persist, + lastFetchedAt: persist, + queryCache: persist, + queryKey: persist, + status: persist, + + /* Internal state and methods to discard */ + fetch: discard, + getData: discard, + getStatus: discard, + isDataExpired: discard, + isStale: discard, + reset: discard, + subscriptionCount: discard, +} satisfies Record; + +export const time = { + seconds: (n: number) => n * 1000, + minutes: (n: number) => time.seconds(n * 60), + hours: (n: number) => time.minutes(n * 60), + days: (n: number) => time.hours(n * 24), + weeks: (n: number) => time.days(n * 7), +}; + +const MIN_STALE_TIME = time.seconds(5); + +export function createQueryStore< + TQueryFnData, + TParams extends Record = Record, + U = unknown, + TData = TQueryFnData, +>( + config: RainbowQueryStoreConfig & PrivateStoreState & U> & { + params?: { [K in keyof TParams]: ParamResolvable & PrivateStoreState & U, TData> }; + }, + customStateCreator: StateCreator, + persistConfig?: RainbowPersistConfig & PrivateStoreState & U> +): QueryStore & U>; + +export function createQueryStore< + TQueryFnData, + TParams extends Record = Record, + U = unknown, + TData = TQueryFnData, +>( + config: RainbowQueryStoreConfig & PrivateStoreState & U> & { + params?: { [K in keyof TParams]: ParamResolvable & PrivateStoreState & U, TData> }; + }, + persistConfig?: RainbowPersistConfig & PrivateStoreState & U> +): QueryStore & U>; + +/** + * Creates a query-enabled Rainbow store with data fetching capabilities. + * @template TQueryFnData - The raw data type returned by the fetcher + * @template TParams - Parameters passed to the fetcher function + * @template U - User-defined custom store state + * @template TData - The transformed data type, if applicable (defaults to `TQueryFnData`) + */ +export function createQueryStore< + TQueryFnData, + TParams extends Record = Record, + U = unknown, + TData = TQueryFnData, +>( + config: RainbowQueryStoreConfig & PrivateStoreState & U> & { + params?: { [K in keyof TParams]: ParamResolvable & PrivateStoreState & U, TData> }; + }, + arg1?: + | StateCreator + | RainbowPersistConfig & PrivateStoreState & U>, + arg2?: RainbowPersistConfig & PrivateStoreState & U> +): QueryStore & U> { + type S = StoreState & PrivateStoreState & U; + + /* If arg1 is a function, it's the customStateCreator; otherwise, it's the persistConfig. */ + const customStateCreator = typeof arg1 === 'function' ? arg1 : () => ({}) as U; + const persistConfig = typeof arg1 === 'object' ? arg1 : arg2; + + const { + fetcher, + onFetched, + transform, + cacheTime = time.days(7), + disableCache = false, + enabled = true, + maxRetries = 3, + onError, + params, + retryDelay = time.seconds(5), + setData, + staleTime = time.minutes(2), + suppressStaleTimeWarning = false, + } = config; + + if (IS_DEV && !suppressStaleTimeWarning && staleTime < MIN_STALE_TIME) { + console.warn( + `[RainbowQueryStore${persistConfig?.storageKey ? `: ${persistConfig.storageKey}` : ''}] ❌ Stale times under ${ + MIN_STALE_TIME / 1000 + } seconds are not recommended.` + ); + } + + let directValues: Partial = {}; + let paramAttachVals: Partial>> = {}; + + let activeFetchPromise: Promise | null = null; + let activeRefetchTimeout: NodeJS.Timeout | null = null; + let lastFetchKey: string | null = null; + + const initialData = { + enabled, + error: null, + lastFetchedAt: null, + queryCache: {}, + queryKey: '', + status: QueryStatuses.Idle, + subscriptionCount: 0, + }; + + const getQueryKey = (params: TParams): string => JSON.stringify(Object.values(params)); + + const getCurrentResolvedParams = () => { + const currentParams = { ...directValues }; + for (const k in paramAttachVals) { + const attachVal = paramAttachVals[k as keyof TParams]; + if (!attachVal) continue; + currentParams[k as keyof TParams] = attachVal.value as TParams[keyof TParams]; + } + return currentParams as TParams; + }; + + const createState: StateCreator = (set, get, api) => { + const pruneCache = (state: S): S => { + const newCache: Record> = {}; + Object.entries(state.queryCache).forEach(([key, entry]) => { + if (entry && Date.now() - entry.lastFetchedAt <= cacheTime) { + newCache[key] = entry; + } + }); + return { ...state, queryCache: newCache }; + }; + + const scheduleNextFetch = (params: TParams, options: FetchOptions | undefined) => { + const effectiveStaleTime = options?.staleTime ?? staleTime; + if (effectiveStaleTime <= 0 || effectiveStaleTime === Infinity) return; + if (activeRefetchTimeout) { + clearTimeout(activeRefetchTimeout); + activeRefetchTimeout = null; + } + const currentQueryKey = get().queryKey; + const lastFetchedAt = + get().queryCache[currentQueryKey]?.lastFetchedAt || (disableCache && lastFetchKey === currentQueryKey ? get().lastFetchedAt : null); + const timeUntilRefetch = lastFetchedAt ? effectiveStaleTime - (Date.now() - lastFetchedAt) : effectiveStaleTime; + + activeRefetchTimeout = setTimeout(() => { + const { enabled, fetch, subscriptionCount } = get(); + if (enabled && subscriptionCount > 0) { + fetch(params, { force: true }); + } + }, timeUntilRefetch); + }; + + const baseMethods = { + async fetch(params: TParams | undefined, options: FetchOptions | undefined) { + if (!options?.force && !get().enabled) return; + + const effectiveParams = params ?? getCurrentResolvedParams(); + const { queryKey: currentQueryKey, status } = get(); + const isLoading = status === QueryStatuses.Loading; + + if (activeFetchPromise && lastFetchKey === currentQueryKey && isLoading && !options?.force) { + return activeFetchPromise; + } + + if (!options?.force && !disableCache) { + const { errorInfo, lastFetchedAt } = get().queryCache[currentQueryKey] ?? {}; + const errorRetriesExhausted = errorInfo && errorInfo.retryCount >= maxRetries; + + if ((!errorInfo || errorRetriesExhausted) && lastFetchedAt && Date.now() - lastFetchedAt <= (options?.staleTime ?? staleTime)) { + return; + } + } + + set(state => ({ ...state, error: null, status: QueryStatuses.Loading })); + lastFetchKey = currentQueryKey; + + const fetchOperation = async () => { + try { + const rawResult = await fetcher(effectiveParams); + let transformedData: TData; + try { + transformedData = transform ? transform(rawResult) : (rawResult as TData); + } catch (transformError) { + throw new RainbowError(`[createQueryStore: ${persistConfig?.storageKey || currentQueryKey}]: transform failed`, { + cause: transformError, + }); + } + + set(state => { + const lastFetchedAt = Date.now(); + let newState: S = { + ...state, + error: null, + lastFetchedAt, + status: QueryStatuses.Success, + }; + + if (!setData && !disableCache) { + newState.queryCache = { + ...newState.queryCache, + [currentQueryKey]: { + data: transformedData, + errorInfo: null, + lastFetchedAt, + }, + }; + } else if (setData) { + setData({ + data: transformedData, + params: effectiveParams, + queryKey: currentQueryKey, + set: (partial: S | Partial | ((state: S) => S | Partial)) => { + newState = typeof partial === 'function' ? { ...newState, ...partial(newState) } : { ...newState, ...partial }; + }, + }); + if (!disableCache) { + newState.queryCache = { + ...newState.queryCache, + [currentQueryKey]: { + data: null, + errorInfo: null, + lastFetchedAt, + }, + }; + } + } + + return disableCache || cacheTime === Infinity ? newState : pruneCache(newState); + }); + + scheduleNextFetch(effectiveParams, options); + + if (onFetched) { + try { + onFetched(transformedData, set); + } catch (onFetchedError) { + logger.error( + new RainbowError(`[createQueryStore: ${persistConfig?.storageKey || currentQueryKey}]: onFetched callback failed`, { + cause: onFetchedError, + }) + ); + } + } + } catch (error) { + const typedError = error instanceof Error ? error : new Error(String(error)); + const entry = get().queryCache[currentQueryKey]; + const currentRetryCount = entry?.errorInfo?.retryCount ?? 0; + + onError?.(typedError, currentRetryCount); + + if (currentRetryCount < maxRetries) { + if (get().subscriptionCount > 0) { + const errorRetryDelay = typeof retryDelay === 'function' ? retryDelay(currentRetryCount, typedError) : retryDelay; + + activeRefetchTimeout = setTimeout(() => { + const { enabled, fetch, subscriptionCount } = get(); + if (enabled && subscriptionCount > 0) { + fetch(params, { force: true }); + } + }, errorRetryDelay); + } + + set(state => ({ + ...state, + error: typedError, + status: QueryStatuses.Error, + queryCache: { + ...state.queryCache, + [currentQueryKey]: { + ...entry, + errorState: { + error: typedError, + retryCount: currentRetryCount + 1, + }, + }, + }, + })); + } else { + set(state => ({ + ...state, + status: QueryStatuses.Error, + queryCache: { + ...state.queryCache, + [currentQueryKey]: { + ...entry, + errorState: { + error: typedError, + retryCount: currentRetryCount, + }, + }, + }, + })); + } + + logger.error(new RainbowError(`[createQueryStore: ${persistConfig?.storageKey || currentQueryKey}]: Failed to fetch data`), { + error: typedError, + }); + } finally { + activeFetchPromise = null; + lastFetchKey = null; + } + }; + + activeFetchPromise = fetchOperation(); + return activeFetchPromise; + }, + + getData(params?: TParams) { + if (disableCache) return null; + const currentQueryKey = params ? getQueryKey(params) : get().queryKey; + return get().queryCache[currentQueryKey]?.data ?? null; + }, + + getStatus() { + const { queryKey, status } = get(); + const lastFetchedAt = + get().queryCache[queryKey]?.lastFetchedAt || (disableCache && lastFetchKey === queryKey ? get().lastFetchedAt : null); + + return { + isError: status === QueryStatuses.Error, + isFetching: status === QueryStatuses.Loading, + isIdle: status === QueryStatuses.Idle, + isInitialLoading: !lastFetchedAt && status === QueryStatuses.Loading, + isSuccess: status === QueryStatuses.Success, + }; + }, + + isDataExpired(cacheTimeOverride?: number) { + const { queryKey } = get(); + const lastFetchedAt = + get().queryCache[queryKey]?.lastFetchedAt || (disableCache && lastFetchKey === queryKey ? get().lastFetchedAt : null); + + if (!lastFetchedAt) return true; + const effectiveCacheTime = cacheTimeOverride ?? cacheTime; + return Date.now() - lastFetchedAt > effectiveCacheTime; + }, + + isStale(staleTimeOverride?: number) { + const { queryKey } = get(); + const lastFetchedAt = + get().queryCache[queryKey]?.lastFetchedAt || (disableCache && lastFetchKey === queryKey ? get().lastFetchedAt : null); + + if (!lastFetchedAt) return true; + const effectiveStaleTime = staleTimeOverride ?? staleTime; + return Date.now() - lastFetchedAt > effectiveStaleTime; + }, + + reset() { + if (activeRefetchTimeout) { + clearTimeout(activeRefetchTimeout); + activeRefetchTimeout = null; + } + activeFetchPromise = null; + lastFetchKey = null; + set(state => ({ ...state, ...initialData, queryKey: getQueryKey(getCurrentResolvedParams()) })); + }, + }; + + const subscribeWithSelector = api.subscribe; + + api.subscribe = (listener: (state: S, prevState: S) => void) => { + set(state => ({ ...state, subscriptionCount: state.subscriptionCount + 1 })); + const unsubscribe = subscribeWithSelector(listener); + + const handleSetEnabled = subscribeWithSelector((state: S, prevState: S) => { + if (state.enabled !== prevState.enabled) { + if (state.enabled) { + const currentParams = getCurrentResolvedParams(); + const currentKey = state.queryKey; + if (currentKey !== lastFetchKey) { + state.fetch(currentParams, { force: true }); + } else if (!state.queryCache[currentKey] || state.isStale()) { + state.fetch(); + } else { + scheduleNextFetch(currentParams, undefined); + } + } else { + if (activeRefetchTimeout) { + clearTimeout(activeRefetchTimeout); + activeRefetchTimeout = null; + } + } + } + }); + + const { enabled, fetch, isStale, queryKey } = get(); + + const currentParams = getCurrentResolvedParams(); + set(state => ({ ...state, queryKey: getQueryKey(currentParams) })); + + if (!get().queryCache[queryKey] || isStale()) { + fetch(currentParams); + } else if (enabled) { + scheduleNextFetch(currentParams, undefined); + } + + return () => { + handleSetEnabled(); + unsubscribe(); + set(state => { + const newCount = Math.max(state.subscriptionCount - 1, 0); + if (newCount === 0) { + if (activeRefetchTimeout) { + clearTimeout(activeRefetchTimeout); + activeRefetchTimeout = null; + } + } + return { ...state, subscriptionCount: newCount }; + }); + }; + }; + + const userState = customStateCreator?.(set, get, api) ?? ({} as U); + + /* Merge base data, user state, and methods into the final store state */ + return { + ...initialData, + ...userState, + ...baseMethods, + }; + }; + + const combinedPersistConfig = persistConfig + ? { + ...persistConfig, + partialize: createBlendedPartialize(persistConfig.partialize), + } + : undefined; + + const baseStore = persistConfig?.storageKey + ? createRainbowStore(createState, combinedPersistConfig) + : create(subscribeWithSelector(createState)); + + const queryCapableStore: QueryStore = Object.assign(baseStore, { + enabled, + queryKey: baseStore.getState().queryKey, + fetch: (params?: TParams, options?: FetchOptions) => baseStore.getState().fetch(params, options), + getData: () => baseStore.getState().getData(), + getStatus: () => baseStore.getState().getStatus(), + isDataExpired: (override?: number) => baseStore.getState().isDataExpired(override), + isStale: (override?: number) => baseStore.getState().isStale(override), + reset: () => { + for (const unsub of paramUnsubscribes) unsub(); + paramUnsubscribes = []; + queryCapableStore.getState().reset(); + queryCapableStore.setState(state => ({ ...state, ...initialData, queryKey: getQueryKey(getCurrentResolvedParams()) })); + }, + }); + + if (params) { + const result = resolveParams(params, queryCapableStore); + paramAttachVals = result.paramAttachVals; + directValues = result.directValues; + } + + const onParamChange = () => { + const newQueryKey = getQueryKey(getCurrentResolvedParams()); + queryCapableStore.setState(state => ({ ...state, queryKey: newQueryKey })); + queryCapableStore.fetch(); + }; + + let paramUnsubscribes: Unsubscribe[] = []; + + for (const k in paramAttachVals) { + const attachVal = paramAttachVals[k]; + if (!attachVal) continue; + + const subscribeFn = attachValueSubscriptionMap.get(attachVal); + if (ENABLE_LOGS) console.log('[πŸŒ€ ParamSubscription πŸŒ€] Subscribed to param:', k); + + if (subscribeFn) { + let oldVal = attachVal.value; + const unsub = subscribeFn(() => { + const newVal = attachVal.value; + if (!Object.is(oldVal, newVal)) { + oldVal = newVal; + if (ENABLE_LOGS) console.log('[πŸŒ€ ParamChange πŸŒ€] Param changed:', k); + onParamChange(); + } + }); + paramUnsubscribes.push(unsub); + } + } + + return queryCapableStore; +} + +function resolveParams, S extends StoreState & U, TData, U = unknown>( + params: { [K in keyof TParams]: ParamResolvable }, + store: QueryStore +): ResolvedParamsResult { + const directValues: Partial = {}; + const paramAttachVals: Partial>> = {}; + const resolvedParams = {} as TParams; + + for (const key in params) { + const param = params[key]; + if (typeof param === 'function') { + const attachVal = param($, store); + resolvedParams[key] = attachVal.value as TParams[typeof key]; + paramAttachVals[key] = attachVal; + } else { + resolvedParams[key] = param as TParams[typeof key]; + directValues[key] = param as TParams[typeof key]; + } + } + + return { directValues, paramAttachVals, resolvedParams }; +} + +function createBlendedPartialize, S extends StoreState & U, U = unknown>( + userPartialize: ((state: StoreState & U) => Partial & U>) | undefined +) { + return (state: S) => { + const clonedState = { ...state }; + const internalStateToPersist: Partial = {}; + + for (const key in clonedState) { + if (key in SHOULD_PERSIST_INTERNAL_STATE_MAP) { + if (SHOULD_PERSIST_INTERNAL_STATE_MAP[key]) internalStateToPersist[key] = clonedState[key]; + delete clonedState[key]; + } + } + + return { + ...(userPartialize ? userPartialize(clonedState) : omitStoreMethods(clonedState)), + ...internalStateToPersist, + } satisfies Partial; + }; +} diff --git a/src/state/internal/createRainbowStore.ts b/src/state/internal/createRainbowStore.ts index e0c14b79ead..6f6ef4f05f5 100644 --- a/src/state/internal/createRainbowStore.ts +++ b/src/state/internal/createRainbowStore.ts @@ -1,8 +1,7 @@ import { debounce } from 'lodash'; import { MMKV } from 'react-native-mmkv'; -import { create } from 'zustand'; +import { StateCreator, create } from 'zustand'; import { PersistOptions, StorageValue, persist, subscribeWithSelector } from 'zustand/middleware'; -import { StateCreator } from 'zustand/vanilla'; import { RainbowError, logger } from '@/logger'; const PERSIST_RATE_LIMIT_MS = 3000; @@ -12,12 +11,23 @@ const rainbowStorage = new MMKV({ id: 'rainbow-storage' }); /** * Configuration options for creating a persistable Rainbow store. */ -interface RainbowPersistConfig { +export interface RainbowPersistConfig { /** * A function to convert the serialized string back into the state object. * If not provided, the default deserializer is used. */ deserializer?: (serializedState: string) => StorageValue>; + /** + * A function to perform persisted state migration. + * This function will be called when persisted state versions mismatch with the one specified here. + */ + migrate?: (persistedState: unknown, version: number) => S | Promise; + /** + * A function returning another (optional) function. + * The main function will be called before the state rehydration. + * The returned function will be called after the state rehydration or when an error occurred. + */ + onRehydrateStorage?: PersistOptions>['onRehydrateStorage']; /** * A function that determines which parts of the state should be persisted. * By default, the entire state is persisted. @@ -38,17 +48,51 @@ interface RainbowPersistConfig { * @default 0 */ version?: number; - /** - * A function to perform persisted state migration. - * This function will be called when persisted state versions mismatch with the one specified here. - */ - migrate?: (persistedState: unknown, version: number) => S | Promise; - /** - * A function returning another (optional) function. - * The main function will be called before the state rehydration. - * The returned function will be called after the state rehydration or when an error occurred. - */ - onRehydrateStorage?: PersistOptions>['onRehydrateStorage']; +} + +/** + * Creates a Rainbow store with optional persistence functionality. + * @param createState - The state creator function for the Rainbow store. + * @param persistConfig - The configuration options for the persistable Rainbow store. + * @returns A Zustand store with the specified state and optional persistence. + */ +export function createRainbowStore( + createState: StateCreator, + persistConfig?: RainbowPersistConfig +) { + if (!persistConfig) return create()(subscribeWithSelector(createState)); + + const { persistStorage, version } = createPersistStorage(persistConfig); + + return create()( + subscribeWithSelector( + persist(createState, { + migrate: persistConfig.migrate, + name: persistConfig.storageKey, + onRehydrateStorage: persistConfig.onRehydrateStorage, + partialize: persistConfig.partialize || omitStoreMethods, + storage: persistStorage, + version, + }) + ) + ); +} + +/** + * Default partialize function if none is provided. It omits top-level store + * methods and keeps all other state. + */ +export function omitStoreMethods(state: S): Partial { + if (state !== null && typeof state === 'object') { + const result: Record = {}; + Object.entries(state).forEach(([key, val]) => { + if (typeof val !== 'function') { + result[key] = val; + } + }); + return result as Partial; + } + return state; } /** @@ -66,7 +110,7 @@ function createPersistStorage(config: RainbowPersistConfig) { if (!serializedValue) return null; return deserializer(serializedValue); }, - setItem: (name, value) => + setItem: (name: string, value: StorageValue>) => lazyPersist({ serializer, storageKey, @@ -118,7 +162,7 @@ const lazyPersist = ({ name, serializer, storageKey, value }: LazyPersistPara */ function defaultSerializeState(state: StorageValue>['state'], version: StorageValue>['version']): string { try { - return JSON.stringify({ state, version }); + return JSON.stringify({ state, version }, replacer); } catch (error) { logger.error(new RainbowError(`[createRainbowStore]: Failed to serialize Rainbow store data`), { error }); throw error; @@ -132,39 +176,51 @@ function defaultSerializeState(state: StorageValue>['state'], vers */ function defaultDeserializeState(serializedState: string): StorageValue> { try { - return JSON.parse(serializedState); + return JSON.parse(serializedState, reviver); } catch (error) { logger.error(new RainbowError(`[createRainbowStore]: Failed to deserialize persisted Rainbow store data`), { error }); throw error; } } +interface SerializedMap { + __type: 'Map'; + entries: [unknown, unknown][]; +} + +function isSerializedMap(value: unknown): value is SerializedMap { + return typeof value === 'object' && value !== null && (value as Record).__type === 'Map'; +} + +interface SerializedSet { + __type: 'Set'; + values: unknown[]; +} + +function isSerializedSet(value: unknown): value is SerializedSet { + return typeof value === 'object' && value !== null && (value as Record).__type === 'Set'; +} + /** - * Creates a Rainbow store with optional persistence functionality. - * @param createState - The state creator function for the Rainbow store. - * @param persistConfig - The configuration options for the persistable Rainbow store. - * @returns A Zustand store with the specified state and optional persistence. + * Replacer function to handle serialization of Maps and Sets. */ -export function createRainbowStore( - createState: StateCreator, - persistConfig?: RainbowPersistConfig -) { - if (!persistConfig) { - return create()(subscribeWithSelector(createState)); +function replacer(key: string, value: unknown): unknown { + if (value instanceof Map) { + return { __type: 'Map', entries: Array.from(value.entries()) }; + } else if (value instanceof Set) { + return { __type: 'Set', values: Array.from(value) }; } + return value; +} - const { persistStorage, version } = createPersistStorage(persistConfig); - - return create()( - subscribeWithSelector( - persist(createState, { - name: persistConfig.storageKey, - partialize: persistConfig.partialize || (state => state), - storage: persistStorage, - version, - migrate: persistConfig.migrate, - onRehydrateStorage: persistConfig.onRehydrateStorage, - }) - ) - ); +/** + * Reviver function to handle deserialization of Maps and Sets. + */ +function reviver(key: string, value: unknown): unknown { + if (isSerializedMap(value)) { + return new Map(value.entries); + } else if (isSerializedSet(value)) { + return new Set(value.values); + } + return value; } diff --git a/src/state/internal/createStore.ts b/src/state/internal/createStore.ts index 3c49c5eb18a..05006491549 100644 --- a/src/state/internal/createStore.ts +++ b/src/state/internal/createStore.ts @@ -8,6 +8,9 @@ export type StoreWithPersist = Mutate, [['zustand/persi initializer: Initializer; }; +/** + * @deprecated This is a legacy store creator. Use `createRainbowStore` instead. + */ export function createStore( initializer: Initializer, { persist: persistOptions }: { persist?: PersistOptions } = {} diff --git a/src/state/internal/signal.ts b/src/state/internal/signal.ts new file mode 100644 index 00000000000..536f0ba5690 --- /dev/null +++ b/src/state/internal/signal.ts @@ -0,0 +1,166 @@ +import { StoreApi } from 'zustand'; + +const ENABLE_LOGS = false; + +/* Store subscribe function so we can handle param changes on any attachVal (root or nested) */ +export const attachValueSubscriptionMap = new WeakMap, Subscribe>(); + +/* Global caching for top-level attachValues */ +const storeSignalCache = new WeakMap< + StoreApi, + Map<(state: unknown) => unknown, Map<(a: unknown, b: unknown) => boolean, AttachValue>> +>(); + +export type AttachValue = T & { value: T } & { + readonly [K in keyof T]: AttachValue; +}; + +export type SignalFunction = { + (store: StoreApi): AttachValue; + (store: StoreApi, selector: (state: T) => S, equalityFn?: (a: S, b: S) => boolean): AttachValue; +}; + +export type Unsubscribe = () => void; +export type Subscribe = (callback: () => void) => Unsubscribe; +export type GetValue = () => unknown; +export type SetValue = (path: unknown[], value: unknown) => void; + +export function $(store: StoreApi): AttachValue; +export function $(store: StoreApi, selector: (state: T) => S, equalityFn?: (a: S, b: S) => boolean): AttachValue; +export function $( + store: StoreApi, + selector: (state: unknown) => unknown = identity, + equalityFn: (a: unknown, b: unknown) => boolean = Object.is +) { + return getOrCreateAttachValue(store, selector, equalityFn); +} + +const identity = (x: T): T => x; + +const updateValue = (obj: T, path: unknown[], value: unknown): T => { + if (!path.length) { + return value as T; + } + const [first, ...rest] = path; + const prevValue = (obj as Record)[first as string]; + const nextValue = updateValue(prevValue, rest, value); + if (Object.is(prevValue, nextValue)) { + return obj; + } + const copied = Array.isArray(obj) ? obj.slice() : { ...obj }; + (copied as Record)[first as string] = nextValue; + return copied as T; +}; + +export const createSignal = ( + store: StoreApi, + selector: (state: T) => S, + equalityFn: (a: S, b: S) => boolean +): [Subscribe, GetValue, SetValue] => { + let selected = selector(store.getState()); + const listeners = new Set<() => void>(); + let unsubscribe: Unsubscribe | undefined; + + const sub: Subscribe = callback => { + if (!listeners.size) { + unsubscribe = store.subscribe(() => { + const nextSelected = selector(store.getState()); + if (!equalityFn(selected, nextSelected)) { + selected = nextSelected; + listeners.forEach(listener => listener()); + } + }); + } + listeners.add(callback); + return () => { + listeners.delete(callback); + if (!listeners.size && unsubscribe) { + unsubscribe(); + unsubscribe = undefined; + } + }; + }; + + const get: GetValue = () => { + if (!listeners.size) { + selected = selector(store.getState()); + } + return selected; + }; + + const set: SetValue = (path, value) => { + if (selector !== identity) { + throw new Error('Cannot set a value with a selector'); + } + store.setState(prev => updateValue(prev, path, value), true); + }; + + return [sub, get, set]; +}; + +function getOrCreateAttachValue(store: StoreApi, selector: (state: T) => S, equalityFn: (a: S, b: S) => boolean): AttachValue { + let bySelector = storeSignalCache.get(store); + if (!bySelector) { + bySelector = new Map(); + storeSignalCache.set(store, bySelector); + } + + let byEqFn = bySelector.get(selector as (state: unknown) => unknown); + if (!byEqFn) { + byEqFn = new Map(); + bySelector.set(selector as (state: unknown) => unknown, byEqFn); + } + + const existing = byEqFn.get(equalityFn as (a: unknown, b: unknown) => boolean); + if (existing) { + return existing as AttachValue; + } + + const [subscribe, getVal, setVal] = createSignal(store, selector, equalityFn); + + const localCache = new Map>(); + + const createAttachValue = (fullPath: string): AttachValue => { + const handler: ProxyHandler = { + get(_, key) { + if (key === 'value') { + let v = getVal(); + const parts = fullPath.split('.'); + for (const p of parts) { + if (p) v = (v as Record)[p]; + } + return v; + } + const keyString = typeof key === 'string' ? key : key.toString(); + const pathKey = fullPath ? `${fullPath}.${keyString}` : keyString; + const cached = localCache.get(pathKey); + if (cached) { + if (ENABLE_LOGS) console.log('[πŸŒ€ AttachValue πŸŒ€] Cache hit for:', pathKey); + return cached; + } else if (ENABLE_LOGS) { + console.log('[πŸŒ€ AttachValue πŸŒ€] Created root attachValue:', pathKey); + } + const val = createAttachValue(pathKey); + attachValueSubscriptionMap.set(val, subscribe); + localCache.set(pathKey, val); + return val; + }, + set(_, __, value) { + const path = fullPath.split('.'); + if (path[0] === '') path.shift(); + setVal(path, value); + return true; + }, + }; + + return new Proxy(Object.create(null), handler) as AttachValue; + }; + + const rootVal = createAttachValue(''); + subscribe(() => { + return; + }); + attachValueSubscriptionMap.set(rootVal, subscribe); + byEqFn.set(equalityFn as (a: unknown, b: unknown) => boolean, rootVal); + return rootVal as AttachValue; +} diff --git a/src/state/internal/tests/QueryStoreTest.tsx b/src/state/internal/tests/QueryStoreTest.tsx new file mode 100644 index 00000000000..6041d131a0b --- /dev/null +++ b/src/state/internal/tests/QueryStoreTest.tsx @@ -0,0 +1,235 @@ +// ⚠️ Uncomment everything below to experiment with the QueryStore creator +// TODO: Comment out test code below before merging + +import React, { memo, useEffect, useMemo } from 'react'; +import { StyleSheet, View } from 'react-native'; +import { Address } from 'viem'; +import ButtonPressAnimation from '@/components/animations/ButtonPressAnimation'; +import { ImgixImage } from '@/components/images'; +import { Text, useForegroundColor } from '@/design-system'; +import { logger, RainbowError } from '@/logger'; +import { SupportedCurrencyKey } from '@/references'; +import { addysHttp } from '@/resources/addys/claimables/query'; +import { parseUserAssets } from '@/__swaps__/screens/Swap/resources/assets/userAssets'; +import { ParsedAssetsDictByChain } from '@/__swaps__/types/assets'; +import { AddressAssetsReceivedMessage } from '@/__swaps__/types/refraction'; +import { useBackendNetworksStore } from '@/state/backendNetworks/backendNetworks'; +import { createQueryStore, time } from '../createQueryStore'; +import { createRainbowStore } from '../createRainbowStore'; + +const ENABLE_LOGS = false; + +type CurrencyStore = { + nestedParamTest: { + currency: SupportedCurrencyKey; + }; + setCurrency: (currency: SupportedCurrencyKey) => void; +}; + +const useCurrencyStore = createRainbowStore((set, get) => ({ + nestedParamTest: { currency: 'USD' }, + setCurrency: (currency: SupportedCurrencyKey) => { + set({ nestedParamTest: { currency } }); + if (ENABLE_LOGS) console.log('[πŸ‘€ useCurrencyStore πŸ‘€] New currency set:', get().nestedParamTest.currency); + }, +})); + +type UserAssetsTestStore = { + address: Address; + setAddress: (address: Address) => void; +}; + +type UserAssetsQueryParams = { address: Address; currency: SupportedCurrencyKey }; + +const testAddresses: Address[] = [ + '0x2e67869829c734ac13723A138a952F7A8B56e774', + '0xCFB83E14AEd465c79F3F82f4cfF4ff7965897644', + '0x17cd072cBd45031EFc21Da538c783E0ed3b25DCc', +]; + +export const useUserAssetsTestStore = createQueryStore( + { + fetcher: ({ address, currency }) => simpleUserAssetsQuery({ address, currency }), + params: { + address: ($, store) => $(store).address, + currency: $ => $(useCurrencyStore).nestedParamTest.currency, + }, + staleTime: time.minutes(1), + }, + + set => ({ + address: testAddresses[0], + setAddress: (address: Address) => set({ address }), + }), + + { storageKey: 'queryStoreTest' } +); + +export const UserAssetsTest = memo(function UserAssetsTest() { + const data = useUserAssetsTestStore(state => state.getData()); + const enabled = useUserAssetsTestStore(state => state.enabled); + + const firstFiveCoinIconUrls = useMemo(() => (data ? getFirstFiveCoinIconUrls(data) : Array.from({ length: 5 }).map(() => '')), [data]); + const skeletonColor = useForegroundColor('fillQuaternary'); + + useEffect(() => { + if (ENABLE_LOGS && data) { + const first5Tokens = Object.values(data) + .flatMap(chainAssets => Object.values(chainAssets)) + .slice(0, 5); + console.log('[πŸ”” UserAssetsTest πŸ””] userAssets data updated - first 5 tokens:', first5Tokens.map(token => token.symbol).join(', ')); + } + }, [data]); + + useEffect(() => { + if (ENABLE_LOGS) console.log(`[πŸ”” UserAssetsTest πŸ””] enabled updated to: ${enabled ? 'βœ… ENABLED' : 'πŸ›‘ DISABLED'}`); + }, [enabled]); + + return ( + + + {firstFiveCoinIconUrls.map((url, index) => + url ? ( + + ) : ( + + ) + )} + + + {data + ? `Number of assets: ${Object.values(data).reduce((acc, chainAssets) => acc + Object.keys(chainAssets).length, 0)}` + : 'Loading…'} + + + { + const currentAddress = useUserAssetsTestStore.getState().address; + switch (currentAddress) { + case testAddresses[0]: + useUserAssetsTestStore.getState().setAddress(testAddresses[1]); + break; + case testAddresses[1]: + useUserAssetsTestStore.getState().setAddress(testAddresses[2]); + break; + case testAddresses[2]: + useUserAssetsTestStore.getState().setAddress(testAddresses[0]); + break; + } + }} + style={styles.button} + > + + Shuffle Address + + + { + useUserAssetsTestStore.setState({ enabled: !enabled }); + }} + style={styles.button} + > + + {useUserAssetsTestStore.getState().enabled ? 'Disable Fetching' : 'Enable Fetching'} + + + + + ); +}); + +if (ENABLE_LOGS) console.log('[πŸ’Ύ UserAssetsTest πŸ’Ύ] Initial data exists:', !!useUserAssetsTestStore.getState().getData()); + +// eslint-disable-next-line @typescript-eslint/no-unused-vars +function logFetchInfo(params: UserAssetsQueryParams) { + const formattedTimeWithSeconds = new Date(Date.now()).toLocaleTimeString('en-US', { + hour: '2-digit', + minute: '2-digit', + second: '2-digit', + }); + console.log('[πŸ”„ UserAssetsTest - logFetchInfo πŸ”„]', '\nTime:', formattedTimeWithSeconds, '\nParams:', { + address: params.address, + currency: params.currency, + raw: JSON.stringify(Object.values(params), null, 2), + }); +} + +function getFirstFiveCoinIconUrls(data: ParsedAssetsDictByChain) { + const result: string[] = []; + outer: for (const chainAssets of Object.values(data)) { + for (const token of Object.values(chainAssets)) { + if (token.icon_url) { + result.push(token.icon_url); + if (result.length === 5) { + break outer; + } + } + } + } + return result; +} + +type FetchUserAssetsArgs = { + address: Address | string; + currency: SupportedCurrencyKey; + testnetMode?: boolean; +}; + +export async function simpleUserAssetsQuery({ address, currency }: FetchUserAssetsArgs): Promise { + if (!address) return {}; + try { + const url = `/${useBackendNetworksStore.getState().getSupportedChainIds().join(',')}/${address}/assets?currency=${currency.toLowerCase()}`; + const res = await addysHttp.get(url, { + timeout: time.seconds(20), + }); + const chainIdsInResponse = res?.data?.meta?.chain_ids || []; + const assets = res?.data?.payload?.assets?.filter(asset => !asset.asset.defi_position) || []; + + if (assets.length && chainIdsInResponse.length) { + return parseUserAssets({ + assets, + chainIds: chainIdsInResponse, + currency, + }); + } + return {}; + } catch (e) { + logger.error(new RainbowError('[simpleUserAssetsQuery]: Failed to fetch user assets'), { + message: (e as Error)?.message, + }); + return {}; + } +} + +const styles = StyleSheet.create({ + button: { + alignItems: 'center', + backgroundColor: 'blue', + borderRadius: 22, + height: 44, + justifyContent: 'center', + paddingHorizontal: 20, + }, + buttonGroup: { + alignItems: 'center', + flexDirection: 'column', + gap: 24, + justifyContent: 'center', + }, + coinIcon: { + borderRadius: 16, + height: 32, + width: 32, + }, + coinIconContainer: { + flexDirection: 'row', + gap: 12, + }, + container: { + alignItems: 'center', + flex: 1, + flexDirection: 'column', + gap: 32, + justifyContent: 'center', + }, +});