Merge pull request #11 from jotaijs/feat/refactoring-core-functionality

Refactoring Core Functionality

src/hydrate-atoms/HydrationBoundary.tsx

@@ -0,0 +1,24 @@
+'use client';
+
+import type { PropsWithChildren } from 'react';
+import type {
+  AnyWritableAtom,
+  HydrateAtomOptions,
+  InferAtomTuples,
+} from './types.js';
+import { useHydrateAtoms } from './use-hydrate-atoms.js';
+
+export function HydrationBoundary<
+  T extends (readonly [AnyWritableAtom, ...unknown[]])[],
+>({
+  children,
+  hydrateAtoms,
+  options,
+}: PropsWithChildren<{
+  hydrateAtoms: InferAtomTuples<T>;
+  options?: HydrateAtomOptions | undefined;
+}>) {
+  useHydrateAtoms(hydrateAtoms, options);
+
+  return <>{children}</>;
+}

src/hydrate-atoms/index.ts

@@ -0,0 +1,2 @@
+export * from './use-hydrate-atoms.js';
+export * from './HydrationBoundary.js';

src/hydrate-atoms/types.ts

@@ -0,0 +1,18 @@
+import type { createStore, WritableAtom } from 'jotai';
+
+export type Store = ReturnType<typeof createStore>;
+
+export type AnyWritableAtom = WritableAtom<unknown, never[], unknown>;
+
+export type InferAtomTuples<T> = {
+  [K in keyof T]: T[K] extends readonly [infer A, ...unknown[]]
+    ? A extends WritableAtom<unknown, infer Args, unknown>
+      ? readonly [A, ...Args]
+      : T[K]
+    : never;
+};
+
+export type HydrateAtomOptions = {
+  store?: Store | undefined;
+  enableReHydrate?: boolean | undefined;
+};

src/hydrate-atoms/use-hydrate-atoms.ts

@@ -0,0 +1,44 @@
+import { atom, useStore } from 'jotai';
+import type {
+  AnyWritableAtom,
+  HydrateAtomOptions,
+  InferAtomTuples,
+} from './types.js';
+import { useCallback, useEffect, useRef } from 'react';
+
+export function useHydrateAtoms<
+  T extends (readonly [AnyWritableAtom, ...unknown[]])[],
+>(hydrateAtoms: InferAtomTuples<T>, options?: HydrateAtomOptions) {
+  const isHydratedRef = useRef(false);
+  const lastHydrateAtoms = useRef(hydrateAtoms);
+  const store = useStore(
+    options?.store != null ? { store: options?.store } : undefined,
+  );
+
+  const hydrate = useCallback(() => {
+    store.set(isHydratingPrimitiveAtom, true);
+    for (const [atom, ...args] of hydrateAtoms) {
+      store.set(atom, ...args);
+    }
+    store.set(isHydratingPrimitiveAtom, false);
+    isHydratedRef.current = true;
+  }, [store, hydrateAtoms]);
+
+  if (!isHydratedRef.current) {
+    hydrate();
+  }
+
+  useEffect(() => {
+    if (!options?.enableReHydrate) {
+      return;
+    }
+    if (hydrateAtoms !== lastHydrateAtoms.current) {
+      lastHydrateAtoms.current = hydrateAtoms;
+      hydrate();
+    }
+  }, [hydrate, options?.enableReHydrate, hydrateAtoms]);
+}
+
+const isHydratingPrimitiveAtom = atom(false);
+
+export const isHydratingAtom = atom((get) => get(isHydratingPrimitiveAtom));

src/index.ts

@@ -1,3 +1 @@
-export * from './HydrationBoundary.js';
-export * from './RenderingBoundary.js';
-export * from './SuspenseBoundary.js';
+export * from './hydrate-atoms/index.js';