diff --git a/app/gui/src/dashboard/hooks/storeHooks.ts b/app/gui/src/dashboard/hooks/storeHooks.ts new file mode 100644 index 000000000000..5b712a4789c2 --- /dev/null +++ b/app/gui/src/dashboard/hooks/storeHooks.ts @@ -0,0 +1,159 @@ +/** + * @file + * + * This file contains hooks for using Zustand store with tearing transitions. + */ +import type { DispatchWithoutAction, Reducer, RefObject } from 'react' +import { useEffect, useReducer, useRef } from 'react' +import { type StoreApi } from 'zustand' +import { useStoreWithEqualityFn } from 'zustand/traditional' +import { objectEquality, refEquality, shallowEquality } from '../utilities/equalities' + +/** + * A type that allows to choose between different equality functions. + */ +export type AreEqual = EqualityFunction | EqualityFunctionName +/** + * Custom equality function. + */ +export type EqualityFunction = (a: T, b: T) => boolean +/** + * Equality function name from a list of predefined ones. + */ +export type EqualityFunctionName = 'object' | 'shallow' | 'strict' + +const EQUALITY_FUNCTIONS: Record boolean> = { + object: objectEquality, + shallow: shallowEquality, + strict: refEquality, +} + +/** Options for the `useStore` hook. */ +export interface UseStoreOptions { + /** + * Adds support for React transitions. + * + * Use it with caution, as it may lead to inconsistent state during transitions. + */ + readonly unsafeEnableTransition?: boolean + /** + * Specifies the equality function to use. + * @default 'Object.is' + */ + readonly areEqual?: AreEqual +} + +/** + * A wrapper that allows to choose between tearing transition and standard Zustand store. + * + * # `options.unsafeEnableTransition` must not be changed during the component lifecycle. + */ +export function useStore( + store: StoreApi, + selector: (state: State) => Slice, + options: UseStoreOptions = {}, +) { + const { unsafeEnableTransition = false, areEqual } = options + + const prevUnsafeEnableTransition = useRef(unsafeEnableTransition) + + const equalityFunction = resolveAreEqual(areEqual) + + return useNonCompilableConditionalStore( + store, + selector, + unsafeEnableTransition, + equalityFunction, + prevUnsafeEnableTransition, + ) +} + +/** A hook that allows to use React transitions with Zustand store. */ +export function useTearingTransitionStore( + store: StoreApi, + selector: (state: State) => Slice, + areEqual: AreEqual = 'object', +) { + const state = store.getState() + + const equalityFunction = resolveAreEqual(areEqual) + + const [[sliceFromReducer, storeFromReducer], rerender] = useReducer< + Reducer< + readonly [Slice, StoreApi, State], + readonly [Slice, StoreApi, State] | undefined + >, + undefined + >( + (prev, fromSelf) => { + if (fromSelf) { + return fromSelf + } + const nextState = store.getState() + if (Object.is(prev[2], nextState) && prev[1] === store) { + return prev + } + const nextSlice = selector(nextState) + if (equalityFunction(prev[0], nextSlice) && prev[1] === store) { + return prev + } + return [nextSlice, store, nextState] + }, + undefined, + () => [selector(state), store, state], + ) + + useEffect(() => { + const unsubscribe = store.subscribe(() => { + // eslint-disable-next-line no-restricted-syntax + ;(rerender as DispatchWithoutAction)() + }) + // eslint-disable-next-line no-restricted-syntax + ;(rerender as DispatchWithoutAction)() + return unsubscribe + }, [store]) + + if (storeFromReducer !== store) { + const slice = selector(state) + rerender([slice, store, state]) + return slice + } + + return sliceFromReducer +} + +/** Resolves the equality function. */ +function resolveAreEqual(areEqual: AreEqual | null | undefined) { + return ( + areEqual == null ? EQUALITY_FUNCTIONS.object + : typeof areEqual === 'string' ? EQUALITY_FUNCTIONS[areEqual] + : areEqual + ) +} + +/** + * Internal hook that isolates the conditional store logic from the `useStore` hook. + * To enable compiler optimizations for the `useStore` hook. + * @internal + * @throws An error if the `unsafeEnableTransition` option is changed during the component lifecycle. + */ +function useNonCompilableConditionalStore( + store: StoreApi, + selector: (state: State) => Slice, + unsafeEnableTransition: boolean, + equalityFunction: EqualityFunction, + prevUnsafeEnableTransition: RefObject, +) { + /* eslint-disable react-compiler/react-compiler */ + /* eslint-disable react-hooks/rules-of-hooks */ + if (prevUnsafeEnableTransition.current !== unsafeEnableTransition) { + throw new Error( + 'useStore shall not change the `unsafeEnableTransition` option during the component lifecycle', + ) + } + return unsafeEnableTransition ? + useTearingTransitionStore(store, selector, equalityFunction) + : useStoreWithEqualityFn(store, selector, equalityFunction) + /* eslint-enable react-compiler/react-compiler */ + /* eslint-enable react-hooks/rules-of-hooks */ +} diff --git a/app/gui/src/dashboard/utilities/equalities.ts b/app/gui/src/dashboard/utilities/equalities.ts new file mode 100644 index 000000000000..ef34b7fcff61 --- /dev/null +++ b/app/gui/src/dashboard/utilities/equalities.ts @@ -0,0 +1,51 @@ +/** + * @file + * + * This file contains functions for checking equality between values. + */ + +/** + * Strict equality check. + */ +export function refEquality(a: T, b: T) { + return a === b +} + +/** + * Object.is equality check. + */ +export function objectEquality(a: T, b: T) { + return Object.is(a, b) +} + +/** + * Shallow equality check. + */ +export function shallowEquality(a: T, b: T) { + if (Object.is(a, b)) { + return true + } + + if (typeof a !== 'object' || a == null || typeof b !== 'object' || b == null) { + return false + } + + const keysA = Object.keys(a) + + if (keysA.length !== Object.keys(b).length) { + return false + } + + for (let i = 0; i < keysA.length; i++) { + const key = keysA[i] + + if (key != null) { + // @ts-expect-error Typescript doesn't know that key is in a and b, but it doesn't matter here + if (!Object.prototype.hasOwnProperty.call(b, key) || !Object.is(a[key], b[key])) { + return false + } + } + } + + return true +} diff --git a/app/gui/src/dashboard/utilities/zustand.ts b/app/gui/src/dashboard/utilities/zustand.ts new file mode 100644 index 000000000000..94d4a40606ef --- /dev/null +++ b/app/gui/src/dashboard/utilities/zustand.ts @@ -0,0 +1,14 @@ +/** + * @file + * + * Re-exporting zustand functions and types. + * Overrides the default `useStore` with a custom one, that supports equality functions and React.transition + */ +export { useStore, useTearingTransitionStore } from '#/hooks/storeHooks' +export type { + AreEqual, + EqualityFunction, + EqualityFunctionName, + UseStoreOptions, +} from '#/hooks/storeHooks' +export * from 'zustand'