Hide { fn: ... } workaround behind the scenes

packages/jotai-x/src/createAtomStore.spec.tsx

@@ -12,6 +12,7 @@ describe('createAtomStore', () => {
     type MyTestStoreValue = {
       name: string;
       age: number;
+      becomeFriends: () => void;
     };
 
     const INITIAL_NAME = 'John';
@@ -20,12 +21,11 @@ describe('createAtomStore', () => {
     const initialTestStoreValue: MyTestStoreValue = {
       name: INITIAL_NAME,
       age: INITIAL_AGE,
+      becomeFriends: () => {},
     };
 
-    const { useMyTestStoreStore, MyTestStoreProvider } = createAtomStore(
-      initialTestStoreValue,
-      { name: 'myTestStore' as const }
-    );
+    const { myTestStoreStore, useMyTestStoreStore, MyTestStoreProvider } =
+      createAtomStore(initialTestStoreValue, { name: 'myTestStore' as const });
 
     const ReadOnlyConsumer = () => {
       const name = useMyTestStoreStore().get.name();
@@ -71,6 +71,76 @@ describe('createAtomStore', () => {
       );
     };
 
+    const BecomeFriendsProvider = ({ children }: { children: ReactNode }) => {
+      const [becameFriends, setBecameFriends] = useState(false);
+
+      return (
+        <>
+          <MyTestStoreProvider becomeFriends={() => setBecameFriends(true)}>
+            {children}
+          </MyTestStoreProvider>
+
+          <div>becameFriends: {becameFriends.toString()}</div>
+        </>
+      );
+    };
+
+    const BecomeFriendsGetter = () => {
+      // Make sure both of these are actual functions, not wrapped functions
+      const becomeFriends1 = useMyTestStoreStore().get.becomeFriends();
+      const becomeFriends2 = useMyTestStoreStore().get.atom(
+        myTestStoreStore.atom.becomeFriends
+      );
+
+      return (
+        <button
+          type="button"
+          onClick={() => {
+            becomeFriends1();
+            becomeFriends2();
+          }}
+        >
+          Become Friends
+        </button>
+      );
+    };
+
+    const BecomeFriendsSetter = () => {
+      const setBecomeFriends = useMyTestStoreStore().set.becomeFriends();
+      const [becameFriends, setBecameFriends] = useState(false);
+
+      return (
+        <>
+          <button
+            type="button"
+            onClick={() => setBecomeFriends(() => setBecameFriends(true))}
+          >
+            Change Callback
+          </button>
+
+          <div>setterBecameFriends: {becameFriends.toString()}</div>
+        </>
+      );
+    };
+
+    const BecomeFriendsUser = () => {
+      const [, setBecomeFriends] = useMyTestStoreStore().use.becomeFriends();
+      const [becameFriends, setBecameFriends] = useState(false);
+
+      return (
+        <>
+          <button
+            type="button"
+            onClick={() => setBecomeFriends(() => setBecameFriends(true))}
+          >
+            Change Callback
+          </button>
+
+          <div>userBecameFriends: {becameFriends.toString()}</div>
+        </>
+      );
+    };
+
     beforeEach(() => {
       renderHook(() => useMyTestStoreStore().set.name()(INITIAL_NAME));
       renderHook(() => useMyTestStoreStore().set.age()(INITIAL_AGE));
@@ -157,6 +227,46 @@ describe('createAtomStore', () => {
       expect(getByText(INITIAL_NAME)).toBeInTheDocument();
       expect(getByText(WRITE_ONLY_CONSUMER_AGE)).toBeInTheDocument();
     });
+
+    it('provides and gets functions', () => {
+      const { getByText } = render(
+        <BecomeFriendsProvider>
+          <BecomeFriendsGetter />
+        </BecomeFriendsProvider>
+      );
+
+      expect(getByText('becameFriends: false')).toBeInTheDocument();
+      act(() => getByText('Become Friends').click());
+      expect(getByText('becameFriends: true')).toBeInTheDocument();
+    });
+
+    it('sets functions', () => {
+      const { getByText } = render(
+        <BecomeFriendsProvider>
+          <BecomeFriendsSetter />
+          <BecomeFriendsGetter />
+        </BecomeFriendsProvider>
+      );
+
+      act(() => getByText('Change Callback').click());
+      expect(getByText('setterBecameFriends: false')).toBeInTheDocument();
+      act(() => getByText('Become Friends').click());
+      expect(getByText('setterBecameFriends: true')).toBeInTheDocument();
+    });
+
+    it('uses functions', () => {
+      const { getByText } = render(
+        <BecomeFriendsProvider>
+          <BecomeFriendsUser />
+          <BecomeFriendsGetter />
+        </BecomeFriendsProvider>
+      );
+
+      act(() => getByText('Change Callback').click());
+      expect(getByText('userBecameFriends: false')).toBeInTheDocument();
+      act(() => getByText('Become Friends').click());
+      expect(getByText('userBecameFriends: true')).toBeInTheDocument();
+    });
   });
 
   describe('scoped providers', () => {

packages/jotai-x/src/createAtomStore.ts

@@ -17,6 +17,30 @@ export type UseAtomOptions = {
 
 type UseAtomOptionsOrScope = UseAtomOptions | string;
 
+// Jotai does not support functions in atoms, so wrap functions in objects
+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;
+
+const atomWithFn = <T>(initialValue: T): SimpleWritableAtom<T> => {
+  const baseAtom = atom(wrapFn(initialValue));
+
+  return atom(
+    (get) => unwrapFn(get(baseAtom)) as T,
+    (_get, set, value) => set(baseAtom, wrapFn(value))
+  );
+};
+
 type GetRecord<O> = {
   [K in keyof O]: O[K] extends Atom<infer V>
     ? (options?: UseAtomOptionsOrScope) => V
@@ -195,7 +219,7 @@ export const createAtomStore = <
   for (const [key, atomOrValue] of Object.entries(initialState)) {
     const atomConfig: Atom<unknown> = isAtom(atomOrValue)
       ? atomOrValue
-      : atom(atomOrValue);
+      : atomWithFn(atomOrValue);
     atomsWithoutExtend[key as keyof MyStoreAtomsWithoutExtend] =
       atomConfig as any;
 

packages/jotai-x/src/useHydrateStore.ts

@@ -2,7 +2,7 @@ import { useEffect } from 'react';
 import { useSetAtom } from 'jotai';
 import { useHydrateAtoms } from 'jotai/utils';
 
-import type {
+import {
   SimpleWritableAtomRecord,
   UseHydrateAtoms,
   UseSyncAtoms,
@@ -22,12 +22,7 @@ export const useHydrateStore = (
     const initialValue = initialValues[key];
 
     if (initialValue !== undefined) {
-      values.push([
-        atoms[key],
-        typeof initialValue === 'function'
-          ? { fn: initialValue }
-          : initialValue,
-      ]);
+      values.push([atoms[key], initialValue]);
     }
   }
 
@@ -50,7 +45,7 @@ export const useSyncStore = (
 
     useEffect(() => {
       if (value !== undefined && value !== null) {
-        set(typeof value === 'function' ? { fn: value } : value);
+        set(value);
       }
     }, [set, value]);
   }