reactHooksUtils.ts

reactHooksUtils.ts

/**
 * Utility functions for contexts and jotai store scoped
 * */
import { atom, useAtomValue, useSetAtom, createStore } from 'jotai';
import {
  useContext,
  createContext,
  useMemo,
  SetStateAction,
  useRef,
} from 'react';

const jotaiStore = createStore();

/**
 * ------------------------------------
 *  Create a safe context. Use this instead of `createContext` from `react` wherever possible
 *
 *  Creating contexts:
 * ```typescriptreact
 * const [FloatingDockContextProvider, useFloatingDockContext] =
 *   createSafeContext<{ mouseX: number }>(
 *     'useFloatingDockContext must be used within a FloatingDockProvider'
 *   );
 * ```
 *
 * using context:
 * ```typescriptreact
 * const { mouseX } = useFloatingDockContext();
 * ```
 * */
export function createSafeContext<ContextValue>(errorMessage: string) {
  const Context = createContext<ContextValue | null>(null);

  const useSafeContext = () => {
    const ctx = useContext(Context);

    if (ctx === null) {
      throw new Error(errorMessage);
    }

    return ctx;
  };

  const Provider = ({
    children,
    value,
  }: {
    value: ContextValue;
    children: React.ReactNode;
  }) => <Context.Provider value={value}>{children}</Context.Provider>;

  return [Provider, useSafeContext] as const;
}

// --------------------------------------------

/** `(state: T) => void` Copied from jotai utils */
type SetAtom<Args extends unknown[], Result> = (...args: Args) => Result;
type UseScopedSetAtom<T> = () => SetAtom<[SetStateAction<T>], void>;

/**
 * - Call with no selector: returns TValue
 * - Call with selector: returns TSelected (the selector's result type)
 **/
type UseJotaiStore<TValue> = {
  (): TValue;
  <TSelected>(selector: (state: TValue) => TSelected): TSelected;
};
type UseStoreSubscribe<T> = () => (callback: (value: T) => void) => () => void;

type UseGetStore<T> = () => () => T;

/**
 * Create a scoped jotai store
 * ---
 * Use this when you want to create atoms that are not global, specially for handling UI
 * when you want fine-grained reactivity.
 *
 * Usage:
 * ```typescriptreact
 * interface BearStore {
 *   bears: string[];
 * }
 * const {
 *   StoreProvider: BearStoreProvider,
 *   useValue: useBearStore,
 *   useSetValue: useSetBearStore
 *   useSubscribe: useBearStoreSubscribe,
 *   useGetValue: useGetBearStore
 * } = createScopedJotaiStore<BearStore>({
 *     bears: [],
 * });
 *
 * // In your component
 * function BearCounter() {
 *   const bears = useBearStore(state => state.bears);
 *   return <div>{bears.length} bears</div>;
 * }
 *
 * function BearAdder() {
 *   // you can use this with a selector or no selector
 *   const setBearStore = useSetBearStore();
 *
 *   const addBear = (id: string) => {
 *     setBearStore((prev) => ({
 *       ...prev,
 *       bears: [...prev.bears, id]
 *     }));
 *   };
 *
 *   const removeBear = (id: string) => {
 *     setBearStore((prev) => ({
 *       ...prev,
 *       bears: prev.bears.filter(bearId => bearId !== id)
 *     }));
 *   };
 *
 *   return <button onClick={() => addBear('bear-' + Date.now())}>Add Bear</button>;
 * }
 *
 * // In your app
 * function App() {
 *   return (
 *     <BearStoreProvider>
 *       <BearCounter />
 *       <BearAdder />
 *     </BearStoreProvider>
 *   );
 * }
 * ```
 * */
export function createScopedStore<T>(initialValue: T): {
  StoreProvider: React.FC<{ children: React.ReactNode }>;
  useValue: UseJotaiStore<T>;
  useSetValue: UseScopedSetAtom<T>;
  useSubscribe: UseStoreSubscribe<T>;
  useGetValue: UseGetStore<T>;
} {
  // Create a context to hold the atom
  const JotaiAtomContext = createContext<ReturnType<typeof atom<T>> | null>(
    null
  );

  function StoreProvider({ children }: { children: React.ReactNode }) {
    // Create the atom with initial value
    const storeAtom = useMemo(() => atom<T>(initialValue), []);

    return (
      <JotaiAtomContext.Provider value={storeAtom}>
        {children}
      </JotaiAtomContext.Provider>
    );
  }

  /**
   * ```js
   * const store = useJotaiStoreWithSelector(); // returns the entire store
   * const selectedValue = useJotaiStoreWithSelector(state => state.selectedValue); // returns the selected value
   * ```
   * */
  function useValue<R>(selector?: (state: T) => R) {
    const storeAtom = useContext(JotaiAtomContext);
    if (!storeAtom)
      throw new Error('useJotaiStore must be used within provider');

    // Create a derived atom for the selector
    const derivedAtom = useMemo(
      () =>
        atom((get) => (selector ? selector(get(storeAtom)) : get(storeAtom))),
      [selector, storeAtom]
    );

    // Use the derived atom
    return useAtomValue(derivedAtom);
  }

  /**
   * ```js
   * const setStore = useSetJotaiStore();
   * setStore((prev) => ({ ...prev, bears: [...prev.bears, 'bear-' + Date.now()] }));
   * ```
   * */
  function useSetValue() {
    const storeAtom = useContext(JotaiAtomContext);
    if (!storeAtom)
      throw new Error('useSetJotaiStore must be used within provider');

    return useSetAtom(storeAtom);
  }

  /**
   * ```js
   * const subscribe = useSubscribe();
   * useEffect(() => {
   *   const unsub = subscribe((value) => {
   *     console.log('Store changed', value);
   *   });
   *   return unsub;
   * }, []);
   * ```
   * */
  function useSubscribe() {
    const storeAtom = useContext(JotaiAtomContext);
    if (!storeAtom)
      throw new Error('useJotaiStore must be used within provider');

    return (callback: (value: T) => void) => {
      const unsub = jotaiStore.sub(storeAtom, () => {
        callback(jotaiStore.get(storeAtom));
      });
      return unsub;
    };
  }

  function useGetValue() {
    const storeAtom = useContext(JotaiAtomContext);
    if (!storeAtom)
      throw new Error('useJotaiStore must be used within provider');
    return () => jotaiStore.get(storeAtom);
  }

  return { StoreProvider, useValue, useSetValue, useSubscribe, useGetValue };
}

export type UseMultiRefs<E> = FunctionConstructor & {
  current: E | null;
};

/**
 * This is a hook that allows you to pass multiple refs to a component.
 * When you want to use `<Component ref={setRef} />`
 * Usage:
 * ```tsx
 * const ref = useMultiRefs<HTMLDivElement>(ref1, ref2, ref3, ...refs);
 * <Component ref={ref} />
 * const ele = ref.current
 * ```
 * */
export function useMultiRefs<E>(
  ...forwardRefs: (React.RefObject<E> | React.ForwardedRef<E>)[]
) {
  const originalRef = useRef<E | null>(null);

  const setRef = (value: E) => {
    originalRef.current = value;
    for (const ref of forwardRefs) {
      if (typeof ref === 'function') {
        ref(value);
      } else if (ref) {
        (ref as React.MutableRefObject<E | null>).current = value;
      }
    }
  };

  setRef.current = originalRef.current;

  return setRef;
}