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,17 @@
+import type { WritableAtom } from 'jotai';
+import type { Store } from '../shared.js';
+
+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,2 @@
-export * from './HydrationBoundary.js';
-export * from './RenderingBoundary.js';
-export * from './SuspenseBoundary.js';
+export * from './hydrate-atoms/index.js';
+export * from './ssr-provider/index.js';

src/shared.ts

@@ -0,0 +1,3 @@
+import type { createStore } from 'jotai';
+
+export type Store = ReturnType<typeof createStore>;

src/ssr-provider/SSRProvider.tsx

@@ -0,0 +1,33 @@
+'use client';
+
+import { Provider } from 'jotai';
+import type { PropsWithChildren } from 'react';
+import { useCreateStore } from './use-create-store.js';
+import type { Store } from '../shared.js';
+
+function SSRProviderWithGivenStore({
+  store,
+  children,
+}: PropsWithChildren<{ store: Store }>) {
+  return <Provider store={store}>{children}</Provider>;
+}
+
+function SSRProviderWithNewStore({ children }: PropsWithChildren) {
+  const store = useCreateStore();
+  return <Provider store={store}>{children}</Provider>;
+}
+
+export function SSRProvider({
+  store,
+  children,
+}: PropsWithChildren<{ store?: Store }>) {
+  if (store) {
+    return (
+      <SSRProviderWithGivenStore store={store}>
+        {children}
+      </SSRProviderWithGivenStore>
+    );
+  } else {
+    return <SSRProviderWithNewStore>{children}</SSRProviderWithNewStore>;
+  }
+}

src/ssr-provider/index.ts

@@ -0,0 +1,2 @@
+export * from './use-create-store.js';
+export * from './SSRProvider.js';

src/ssr-provider/use-create-store.ts

@@ -0,0 +1,7 @@
+import { createStore } from 'jotai';
+import { useState } from 'react';
+
+export const useCreateStore = () => {
+  const [store] = useState(() => createStore());
+  return store;
+};