feat(jotai): Added createAtomStore functionality

packages/jotai/package.json

@@ -10,12 +10,16 @@
   },
   "type": "module",
   "devDependencies": {
+    "@types/react": "18.2.43",
+    "@types/react-dom": "18.2.17",
     "jotai": "^2.6.0",
-    "react": "^18.2.0"
+    "react": "^18.2.0",
+    "react-dom": "18.2.0"
   },
   "peerDependencies": {
     "jotai": "^2.6.0",
-    "react": "^18.2.0"
+    "react": "^18.2.0",
+    "react-dom": "18.2.0"
   },
   "publishConfig": {
     "access": "public"

packages/jotai/project.json

@@ -5,14 +5,15 @@
   "sourceRoot": "packages/jotai/src",
   "targets": {
     "build": {
-      "executor": "@storm-software/workspace-tools:tsup-neutral",
+      "executor": "@storm-software/workspace-tools:tsup-browser",
       "outputs": ["{options.outputPath}"],
       "options": {
         "entry": "packages/jotai/src/index.ts",
         "outputPath": "dist/packages/jotai",
         "tsConfig": "packages/jotai/tsconfig.json",
         "project": "packages/jotai/package.json",
         "defaultConfiguration": "production",
+        "platform": "browser",
         "assets": [
           {
             "input": "packages/jotai",
@@ -24,8 +25,7 @@
             "glob": "LICENSE",
             "output": "/"
           }
-        ],
-        "platform": "neutral"
+        ]
       },
       "configurations": {
         "production": {

packages/jotai/src/atoms/atom-with-wrapper.ts

@@ -0,0 +1,40 @@
+import { atom } from "jotai";
+
+import type { WritableAtom } from "jotai/vanilla";
+
+type WrapFn<T> = T extends (...args: infer _A) => infer _R ? { __fn: T } : T;
+
+const wrapFn = <T>(fnOrValue: T): WrapFn<T> =>
+  (typeof fnOrValue === "function" ? { __fn: fnOrValue } : fnOrValue) as any;
+
+type UnwrapFn<T> = T extends { __fn: infer U } ? U : T;
+
+const unwrapFn = <T>(wrappedFnOrValue: T): UnwrapFn<T> =>
+  (wrappedFnOrValue &&
+  typeof wrappedFnOrValue === "object" &&
+  "__fn" in wrappedFnOrValue
+    ? wrappedFnOrValue.__fn
+    : wrappedFnOrValue) as any;
+
+/**
+ * Create an atom with a wrapper that allows functions as values.
+ *
+ * @remarks
+ * Jotai atoms don't allow functions as values by default. This function is a
+ * drop-in replacement for `atom` that wraps functions in an object while
+ * leaving non-functions unchanged. The wrapper object should be completely
+ * invisible to consumers of the atom.
+ *
+ * @param initialValue - The initial value of the atom
+ * @returns An atom with a wrapper that allows functions as values.
+ */
+export const atomWithWrapper = <TValue>(
+  initialValue: TValue
+): WritableAtom<TValue, [TValue], void> => {
+  const baseAtom = atom(wrapFn(initialValue));
+
+  return atom(
+    get => unwrapFn(get(baseAtom)) as TValue,
+    (_get, set, value) => set(baseAtom, wrapFn(value))
+  );
+};

packages/jotai/src/atoms/index.ts

@@ -2,3 +2,4 @@ export * from "./atom-with-broadcast";
 export * from "./atom-with-effect";
 export * from "./atom-with-pending";
 export * from "./atom-with-web-storage";
+export * from "./atom-with-wrapper";

packages/jotai/src/hooks/index.ts

@@ -1 +1,3 @@
+export * from "./use-hydrate-store";
 export * from "./use-prepare-atoms";
+export * from "./use-sync-store";

packages/jotai/src/hooks/use-hydrate-store.ts

@@ -0,0 +1,26 @@
+import { useHydrateAtoms } from "jotai/utils";
+import {
+  SimpleWritableAtomRecord,
+  UseHydrateAtoms
+} from "../utilities/create-atom-store";
+
+/**
+ * Hydrate atoms with initial values for SSR.
+ */
+export const useHydrateStore = (
+  atoms: SimpleWritableAtomRecord<any>,
+  initialValues: Parameters<UseHydrateAtoms<any>>[0],
+  options: Parameters<UseHydrateAtoms<any>>[1] = {}
+) => {
+  const values: any[] = [];
+
+  for (const key of Object.keys(atoms)) {
+    const initialValue = initialValues[key];
+
+    if (initialValue !== undefined) {
+      values.push([atoms[key], initialValue]);
+    }
+  }
+
+  useHydrateAtoms(values, options);
+};

packages/jotai/src/hooks/use-sync-store.ts

@@ -0,0 +1,26 @@
+import { useSetAtom } from "jotai";
+import { useEffect } from "react";
+import {
+  SimpleWritableAtomRecord,
+  UseSyncAtoms
+} from "../utilities/create-atom-store";
+
+/**
+ * Update atoms with new values on changes.
+ */
+export const useSyncStore = (
+  atoms: SimpleWritableAtomRecord<any>,
+  values: any,
+  { store }: Parameters<UseSyncAtoms<any>>[1] = {}
+) => {
+  for (const key of Object.keys(atoms)) {
+    const value = values[key];
+    const set = useSetAtom(atoms[key]!, { store });
+
+    useEffect(() => {
+      if (value !== undefined && value !== null) {
+        set(value);
+      }
+    }, [set, value]);
+  }
+};

packages/jotai/src/utilities/create-atom-provider.tsx

@@ -0,0 +1,135 @@
+import { Provider as AtomProvider } from "jotai";
+import { createStore } from "jotai/vanilla";
+import {
+  ComponentProps,
+  createContext,
+  FC,
+  useContext,
+  useEffect,
+  useMemo,
+  useState
+} from "react";
+import { useHydrateStore, useSyncStore } from "../hooks";
+import { JotaiStore, SimpleWritableAtomRecord } from "./create-atom-store";
+
+type AtomProviderProps = ComponentProps<typeof AtomProvider>;
+
+const getFullyQualifiedScope = (storeName: string, scope: string) => {
+  return `${storeName}:${scope}`;
+};
+
+/**
+ * Context mapping store name and scope to store. The 'provider' scope is used
+ * to reference any provider belonging to the store, regardless of scope.
+ */
+const PROVIDER_SCOPE = "provider";
+const AtomStoreContext = createContext<Map<string, JotaiStore>>(new Map());
+
+/**
+ * Tries to find a store in each of the following places, in order:
+ * 1. The store context, matching the store name and scope
+ * 2. The store context, matching the store name and 'provider' scope
+ * 3. Otherwise, return undefined
+ */
+export const useAtomStore = (
+  storeName: string,
+  scope: string = PROVIDER_SCOPE,
+  warnIfUndefined: boolean = true
+): JotaiStore | undefined => {
+  const storeContext = useContext(AtomStoreContext);
+  const store =
+    storeContext.get(getFullyQualifiedScope(storeName, scope)) ??
+    storeContext.get(getFullyQualifiedScope(storeName, PROVIDER_SCOPE));
+
+  if (!store && warnIfUndefined) {
+    console.warn(
+      `Tried to access jotai store '${storeName}' outside of a matching provider.`
+    );
+  }
+
+  return store;
+};
+
+export type ProviderProps<T extends object> = AtomProviderProps &
+  Partial<T> & {
+    scope?: string;
+    initialValues?: Partial<T>;
+    resetKey?: any;
+  };
+
+export const HydrateAtoms = <T extends object>({
+  initialValues,
+  children,
+  store,
+  atoms,
+  ...props
+}: Omit<ProviderProps<T>, "scope"> & {
+  atoms: SimpleWritableAtomRecord<T>;
+}) => {
+  useHydrateStore(atoms, { ...initialValues, ...props } as any, {
+    store
+  });
+  useSyncStore(atoms, props as any, {
+    store
+  });
+
+  return <>{children}</>;
+};
+
+/**
+ * Creates a generic provider for a jotai store.
+ * - `initialValues`: Initial values for the store.
+ * - `props`: Dynamic values for the store.
+ */
+export const createAtomProvider = <T extends object, N extends string = "">(
+  storeScope: N,
+  atoms: SimpleWritableAtomRecord<T>,
+  options: { effect?: FC } = {}
+) => {
+  const Effect = options.effect;
+
+  // eslint-disable-next-line react/display-name
+  return ({ store, scope, children, resetKey, ...props }: ProviderProps<T>) => {
+    const [storeState, setStoreState] = useState<JotaiStore>(createStore());
+
+    useEffect(() => {
+      if (resetKey) {
+        setStoreState(createStore());
+      }
+    }, [resetKey]);
+
+    const previousStoreContext = useContext(AtomStoreContext);
+
+    const storeContext = useMemo(() => {
+      const newStoreContext = new Map(previousStoreContext);
+
+      if (scope) {
+        // Make the store findable by its fully qualified scope
+        newStoreContext.set(
+          getFullyQualifiedScope(storeScope, scope),
+          storeState
+        );
+      }
+
+      // Make the store findable by its store name alone
+      newStoreContext.set(
+        getFullyQualifiedScope(storeScope, PROVIDER_SCOPE),
+        storeState
+      );
+
+      return newStoreContext;
+    }, [previousStoreContext, scope, storeState]);
+
+    return (
+      <AtomStoreContext.Provider value={storeContext}>
+        <AtomProvider store={storeState}>
+          <HydrateAtoms store={storeState} atoms={atoms} {...(props as any)}>
+            {!!Effect && <Effect />}
+
+            {children}
+          </HydrateAtoms>
+        </AtomProvider>
+      </AtomStoreContext.Provider>
+    );
+  };
+};