From 347bcbe1acab4c6e7f9ca580f07e8883cce978a6 Mon Sep 17 00:00:00 2001
From: Joe Anderson <joe@mousetrapped.co.uk>
Date: Wed, 20 Dec 2023 19:48:35 +0000
Subject: [PATCH 01/20] Add extend option

---
 packages/jotai-x/src/createAtomProvider.tsx   |   6 +-
 packages/jotai-x/src/createAtomStore.spec.tsx |  71 ++++++++-
 packages/jotai-x/src/createAtomStore.ts       | 149 ++++++++----------
 packages/jotai-x/src/useHydrateStore.ts       |   6 +-
 4 files changed, 140 insertions(+), 92 deletions(-)

diff --git a/packages/jotai-x/src/createAtomProvider.tsx b/packages/jotai-x/src/createAtomProvider.tsx
index 5412023..f4930a8 100644
--- a/packages/jotai-x/src/createAtomProvider.tsx
+++ b/packages/jotai-x/src/createAtomProvider.tsx
@@ -9,7 +9,7 @@ import React, {
 import { createStore } from 'jotai/vanilla';
 
 import { AtomProvider, AtomProviderProps } from './atomProvider';
-import { AtomRecord, JotaiStore } from './createAtomStore';
+import { JotaiStore, PrimitiveAtomRecord } from './createAtomStore';
 import { useHydrateStore, useSyncStore } from './useHydrateStore';
 
 const getFullyQualifiedScope = (storeName: string, scope: string) => {
@@ -62,7 +62,7 @@ export const HydrateAtoms = <T extends object>({
   atoms,
   ...props
 }: Omit<ProviderProps<T>, 'scope'> & {
-  atoms: AtomRecord<T>;
+  atoms: PrimitiveAtomRecord<T>;
 }) => {
   useHydrateStore(atoms, { ...initialValues, ...props } as any, {
     store,
@@ -81,7 +81,7 @@ export const HydrateAtoms = <T extends object>({
  */
 export const createAtomProvider = <T extends object, N extends string = ''>(
   storeScope: N,
-  atoms: AtomRecord<T>,
+  atoms: PrimitiveAtomRecord<T>,
   options: { effect?: FC } = {}
 ) => {
   const Effect = options.effect;
diff --git a/packages/jotai-x/src/createAtomStore.spec.tsx b/packages/jotai-x/src/createAtomStore.spec.tsx
index 9d911f2..c68e692 100644
--- a/packages/jotai-x/src/createAtomStore.spec.tsx
+++ b/packages/jotai-x/src/createAtomStore.spec.tsx
@@ -2,6 +2,7 @@ import '@testing-library/jest-dom';
 
 import React, { ReactNode, useState } from 'react';
 import { act, render, renderHook } from '@testing-library/react';
+import { atom, useAtomValue } from 'jotai';
 
 import { createAtomStore } from './createAtomStore';
 
@@ -26,8 +27,8 @@ describe('createAtomStore', () => {
     );
 
     const ReadOnlyConsumer = () => {
-      const [name] = useMyTestStoreStore().use.name();
-      const [age] = useMyTestStoreStore().use.age();
+      const name = useMyTestStoreStore().get.name();
+      const age = useMyTestStoreStore().get.age();
 
       return (
         <div>
@@ -170,7 +171,7 @@ describe('createAtomStore', () => {
       });
 
     const ReadOnlyConsumer = ({ scope }: { scope: string }) => {
-      const [age] = useMyScopedTestStoreStore().use.age({ scope });
+      const age = useMyScopedTestStoreStore().get.age({ scope });
 
       return (
         <div>
@@ -184,7 +185,7 @@ describe('createAtomStore', () => {
     }: {
       scope: string;
     }) => {
-      const [age] = useMyScopedTestStoreStore(scope).use.age();
+      const age = useMyScopedTestStoreStore(scope).get.age();
 
       return (
         <div>
@@ -269,7 +270,7 @@ describe('createAtomStore', () => {
       });
 
     const FirstReadOnlyConsumer = () => {
-      const [name] = useMyFirstTestStoreStore().use.name();
+      const name = useMyFirstTestStoreStore().get.name();
 
       return (
         <div>
@@ -279,7 +280,7 @@ describe('createAtomStore', () => {
     };
 
     const SecondReadOnlyConsumer = () => {
-      const [age] = useMySecondTestStoreStore().use.age();
+      const age = useMySecondTestStoreStore().get.age();
 
       return (
         <div>
@@ -302,4 +303,62 @@ describe('createAtomStore', () => {
       expect(getByText('98')).toBeInTheDocument();
     });
   });
+
+  describe('extended stores', () => {
+    type User = {
+      name: string;
+      age: number;
+    };
+
+    const initialUser: User = {
+      name: 'Jane',
+      age: 98,
+    };
+
+    const { userStore, useUserStore, UserProvider } = createAtomStore(
+      initialUser,
+      {
+        name: 'user' as const,
+        extend: ({ name, age }) => ({
+          bio: atom((get) => `${get(name)} is ${get(age)} years old`),
+        }),
+      }
+    );
+
+    const ReadOnlyConsumer = () => {
+      const bio = useUserStore().get.bio();
+
+      return <div>{bio}</div>;
+    };
+
+    it('includes extended atom in store object', () => {
+      const { result } = renderHook(() => useAtomValue(userStore.atom.bio));
+      expect(result.current).toBe('Jane is 98 years old');
+    });
+
+    it('includes extended atom in get hooks', () => {
+      const { result } = renderHook(() => useUserStore().get.bio());
+      expect(result.current).toBe('Jane is 98 years old');
+    });
+
+    it('does not include extended atom in set hooks', () => {
+      const { result } = renderHook(() => Object.keys(useUserStore().set));
+      expect(result.current).not.toContain('bio');
+    });
+
+    it('does not include extended atom in use hooks', () => {
+      const { result } = renderHook(() => Object.keys(useUserStore().use));
+      expect(result.current).not.toContain('bio');
+    });
+
+    it('computes extended atom based on current state', () => {
+      const { getByText } = render(
+        <UserProvider name="John" age={42}>
+          <ReadOnlyConsumer />
+        </UserProvider>
+      );
+
+      expect(getByText('John is 42 years old')).toBeInTheDocument();
+    });
+  });
 });
diff --git a/packages/jotai-x/src/createAtomStore.ts b/packages/jotai-x/src/createAtomStore.ts
index fa675c8..a63f242 100644
--- a/packages/jotai-x/src/createAtomStore.ts
+++ b/packages/jotai-x/src/createAtomStore.ts
@@ -4,14 +4,9 @@ import { createAtomProvider, useAtomStore } from './createAtomProvider';
 
 import type { ProviderProps } from './createAtomProvider';
 import type { FC } from 'react';
-import type { PrimitiveAtom } from 'jotai';
 import type { useHydrateAtoms } from 'jotai/utils';
-import type { createStore } from 'jotai/vanilla';
+import type { Atom, createStore, PrimitiveAtom } from 'jotai/vanilla';
 
-type WithInitialValue<Value> = {
-  init: Value;
-};
-type Atom<Value> = PrimitiveAtom<Value> & WithInitialValue<Value>;
 export type JotaiStore = ReturnType<typeof createStore>;
 
 export type UseAtomOptions = {
@@ -33,6 +28,9 @@ export type UseRecord<O> = {
     options?: UseAtomOptionsOrScope
   ) => [O[K], (value: O[K]) => void];
 };
+export type PrimitiveAtomRecord<O> = {
+  [K in keyof O]: PrimitiveAtom<O[K]>;
+};
 export type AtomRecord<O> = {
   [K in keyof O]: Atom<O[K]>;
 };
@@ -51,32 +49,37 @@ export type UseSyncAtoms<T> = (
   }
 ) => void;
 
-export type StoreApi<T extends object, N extends string = ''> = {
-  atom: AtomRecord<T>;
+export type StoreApi<
+  T extends object,
+  E extends AtomRecord<object>,
+  N extends string = '',
+> = {
+  atom: PrimitiveAtomRecord<T> & E;
   name: N;
-  extend: <ET, EN>(
-    extendedState: ET,
-    options?: Omit<
-      CreateAtomStoreOptions<object, EN extends string ? EN : N>,
-      'initialStore'
-    >
-  ) => AtomStoreApi<T & ET, EN extends string ? EN : N>;
 };
 
-export type UseStoreApi<T> = (options?: UseAtomOptionsOrScope) => {
-  get: GetRecord<T>;
+export type UseStoreApi<T, E> = (options?: UseAtomOptionsOrScope) => {
+  get: GetRecord<
+    T & {
+      [key in keyof E]: E[key] extends Atom<infer U> ? U : never;
+    }
+  >;
   set: SetRecord<T>;
   use: UseRecord<T>;
 };
 
-export type AtomStoreApi<T extends object, N extends string = ''> = {
+export type AtomStoreApi<
+  T extends object,
+  E extends AtomRecord<object>,
+  N extends string = '',
+> = {
   name: N;
 } & {
   [key in keyof Record<NameProvider<N>, object>]: FC<ProviderProps<T>>;
 } & {
-  [key in keyof Record<NameStore<N>, object>]: StoreApi<T, N>;
+  [key in keyof Record<NameStore<N>, object>]: StoreApi<T, E, N>;
 } & {
-  [key in keyof Record<UseNameStore<N>, object>]: UseStoreApi<T>;
+  [key in keyof Record<UseNameStore<N>, object>]: UseStoreApi<T, E>;
 };
 
 const capitalizeFirstLetter = (str = '') =>
@@ -88,12 +91,12 @@ const getStoreIndex = (name = '') =>
 const getUseStoreIndex = (name = '') =>
   `use${capitalizeFirstLetter(name)}Store`;
 
-const withDefaultOptions = <T>(
-  atomRecord: AtomRecord<T>,
+const withDefaultOptions = <T, R>(
+  fnRecord: { [key in keyof T]: (options?: UseAtomOptions) => R },
   defaultOptions: UseAtomOptions
-): AtomRecord<T> =>
+): typeof fnRecord =>
   Object.fromEntries(
-    Object.entries(atomRecord).map(([key, fn]) => [
+    Object.entries(fnRecord).map(([key, fn]) => [
       key,
       (options: UseAtomOptions = {}) =>
         (fn as any)({ ...defaultOptions, ...options }),
@@ -107,12 +110,16 @@ const convertScopeShorthand = (
     ? { scope: optionsOrScope }
     : optionsOrScope;
 
-export interface CreateAtomStoreOptions<T extends object, N extends string> {
+export interface CreateAtomStoreOptions<
+  T extends object,
+  E extends AtomRecord<object>,
+  N extends string,
+> {
+  name: N;
   store?: UseAtomOptions['store'];
   delay?: UseAtomOptions['delay'];
-  initialStore?: AtomStoreApi<T, N>;
-  name?: N;
   effect?: FC;
+  extend?: (primitiveAtoms: PrimitiveAtomRecord<T>) => E;
 }
 
 /**
@@ -129,94 +136,76 @@ export interface CreateAtomStoreOptions<T extends object, N extends string> {
  */
 export const createAtomStore = <
   T extends object,
-  IT extends object,
+  E extends AtomRecord<object>,
   N extends string = '',
 >(
   initialState: T,
-  {
-    delay: delayRoot,
-    initialStore,
-    name = '' as any,
-    effect,
-  }: CreateAtomStoreOptions<IT, N> = {}
-): AtomStoreApi<T & IT, N> => {
-  const useInitialStoreIndex = getUseStoreIndex(
-    initialStore?.name
-  ) as UseNameStore<N>;
-  const initialStoreIndex = getStoreIndex(initialStore?.name) as NameStore<N>;
+  { name, delay: delayRoot, effect, extend }: CreateAtomStoreOptions<T, E, N>
+): AtomStoreApi<T, E, N> => {
   const providerIndex = getProviderIndex(name) as NameProvider<N>;
   const useStoreIndex = getUseStoreIndex(name) as UseNameStore<N>;
   const storeIndex = getStoreIndex(name) as NameStore<N>;
 
-  // FIXME: These constants have type any
-  const getAtoms = initialStore
-    ? (initialStore[useInitialStoreIndex] as any)().get
-    : ({} as GetRecord<T & IT>);
-  const setAtoms = initialStore
-    ? (initialStore[useInitialStoreIndex] as any)().set
-    : ({} as SetRecord<T & IT>);
-  const useAtoms = initialStore
-    ? (initialStore[useInitialStoreIndex] as any)().use
-    : ({} as UseRecord<T & IT>);
-  const atoms = initialStore
-    ? (initialStore[initialStoreIndex] as any).atom
-    : ({} as AtomRecord<T & IT>);
+  const getAtoms = {} as ReturnType<UseStoreApi<T, E>>['get'];
+  const setAtoms = {} as ReturnType<UseStoreApi<T, E>>['set'];
+  const useAtoms = {} as ReturnType<UseStoreApi<T, E>>['use'];
+  const primitiveAtoms = {} as PrimitiveAtomRecord<T>;
 
   for (const key of Object.keys(initialState)) {
     const atomConfig = atom(initialState[key as keyof T]);
 
-    atoms[key] = atomConfig;
-    getAtoms[key] = (optionsOrScope: UseAtomOptionsOrScope = {}) => {
+    (primitiveAtoms as any)[key] = atomConfig;
+
+    (setAtoms as any)[key] = (optionsOrScope: UseAtomOptionsOrScope = {}) => {
       const options = convertScopeShorthand(optionsOrScope);
-      const contextStore = useAtomStore(name, options.scope, false);
+      const contextStore = useAtomStore(name, options.scope);
 
-      return useAtomValue(atomConfig, {
+      return useSetAtom(atomConfig as any, {
         store: options.store ?? contextStore,
-        delay: options.delay ?? delayRoot,
       });
     };
-    setAtoms[key] = (optionsOrScope: UseAtomOptionsOrScope = {}) => {
+
+    (useAtoms as any)[key] = (optionsOrScope: UseAtomOptionsOrScope = {}) => {
       const options = convertScopeShorthand(optionsOrScope);
       const contextStore = useAtomStore(name, options.scope);
 
-      return useSetAtom(atomConfig as any, {
+      return useAtom(atomConfig, {
         store: options.store ?? contextStore,
+        delay: options.delay ?? delayRoot,
       });
     };
-    useAtoms[key] = (optionsOrScope: UseAtomOptionsOrScope = {}) => {
+  }
+
+  const atoms = {
+    ...primitiveAtoms,
+    ...(extend ? extend(primitiveAtoms) : {}),
+  } as PrimitiveAtomRecord<T> & E;
+
+  for (const key of Object.keys(atoms)) {
+    const atomConfig = atoms[key as keyof T & keyof E];
+
+    (getAtoms as any)[key] = (optionsOrScope: UseAtomOptionsOrScope = {}) => {
       const options = convertScopeShorthand(optionsOrScope);
-      const contextStore = useAtomStore(name, options.scope);
+      const contextStore = useAtomStore(name, options.scope, false);
 
-      return useAtom(atomConfig, {
+      return useAtomValue(atomConfig, {
         store: options.store ?? contextStore,
         delay: options.delay ?? delayRoot,
       });
     };
   }
 
-  const api: any = {
-    [providerIndex]: createAtomProvider(name, atoms, { effect }),
+  return {
+    [providerIndex]: createAtomProvider(name, primitiveAtoms, { effect }),
     [useStoreIndex]: (options: UseAtomOptionsOrScope = {}) => ({
-      get: withDefaultOptions(getAtoms, convertScopeShorthand(options)),
-      set: withDefaultOptions(setAtoms, convertScopeShorthand(options)),
-      use: withDefaultOptions(useAtoms, convertScopeShorthand(options)),
+      get: withDefaultOptions(getAtoms as any, convertScopeShorthand(options)),
+      set: withDefaultOptions(setAtoms as any, convertScopeShorthand(options)),
+      use: withDefaultOptions(useAtoms as any, convertScopeShorthand(options)),
     }),
     [storeIndex]: {
       atom: atoms,
       name,
     },
     name,
-  };
-
-  return {
-    ...api,
-    [storeIndex]: {
-      ...api[storeIndex],
-      extend: (extendedState: any, options: any) =>
-        createAtomStore(extendedState, {
-          initialStore: api,
-          ...options,
-        }),
-    },
-  };
+  } as any;
 };
diff --git a/packages/jotai-x/src/useHydrateStore.ts b/packages/jotai-x/src/useHydrateStore.ts
index 9dbc595..517ba85 100644
--- a/packages/jotai-x/src/useHydrateStore.ts
+++ b/packages/jotai-x/src/useHydrateStore.ts
@@ -3,7 +3,7 @@ import { useSetAtom } from 'jotai';
 import { useHydrateAtoms } from 'jotai/utils';
 
 import type {
-  AtomRecord,
+  PrimitiveAtomRecord,
   UseHydrateAtoms,
   UseSyncAtoms,
 } from './createAtomStore';
@@ -12,7 +12,7 @@ import type {
  * Hydrate atoms with initial values for SSR.
  */
 export const useHydrateStore = (
-  atoms: AtomRecord<any>,
+  atoms: PrimitiveAtomRecord<any>,
   initialValues: Parameters<UseHydrateAtoms<any>>[0],
   options: Parameters<UseHydrateAtoms<any>>[1] = {}
 ) => {
@@ -38,7 +38,7 @@ export const useHydrateStore = (
  * Update atoms with new values on changes.
  */
 export const useSyncStore = (
-  atoms: AtomRecord<any>,
+  atoms: PrimitiveAtomRecord<any>,
   values: any,
   { store }: Parameters<UseSyncAtoms<any>>[1] = {}
 ) => {

From 3cc7be75b624b549d84a55e8af7316abf97f81cf Mon Sep 17 00:00:00 2001
From: Joe Anderson <joe@mousetrapped.co.uk>
Date: Wed, 20 Dec 2023 20:04:27 +0000
Subject: [PATCH 02/20] Let consumer specify custom atom creator function

---
 packages/jotai-x/src/createAtomProvider.tsx   |  6 ++--
 packages/jotai-x/src/createAtomStore.spec.tsx | 28 ++++++++++++++++-
 packages/jotai-x/src/createAtomStore.ts       | 31 +++++++++++++------
 packages/jotai-x/src/useHydrateStore.ts       |  6 ++--
 4 files changed, 54 insertions(+), 17 deletions(-)

diff --git a/packages/jotai-x/src/createAtomProvider.tsx b/packages/jotai-x/src/createAtomProvider.tsx
index f4930a8..1f95e6c 100644
--- a/packages/jotai-x/src/createAtomProvider.tsx
+++ b/packages/jotai-x/src/createAtomProvider.tsx
@@ -9,7 +9,7 @@ import React, {
 import { createStore } from 'jotai/vanilla';
 
 import { AtomProvider, AtomProviderProps } from './atomProvider';
-import { JotaiStore, PrimitiveAtomRecord } from './createAtomStore';
+import { JotaiStore, WritableAtomRecord } from './createAtomStore';
 import { useHydrateStore, useSyncStore } from './useHydrateStore';
 
 const getFullyQualifiedScope = (storeName: string, scope: string) => {
@@ -62,7 +62,7 @@ export const HydrateAtoms = <T extends object>({
   atoms,
   ...props
 }: Omit<ProviderProps<T>, 'scope'> & {
-  atoms: PrimitiveAtomRecord<T>;
+  atoms: WritableAtomRecord<T>;
 }) => {
   useHydrateStore(atoms, { ...initialValues, ...props } as any, {
     store,
@@ -81,7 +81,7 @@ export const HydrateAtoms = <T extends object>({
  */
 export const createAtomProvider = <T extends object, N extends string = ''>(
   storeScope: N,
-  atoms: PrimitiveAtomRecord<T>,
+  atoms: WritableAtomRecord<T>,
   options: { effect?: FC } = {}
 ) => {
   const Effect = options.effect;
diff --git a/packages/jotai-x/src/createAtomStore.spec.tsx b/packages/jotai-x/src/createAtomStore.spec.tsx
index c68e692..945b8a6 100644
--- a/packages/jotai-x/src/createAtomStore.spec.tsx
+++ b/packages/jotai-x/src/createAtomStore.spec.tsx
@@ -2,7 +2,7 @@ import '@testing-library/jest-dom';
 
 import React, { ReactNode, useState } from 'react';
 import { act, render, renderHook } from '@testing-library/react';
-import { atom, useAtomValue } from 'jotai';
+import { atom, PrimitiveAtom, useAtomValue } from 'jotai';
 
 import { createAtomStore } from './createAtomStore';
 
@@ -361,4 +361,30 @@ describe('createAtomStore', () => {
       expect(getByText('John is 42 years old')).toBeInTheDocument();
     });
   });
+
+  describe('custom createAtom function', () => {
+    type CustomAtom<T> = PrimitiveAtom<T> & {
+      isCustomAtom: true;
+    };
+
+    const createCustomAtom = <T,>(value: T): CustomAtom<T> => ({
+      ...atom(value),
+      isCustomAtom: true,
+    });
+
+    const { customStore } = createAtomStore(
+      {
+        x: 5,
+      },
+      {
+        name: 'custom' as const,
+        createAtom: createCustomAtom,
+      }
+    );
+
+    it('uses custom createAtom function', () => {
+      const myAtom = customStore.atom.x as CustomAtom<number>;
+      expect(myAtom.isCustomAtom).toBe(true);
+    });
+  });
 });
diff --git a/packages/jotai-x/src/createAtomStore.ts b/packages/jotai-x/src/createAtomStore.ts
index a63f242..59bed86 100644
--- a/packages/jotai-x/src/createAtomStore.ts
+++ b/packages/jotai-x/src/createAtomStore.ts
@@ -1,11 +1,11 @@
 import { atom, useAtom, useAtomValue, useSetAtom } from 'jotai';
+import { useHydrateAtoms } from 'jotai/utils';
 
 import { createAtomProvider, useAtomStore } from './createAtomProvider';
 
 import type { ProviderProps } from './createAtomProvider';
 import type { FC } from 'react';
-import type { useHydrateAtoms } from 'jotai/utils';
-import type { Atom, createStore, PrimitiveAtom } from 'jotai/vanilla';
+import type { Atom, createStore, WritableAtom } from 'jotai/vanilla';
 
 export type JotaiStore = ReturnType<typeof createStore>;
 
@@ -28,9 +28,13 @@ export type UseRecord<O> = {
     options?: UseAtomOptionsOrScope
   ) => [O[K], (value: O[K]) => void];
 };
-export type PrimitiveAtomRecord<O> = {
-  [K in keyof O]: PrimitiveAtom<O[K]>;
+
+export type SimpleWritableAtom<T> = WritableAtom<T, [T], void>;
+
+export type WritableAtomRecord<O> = {
+  [K in keyof O]: SimpleWritableAtom<O[K]>;
 };
+
 export type AtomRecord<O> = {
   [K in keyof O]: Atom<O[K]>;
 };
@@ -54,7 +58,7 @@ export type StoreApi<
   E extends AtomRecord<object>,
   N extends string = '',
 > = {
-  atom: PrimitiveAtomRecord<T> & E;
+  atom: WritableAtomRecord<T> & E;
   name: N;
 };
 
@@ -119,7 +123,8 @@ export interface CreateAtomStoreOptions<
   store?: UseAtomOptions['store'];
   delay?: UseAtomOptions['delay'];
   effect?: FC;
-  extend?: (primitiveAtoms: PrimitiveAtomRecord<T>) => E;
+  extend?: (primitiveAtoms: WritableAtomRecord<T>) => E;
+  createAtom?: <V>(value: V) => SimpleWritableAtom<V>;
 }
 
 /**
@@ -140,7 +145,13 @@ export const createAtomStore = <
   N extends string = '',
 >(
   initialState: T,
-  { name, delay: delayRoot, effect, extend }: CreateAtomStoreOptions<T, E, N>
+  {
+    name,
+    delay: delayRoot,
+    effect,
+    extend,
+    createAtom = atom,
+  }: CreateAtomStoreOptions<T, E, N>
 ): AtomStoreApi<T, E, N> => {
   const providerIndex = getProviderIndex(name) as NameProvider<N>;
   const useStoreIndex = getUseStoreIndex(name) as UseNameStore<N>;
@@ -149,10 +160,10 @@ export const createAtomStore = <
   const getAtoms = {} as ReturnType<UseStoreApi<T, E>>['get'];
   const setAtoms = {} as ReturnType<UseStoreApi<T, E>>['set'];
   const useAtoms = {} as ReturnType<UseStoreApi<T, E>>['use'];
-  const primitiveAtoms = {} as PrimitiveAtomRecord<T>;
+  const primitiveAtoms = {} as WritableAtomRecord<T>;
 
   for (const key of Object.keys(initialState)) {
-    const atomConfig = atom(initialState[key as keyof T]);
+    const atomConfig = createAtom(initialState[key as keyof T]);
 
     (primitiveAtoms as any)[key] = atomConfig;
 
@@ -179,7 +190,7 @@ export const createAtomStore = <
   const atoms = {
     ...primitiveAtoms,
     ...(extend ? extend(primitiveAtoms) : {}),
-  } as PrimitiveAtomRecord<T> & E;
+  } as WritableAtomRecord<T> & E;
 
   for (const key of Object.keys(atoms)) {
     const atomConfig = atoms[key as keyof T & keyof E];
diff --git a/packages/jotai-x/src/useHydrateStore.ts b/packages/jotai-x/src/useHydrateStore.ts
index 517ba85..5ecbd8f 100644
--- a/packages/jotai-x/src/useHydrateStore.ts
+++ b/packages/jotai-x/src/useHydrateStore.ts
@@ -3,16 +3,16 @@ import { useSetAtom } from 'jotai';
 import { useHydrateAtoms } from 'jotai/utils';
 
 import type {
-  PrimitiveAtomRecord,
   UseHydrateAtoms,
   UseSyncAtoms,
+  WritableAtomRecord,
 } from './createAtomStore';
 
 /**
  * Hydrate atoms with initial values for SSR.
  */
 export const useHydrateStore = (
-  atoms: PrimitiveAtomRecord<any>,
+  atoms: WritableAtomRecord<any>,
   initialValues: Parameters<UseHydrateAtoms<any>>[0],
   options: Parameters<UseHydrateAtoms<any>>[1] = {}
 ) => {
@@ -38,7 +38,7 @@ export const useHydrateStore = (
  * Update atoms with new values on changes.
  */
 export const useSyncStore = (
-  atoms: PrimitiveAtomRecord<any>,
+  atoms: WritableAtomRecord<any>,
   values: any,
   { store }: Parameters<UseSyncAtoms<any>>[1] = {}
 ) => {

From e0ff6cca47ef2dd70c553f4209b36660e2441c61 Mon Sep 17 00:00:00 2001
From: Joe Anderson <joe@mousetrapped.co.uk>
Date: Thu, 21 Dec 2023 10:27:50 +0000
Subject: [PATCH 03/20] Expose `.(get|set|use)Atom` and `.store` on UseStoreApi

---
 packages/jotai-x/src/createAtomStore.spec.tsx |  35 +++++
 packages/jotai-x/src/createAtomStore.ts       | 128 ++++++++++++------
 2 files changed, 123 insertions(+), 40 deletions(-)

diff --git a/packages/jotai-x/src/createAtomStore.spec.tsx b/packages/jotai-x/src/createAtomStore.spec.tsx
index 945b8a6..cac9073 100644
--- a/packages/jotai-x/src/createAtomStore.spec.tsx
+++ b/packages/jotai-x/src/createAtomStore.spec.tsx
@@ -387,4 +387,39 @@ describe('createAtomStore', () => {
       expect(myAtom.isCustomAtom).toBe(true);
     });
   });
+
+  describe('arbitrary atom accessors', () => {
+    type User = {
+      name: string;
+    };
+
+    const initialUser: User = {
+      name: 'Jane',
+    };
+
+    const { userStore, useUserStore, UserProvider } = createAtomStore(
+      initialUser,
+      {
+        name: 'user' as const,
+      }
+    );
+
+    const derivedAtom = atom((get) => `My name is ${get(userStore.atom.name)}`);
+
+    const DerivedAtomConsumer = () => {
+      const message = useUserStore().getAtom(derivedAtom);
+
+      return <div>{message}</div>;
+    };
+
+    it('accesses arbitrary atom within store', () => {
+      const { getByText } = render(
+        <UserProvider name="John">
+          <DerivedAtomConsumer />
+        </UserProvider>
+      );
+
+      expect(getByText('My name is John')).toBeInTheDocument();
+    });
+  });
 });
diff --git a/packages/jotai-x/src/createAtomStore.ts b/packages/jotai-x/src/createAtomStore.ts
index 59bed86..e96bc43 100644
--- a/packages/jotai-x/src/createAtomStore.ts
+++ b/packages/jotai-x/src/createAtomStore.ts
@@ -62,6 +62,18 @@ export type StoreApi<
   name: N;
 };
 
+type GetAtomFn = <V>(atom: Atom<V>, options?: UseAtomOptionsOrScope) => V;
+
+type SetAtomFn = <V, A extends unknown[], R>(
+  atom: WritableAtom<V, A, R>,
+  options?: UseAtomOptionsOrScope
+) => (...args: A) => R;
+
+type UseAtomFn = <V, A extends unknown[], R>(
+  atom: WritableAtom<V, A, R>,
+  options?: UseAtomOptionsOrScope
+) => [V, (...args: A) => R];
+
 export type UseStoreApi<T, E> = (options?: UseAtomOptionsOrScope) => {
   get: GetRecord<
     T & {
@@ -70,6 +82,10 @@ export type UseStoreApi<T, E> = (options?: UseAtomOptionsOrScope) => {
   >;
   set: SetRecord<T>;
   use: UseRecord<T>;
+  getAtom: GetAtomFn;
+  setAtom: SetAtomFn;
+  useAtom: UseAtomFn;
+  store: (options?: UseAtomOptionsOrScope) => JotaiStore | undefined;
 };
 
 export type AtomStoreApi<
@@ -95,10 +111,10 @@ const getStoreIndex = (name = '') =>
 const getUseStoreIndex = (name = '') =>
   `use${capitalizeFirstLetter(name)}Store`;
 
-const withDefaultOptions = <T, R>(
-  fnRecord: { [key in keyof T]: (options?: UseAtomOptions) => R },
+const withDefaultOptions = <T extends object>(
+  fnRecord: T,
   defaultOptions: UseAtomOptions
-): typeof fnRecord =>
+): T =>
   Object.fromEntries(
     Object.entries(fnRecord).map(([key, fn]) => [
       key,
@@ -120,7 +136,6 @@ export interface CreateAtomStoreOptions<
   N extends string,
 > {
   name: N;
-  store?: UseAtomOptions['store'];
   delay?: UseAtomOptions['delay'];
   effect?: FC;
   extend?: (primitiveAtoms: WritableAtomRecord<T>) => E;
@@ -162,29 +177,39 @@ export const createAtomStore = <
   const useAtoms = {} as ReturnType<UseStoreApi<T, E>>['use'];
   const primitiveAtoms = {} as WritableAtomRecord<T>;
 
+  const useStore = (optionsOrScope: UseAtomOptionsOrScope = {}) => {
+    const { scope, store } = convertScopeShorthand(optionsOrScope);
+    const contextStore = useAtomStore(name, scope);
+    return store ?? contextStore;
+  };
+
+  const useAtomValueWithStore: GetAtomFn = (atomConfig, optionsOrScope) => {
+    const store = useStore(optionsOrScope);
+    const { delay = delayRoot } = convertScopeShorthand(optionsOrScope);
+    return useAtomValue(atomConfig, { store, delay });
+  };
+
+  const useSetAtomWithStore: SetAtomFn = (atomConfig, optionsOrScope) => {
+    const store = useStore(optionsOrScope);
+    return useSetAtom(atomConfig, { store });
+  };
+
+  const useAtomWithStore: UseAtomFn = (atomConfig, optionsOrScope) => {
+    const store = useStore(optionsOrScope);
+    const { delay = delayRoot } = convertScopeShorthand(optionsOrScope);
+    return useAtom(atomConfig, { store, delay });
+  };
+
   for (const key of Object.keys(initialState)) {
     const atomConfig = createAtom(initialState[key as keyof T]);
 
     (primitiveAtoms as any)[key] = atomConfig;
 
-    (setAtoms as any)[key] = (optionsOrScope: UseAtomOptionsOrScope = {}) => {
-      const options = convertScopeShorthand(optionsOrScope);
-      const contextStore = useAtomStore(name, options.scope);
-
-      return useSetAtom(atomConfig as any, {
-        store: options.store ?? contextStore,
-      });
-    };
+    (setAtoms as any)[key] = (optionsOrScope: UseAtomOptionsOrScope = {}) =>
+      useSetAtomWithStore(atomConfig, optionsOrScope);
 
-    (useAtoms as any)[key] = (optionsOrScope: UseAtomOptionsOrScope = {}) => {
-      const options = convertScopeShorthand(optionsOrScope);
-      const contextStore = useAtomStore(name, options.scope);
-
-      return useAtom(atomConfig, {
-        store: options.store ?? contextStore,
-        delay: options.delay ?? delayRoot,
-      });
-    };
+    (useAtoms as any)[key] = (optionsOrScope: UseAtomOptionsOrScope = {}) =>
+      useAtomWithStore(atomConfig, optionsOrScope);
   }
 
   const atoms = {
@@ -195,28 +220,51 @@ export const createAtomStore = <
   for (const key of Object.keys(atoms)) {
     const atomConfig = atoms[key as keyof T & keyof E];
 
-    (getAtoms as any)[key] = (optionsOrScope: UseAtomOptionsOrScope = {}) => {
-      const options = convertScopeShorthand(optionsOrScope);
-      const contextStore = useAtomStore(name, options.scope, false);
-
-      return useAtomValue(atomConfig, {
-        store: options.store ?? contextStore,
-        delay: options.delay ?? delayRoot,
-      });
-    };
+    (getAtoms as any)[key] = (optionsOrScope: UseAtomOptionsOrScope = {}) =>
+      useAtomValueWithStore(atomConfig, optionsOrScope);
   }
 
+  const Provider: FC<ProviderProps<T>> = createAtomProvider(
+    name,
+    primitiveAtoms,
+    { effect }
+  );
+
+  const storeApi: StoreApi<T, E, N> = {
+    atom: atoms,
+    name,
+  };
+
+  const useStoreApi: UseStoreApi<T, E> = (defaultOptions = {}) => ({
+    get: withDefaultOptions(getAtoms, convertScopeShorthand(defaultOptions)),
+    set: withDefaultOptions(setAtoms, convertScopeShorthand(defaultOptions)),
+    use: withDefaultOptions(useAtoms, convertScopeShorthand(defaultOptions)),
+    getAtom: (atomConfig, options) =>
+      useAtomValueWithStore(atomConfig, {
+        ...convertScopeShorthand(defaultOptions),
+        ...convertScopeShorthand(options),
+      }),
+    setAtom: (atomConfig, options) =>
+      useSetAtomWithStore(atomConfig, {
+        ...convertScopeShorthand(defaultOptions),
+        ...convertScopeShorthand(options),
+      }),
+    useAtom: (atomConfig, options) =>
+      useAtomWithStore(atomConfig, {
+        ...convertScopeShorthand(defaultOptions),
+        ...convertScopeShorthand(options),
+      }),
+    store: (options) =>
+      useStore({
+        ...convertScopeShorthand(defaultOptions),
+        ...convertScopeShorthand(options),
+      }),
+  });
+
   return {
-    [providerIndex]: createAtomProvider(name, primitiveAtoms, { effect }),
-    [useStoreIndex]: (options: UseAtomOptionsOrScope = {}) => ({
-      get: withDefaultOptions(getAtoms as any, convertScopeShorthand(options)),
-      set: withDefaultOptions(setAtoms as any, convertScopeShorthand(options)),
-      use: withDefaultOptions(useAtoms as any, convertScopeShorthand(options)),
-    }),
-    [storeIndex]: {
-      atom: atoms,
-      name,
-    },
+    [providerIndex]: Provider,
+    [useStoreIndex]: useStoreApi,
+    [storeIndex]: storeApi,
     name,
   } as any;
 };

From 898d26f7bfb40349c1c53022c480246e1c73e37b Mon Sep 17 00:00:00 2001
From: Joe Anderson <joe@mousetrapped.co.uk>
Date: Thu, 21 Dec 2023 13:22:22 +0000
Subject: [PATCH 04/20] Move {getAtom,setAtom,useAtom} to {get,set,use}.atom

---
 packages/jotai-x/src/createAtomStore.spec.tsx |  2 +-
 packages/jotai-x/src/createAtomStore.ts       | 69 ++++++++++---------
 2 files changed, 39 insertions(+), 32 deletions(-)

diff --git a/packages/jotai-x/src/createAtomStore.spec.tsx b/packages/jotai-x/src/createAtomStore.spec.tsx
index cac9073..5f4eb6c 100644
--- a/packages/jotai-x/src/createAtomStore.spec.tsx
+++ b/packages/jotai-x/src/createAtomStore.spec.tsx
@@ -407,7 +407,7 @@ describe('createAtomStore', () => {
     const derivedAtom = atom((get) => `My name is ${get(userStore.atom.name)}`);
 
     const DerivedAtomConsumer = () => {
-      const message = useUserStore().getAtom(derivedAtom);
+      const message = useUserStore().get.atom(derivedAtom);
 
       return <div>{message}</div>;
     };
diff --git a/packages/jotai-x/src/createAtomStore.ts b/packages/jotai-x/src/createAtomStore.ts
index e96bc43..e59ab58 100644
--- a/packages/jotai-x/src/createAtomStore.ts
+++ b/packages/jotai-x/src/createAtomStore.ts
@@ -20,9 +20,17 @@ type UseAtomOptionsOrScope = UseAtomOptions | string;
 export type GetRecord<O> = {
   [K in keyof O]: (options?: UseAtomOptionsOrScope) => O[K];
 };
+
+export type ExtendedGetRecord<O, E> = GetRecord<
+  O & {
+    [key in keyof E]: E[key] extends Atom<infer U> ? U : never;
+  }
+>;
+
 export type SetRecord<O> = {
   [K in keyof O]: (options?: UseAtomOptionsOrScope) => (value: O[K]) => void;
 };
+
 export type UseRecord<O> = {
   [K in keyof O]: (
     options?: UseAtomOptionsOrScope
@@ -75,16 +83,9 @@ type UseAtomFn = <V, A extends unknown[], R>(
 ) => [V, (...args: A) => R];
 
 export type UseStoreApi<T, E> = (options?: UseAtomOptionsOrScope) => {
-  get: GetRecord<
-    T & {
-      [key in keyof E]: E[key] extends Atom<infer U> ? U : never;
-    }
-  >;
-  set: SetRecord<T>;
-  use: UseRecord<T>;
-  getAtom: GetAtomFn;
-  setAtom: SetAtomFn;
-  useAtom: UseAtomFn;
+  get: ExtendedGetRecord<T, E> & { atom: GetAtomFn };
+  set: SetRecord<T> & { atom: SetAtomFn };
+  use: UseRecord<T> & { atom: UseAtomFn };
   store: (options?: UseAtomOptionsOrScope) => JotaiStore | undefined;
 };
 
@@ -172,9 +173,9 @@ export const createAtomStore = <
   const useStoreIndex = getUseStoreIndex(name) as UseNameStore<N>;
   const storeIndex = getStoreIndex(name) as NameStore<N>;
 
-  const getAtoms = {} as ReturnType<UseStoreApi<T, E>>['get'];
-  const setAtoms = {} as ReturnType<UseStoreApi<T, E>>['set'];
-  const useAtoms = {} as ReturnType<UseStoreApi<T, E>>['use'];
+  const getAtoms = {} as ExtendedGetRecord<T, E>;
+  const setAtoms = {} as SetRecord<T>;
+  const useAtoms = {} as UseRecord<T>;
   const primitiveAtoms = {} as WritableAtomRecord<T>;
 
   const useStore = (optionsOrScope: UseAtomOptionsOrScope = {}) => {
@@ -236,24 +237,30 @@ export const createAtomStore = <
   };
 
   const useStoreApi: UseStoreApi<T, E> = (defaultOptions = {}) => ({
-    get: withDefaultOptions(getAtoms, convertScopeShorthand(defaultOptions)),
-    set: withDefaultOptions(setAtoms, convertScopeShorthand(defaultOptions)),
-    use: withDefaultOptions(useAtoms, convertScopeShorthand(defaultOptions)),
-    getAtom: (atomConfig, options) =>
-      useAtomValueWithStore(atomConfig, {
-        ...convertScopeShorthand(defaultOptions),
-        ...convertScopeShorthand(options),
-      }),
-    setAtom: (atomConfig, options) =>
-      useSetAtomWithStore(atomConfig, {
-        ...convertScopeShorthand(defaultOptions),
-        ...convertScopeShorthand(options),
-      }),
-    useAtom: (atomConfig, options) =>
-      useAtomWithStore(atomConfig, {
-        ...convertScopeShorthand(defaultOptions),
-        ...convertScopeShorthand(options),
-      }),
+    get: {
+      ...withDefaultOptions(getAtoms, convertScopeShorthand(defaultOptions)),
+      atom: (atomConfig, options) =>
+        useAtomValueWithStore(atomConfig, {
+          ...convertScopeShorthand(defaultOptions),
+          ...convertScopeShorthand(options),
+        }),
+    },
+    set: {
+      ...withDefaultOptions(setAtoms, convertScopeShorthand(defaultOptions)),
+      atom: (atomConfig, options) =>
+        useSetAtomWithStore(atomConfig, {
+          ...convertScopeShorthand(defaultOptions),
+          ...convertScopeShorthand(options),
+        }),
+    },
+    use: {
+      ...withDefaultOptions(useAtoms, convertScopeShorthand(defaultOptions)),
+      atom: (atomConfig, options) =>
+        useAtomWithStore(atomConfig, {
+          ...convertScopeShorthand(defaultOptions),
+          ...convertScopeShorthand(options),
+        }),
+    },
     store: (options) =>
       useStore({
         ...convertScopeShorthand(defaultOptions),

From b43fdbb925b59d714d8d6c371f5fe49b01af0de7 Mon Sep 17 00:00:00 2001
From: zbeyens <zbeyens@udecode.io>
Date: Thu, 21 Dec 2023 15:25:58 +0100
Subject: [PATCH 05/20] extend set, use

---
 packages/jotai-x/src/createAtomStore.spec.tsx | 180 +++++++++++++++++-
 packages/jotai-x/src/createAtomStore.ts       |  36 ++--
 2 files changed, 202 insertions(+), 14 deletions(-)

diff --git a/packages/jotai-x/src/createAtomStore.spec.tsx b/packages/jotai-x/src/createAtomStore.spec.tsx
index 5f4eb6c..5188efa 100644
--- a/packages/jotai-x/src/createAtomStore.spec.tsx
+++ b/packages/jotai-x/src/createAtomStore.spec.tsx
@@ -1,8 +1,9 @@
 import '@testing-library/jest-dom';
 
 import React, { ReactNode, useState } from 'react';
-import { act, render, renderHook } from '@testing-library/react';
+import { act, queryByText, render, renderHook } from '@testing-library/react';
 import { atom, PrimitiveAtom, useAtomValue } from 'jotai';
+import { splitAtom } from 'jotai/utils';
 
 import { createAtomStore } from './createAtomStore';
 
@@ -422,4 +423,181 @@ describe('createAtomStore', () => {
       expect(getByText('My name is John')).toBeInTheDocument();
     });
   });
+
+  describe('splitAtoms using todoStore.atom.items', () => {
+    const initialState = {
+      items: [] as {
+        task: string;
+        done: boolean;
+      }[],
+    };
+
+    const { todoStore, useTodoStore, TodoProvider } = createAtomStore(
+      initialState,
+      {
+        name: 'todo' as const,
+      }
+    );
+
+    const todoAtomsAtom = splitAtom(todoStore.atom.items);
+
+    type TodoType = (typeof initialState)['items'][number];
+
+    const TodoItem = ({
+      todoAtom,
+      remove,
+    }: {
+      todoAtom: PrimitiveAtom<TodoType>;
+      remove: () => void;
+    }) => {
+      const [todo, setTodo] = useTodoStore().use.atom(todoAtom);
+
+      return (
+        <div>
+          <label>{todo.task}</label>
+          <input
+            type="checkbox"
+            checked={todo.done}
+            onChange={() => {
+              setTodo((oldValue) => ({ ...oldValue, done: !oldValue.done }));
+            }}
+          />
+          {/* eslint-disable-next-line react/button-has-type */}
+          <button onClick={remove}>remove {todo.task}</button>
+        </div>
+      );
+    };
+
+    const TodoList = () => {
+      const [todoAtoms, dispatch] = useTodoStore().use.atom(todoAtomsAtom);
+      return (
+        <ul>
+          {todoAtoms.map((todoAtom) => (
+            <TodoItem
+              key={`${todoAtom}`}
+              todoAtom={todoAtom}
+              remove={() => dispatch({ type: 'remove', atom: todoAtom })}
+            />
+          ))}
+        </ul>
+      );
+    };
+
+    it('should work', () => {
+      const { getByText, container } = render(
+        <TodoProvider
+          initialValues={{
+            items: [
+              {
+                task: 'help the town',
+                done: false,
+              },
+              {
+                task: 'feed the dragon',
+                done: false,
+              },
+            ],
+          }}
+        >
+          <TodoList />
+        </TodoProvider>
+      );
+
+      expect(getByText('help the town')).toBeInTheDocument();
+      expect(getByText('feed the dragon')).toBeInTheDocument();
+
+      act(() => getByText('remove help the town').click());
+
+      expect(queryByText(container, 'help the town')).not.toBeInTheDocument();
+      expect(getByText('feed the dragon')).toBeInTheDocument();
+    });
+  });
+
+  describe('splitAtoms using extend', () => {
+    const initialState = {
+      items: [] as {
+        task: string;
+        done: boolean;
+      }[],
+    };
+
+    const { useTodoStore, TodoProvider } = createAtomStore(initialState, {
+      name: 'todo' as const,
+      extend: ({ items }) => ({
+        itemAtoms: splitAtom(items),
+      }),
+    });
+
+    type TodoType = (typeof initialState)['items'][number];
+
+    const TodoItem = ({
+      todoAtom,
+      remove,
+    }: {
+      todoAtom: PrimitiveAtom<TodoType>;
+      remove: () => void;
+    }) => {
+      const [todo, setTodo] = useTodoStore().use.atom(todoAtom);
+
+      return (
+        <div>
+          <label>{todo.task}</label>
+          <input
+            type="checkbox"
+            checked={todo.done}
+            onChange={() => {
+              setTodo((oldValue) => ({ ...oldValue, done: !oldValue.done }));
+            }}
+          />
+          {/* eslint-disable-next-line react/button-has-type */}
+          <button onClick={remove}>remove {todo.task}</button>
+        </div>
+      );
+    };
+
+    const TodoList = () => {
+      const [todoAtoms, dispatch] = useTodoStore().use.itemAtoms();
+
+      return (
+        <ul>
+          {todoAtoms.map((todoAtom) => (
+            <TodoItem
+              key={`${todoAtom}`}
+              todoAtom={todoAtom}
+              remove={() => dispatch({ type: 'remove', atom: todoAtom })}
+            />
+          ))}
+        </ul>
+      );
+    };
+
+    it('should work', () => {
+      const { getByText, container } = render(
+        <TodoProvider
+          initialValues={{
+            items: [
+              {
+                task: 'help the town',
+                done: false,
+              },
+              {
+                task: 'feed the dragon',
+                done: false,
+              },
+            ],
+          }}
+        >
+          <TodoList />
+        </TodoProvider>
+      );
+
+      expect(getByText('help the town')).toBeInTheDocument();
+      expect(getByText('feed the dragon')).toBeInTheDocument();
+
+      act(() => getByText('remove help the town').click());
+
+      expect(queryByText(container, 'help the town')).not.toBeInTheDocument();
+      expect(getByText('feed the dragon')).toBeInTheDocument();
+    });
+  });
 });
diff --git a/packages/jotai-x/src/createAtomStore.ts b/packages/jotai-x/src/createAtomStore.ts
index e59ab58..912c24c 100644
--- a/packages/jotai-x/src/createAtomStore.ts
+++ b/packages/jotai-x/src/createAtomStore.ts
@@ -31,12 +31,24 @@ export type SetRecord<O> = {
   [K in keyof O]: (options?: UseAtomOptionsOrScope) => (value: O[K]) => void;
 };
 
+export type ExtendedSetRecord<O, E> = SetRecord<
+  O & {
+    [key in keyof E]: E[key] extends Atom<infer U> ? U : never;
+  }
+>;
+
 export type UseRecord<O> = {
   [K in keyof O]: (
     options?: UseAtomOptionsOrScope
   ) => [O[K], (value: O[K]) => void];
 };
 
+export type ExtendedUseRecord<O, E> = UseRecord<
+  O & {
+    [key in keyof E]: E[key] extends Atom<infer U> ? U : never;
+  }
+>;
+
 export type SimpleWritableAtom<T> = WritableAtom<T, [T], void>;
 
 export type WritableAtomRecord<O> = {
@@ -84,8 +96,8 @@ type UseAtomFn = <V, A extends unknown[], R>(
 
 export type UseStoreApi<T, E> = (options?: UseAtomOptionsOrScope) => {
   get: ExtendedGetRecord<T, E> & { atom: GetAtomFn };
-  set: SetRecord<T> & { atom: SetAtomFn };
-  use: UseRecord<T> & { atom: UseAtomFn };
+  set: ExtendedSetRecord<T, E> & { atom: SetAtomFn };
+  use: ExtendedUseRecord<T, E> & { atom: UseAtomFn };
   store: (options?: UseAtomOptionsOrScope) => JotaiStore | undefined;
 };
 
@@ -174,8 +186,8 @@ export const createAtomStore = <
   const storeIndex = getStoreIndex(name) as NameStore<N>;
 
   const getAtoms = {} as ExtendedGetRecord<T, E>;
-  const setAtoms = {} as SetRecord<T>;
-  const useAtoms = {} as UseRecord<T>;
+  const setAtoms = {} as ExtendedSetRecord<T, E>;
+  const useAtoms = {} as ExtendedUseRecord<T, E>;
   const primitiveAtoms = {} as WritableAtomRecord<T>;
 
   const useStore = (optionsOrScope: UseAtomOptionsOrScope = {}) => {
@@ -202,15 +214,7 @@ export const createAtomStore = <
   };
 
   for (const key of Object.keys(initialState)) {
-    const atomConfig = createAtom(initialState[key as keyof T]);
-
-    (primitiveAtoms as any)[key] = atomConfig;
-
-    (setAtoms as any)[key] = (optionsOrScope: UseAtomOptionsOrScope = {}) =>
-      useSetAtomWithStore(atomConfig, optionsOrScope);
-
-    (useAtoms as any)[key] = (optionsOrScope: UseAtomOptionsOrScope = {}) =>
-      useAtomWithStore(atomConfig, optionsOrScope);
+    (primitiveAtoms as any)[key] = createAtom(initialState[key as keyof T]);
   }
 
   const atoms = {
@@ -223,6 +227,12 @@ export const createAtomStore = <
 
     (getAtoms as any)[key] = (optionsOrScope: UseAtomOptionsOrScope = {}) =>
       useAtomValueWithStore(atomConfig, optionsOrScope);
+
+    (setAtoms as any)[key] = (optionsOrScope: UseAtomOptionsOrScope = {}) =>
+      useSetAtomWithStore(atomConfig, optionsOrScope);
+
+    (useAtoms as any)[key] = (optionsOrScope: UseAtomOptionsOrScope = {}) =>
+      useAtomWithStore(atomConfig, optionsOrScope);
   }
 
   const Provider: FC<ProviderProps<T>> = createAtomProvider(

From 301380c0b158724b69407e7a0e9737b61cac3ba9 Mon Sep 17 00:00:00 2001
From: zbeyens <zbeyens@udecode.io>
Date: Thu, 21 Dec 2023 15:33:01 +0100
Subject: [PATCH 06/20] fix test

---
 packages/jotai-x/src/createAtomStore.spec.tsx | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/packages/jotai-x/src/createAtomStore.spec.tsx b/packages/jotai-x/src/createAtomStore.spec.tsx
index 5188efa..434e98e 100644
--- a/packages/jotai-x/src/createAtomStore.spec.tsx
+++ b/packages/jotai-x/src/createAtomStore.spec.tsx
@@ -342,14 +342,14 @@ describe('createAtomStore', () => {
       expect(result.current).toBe('Jane is 98 years old');
     });
 
-    it('does not include extended atom in set hooks', () => {
+    it('does include extended atom in set hooks', () => {
       const { result } = renderHook(() => Object.keys(useUserStore().set));
-      expect(result.current).not.toContain('bio');
+      expect(result.current).toContain('bio');
     });
 
-    it('does not include extended atom in use hooks', () => {
+    it('does include extended atom in use hooks', () => {
       const { result } = renderHook(() => Object.keys(useUserStore().use));
-      expect(result.current).not.toContain('bio');
+      expect(result.current).toContain('bio');
     });
 
     it('computes extended atom based on current state', () => {

From 56dcfc201462ce0fd60f1a20db84a46854d321ec Mon Sep 17 00:00:00 2001
From: Joe Anderson <joe@mousetrapped.co.uk>
Date: Thu, 21 Dec 2023 14:49:50 +0000
Subject: [PATCH 07/20] Fix types for writable atoms in extend

---
 packages/jotai-x/src/createAtomStore.spec.tsx | 10 -----
 packages/jotai-x/src/createAtomStore.ts       | 38 ++++++++-----------
 2 files changed, 15 insertions(+), 33 deletions(-)

diff --git a/packages/jotai-x/src/createAtomStore.spec.tsx b/packages/jotai-x/src/createAtomStore.spec.tsx
index 5188efa..ce520c3 100644
--- a/packages/jotai-x/src/createAtomStore.spec.tsx
+++ b/packages/jotai-x/src/createAtomStore.spec.tsx
@@ -342,16 +342,6 @@ describe('createAtomStore', () => {
       expect(result.current).toBe('Jane is 98 years old');
     });
 
-    it('does not include extended atom in set hooks', () => {
-      const { result } = renderHook(() => Object.keys(useUserStore().set));
-      expect(result.current).not.toContain('bio');
-    });
-
-    it('does not include extended atom in use hooks', () => {
-      const { result } = renderHook(() => Object.keys(useUserStore().use));
-      expect(result.current).not.toContain('bio');
-    });
-
     it('computes extended atom based on current state', () => {
       const { getByText } = render(
         <UserProvider name="John" age={42}>
diff --git a/packages/jotai-x/src/createAtomStore.ts b/packages/jotai-x/src/createAtomStore.ts
index 912c24c..0331801 100644
--- a/packages/jotai-x/src/createAtomStore.ts
+++ b/packages/jotai-x/src/createAtomStore.ts
@@ -17,37 +17,29 @@ export type UseAtomOptions = {
 
 type UseAtomOptionsOrScope = UseAtomOptions | string;
 
-export type GetRecord<O> = {
-  [K in keyof O]: (options?: UseAtomOptionsOrScope) => O[K];
+type SetRecord<O> = {
+  [K in keyof O]: O[K] extends WritableAtom<never, infer A, infer R>
+    ? (options?: UseAtomOptionsOrScope) => (...args: A) => R
+    : never;
 };
 
-export type ExtendedGetRecord<O, E> = GetRecord<
-  O & {
-    [key in keyof E]: E[key] extends Atom<infer U> ? U : never;
-  }
->;
+type ExtendedSetRecord<T, E> = SetRecord<{ [K in keyof T]: Atom<T[K]> } & E>;
 
-export type SetRecord<O> = {
-  [K in keyof O]: (options?: UseAtomOptionsOrScope) => (value: O[K]) => void;
+type GetRecord<O> = {
+  [K in keyof O]: O[K] extends Atom<infer V>
+    ? (options?: UseAtomOptionsOrScope) => V
+    : never;
 };
 
-export type ExtendedSetRecord<O, E> = SetRecord<
-  O & {
-    [key in keyof E]: E[key] extends Atom<infer U> ? U : never;
-  }
->;
+type ExtendedGetRecord<T, E> = GetRecord<{ [K in keyof T]: Atom<T[K]> } & E>;
 
-export type UseRecord<O> = {
-  [K in keyof O]: (
-    options?: UseAtomOptionsOrScope
-  ) => [O[K], (value: O[K]) => void];
+type UseRecord<O> = {
+  [K in keyof O]: O[K] extends WritableAtom<infer V, infer A, infer R>
+    ? (options?: UseAtomOptionsOrScope) => [V, (...args: A) => R]
+    : never;
 };
 
-export type ExtendedUseRecord<O, E> = UseRecord<
-  O & {
-    [key in keyof E]: E[key] extends Atom<infer U> ? U : never;
-  }
->;
+type ExtendedUseRecord<T, E> = UseRecord<{ [K in keyof T]: Atom<T[K]> } & E>;
 
 export type SimpleWritableAtom<T> = WritableAtom<T, [T], void>;
 

From a9db7814f079e5dcb5cd3ad2633a6631782c6e77 Mon Sep 17 00:00:00 2001
From: Joe Anderson <joe@mousetrapped.co.uk>
Date: Thu, 21 Dec 2023 15:03:07 +0000
Subject: [PATCH 08/20] Fix type inference

---
 packages/jotai-x/src/createAtomStore.ts | 18 +++++++++---------
 1 file changed, 9 insertions(+), 9 deletions(-)

diff --git a/packages/jotai-x/src/createAtomStore.ts b/packages/jotai-x/src/createAtomStore.ts
index 0331801..1e21fa1 100644
--- a/packages/jotai-x/src/createAtomStore.ts
+++ b/packages/jotai-x/src/createAtomStore.ts
@@ -17,14 +17,6 @@ export type UseAtomOptions = {
 
 type UseAtomOptionsOrScope = UseAtomOptions | string;
 
-type SetRecord<O> = {
-  [K in keyof O]: O[K] extends WritableAtom<never, infer A, infer R>
-    ? (options?: UseAtomOptionsOrScope) => (...args: A) => R
-    : never;
-};
-
-type ExtendedSetRecord<T, E> = SetRecord<{ [K in keyof T]: Atom<T[K]> } & E>;
-
 type GetRecord<O> = {
   [K in keyof O]: O[K] extends Atom<infer V>
     ? (options?: UseAtomOptionsOrScope) => V
@@ -33,13 +25,21 @@ type GetRecord<O> = {
 
 type ExtendedGetRecord<T, E> = GetRecord<{ [K in keyof T]: Atom<T[K]> } & E>;
 
+type SetRecord<O> = {
+  [K in keyof O]: O[K] extends WritableAtom<infer _V, infer A, infer R>
+    ? (options?: UseAtomOptionsOrScope) => (...args: A) => R
+    : never;
+};
+
+type ExtendedSetRecord<T, E> = SetRecord<{ [K in keyof T]: SimpleWritableAtom<T[K]> } & E>;
+
 type UseRecord<O> = {
   [K in keyof O]: O[K] extends WritableAtom<infer V, infer A, infer R>
     ? (options?: UseAtomOptionsOrScope) => [V, (...args: A) => R]
     : never;
 };
 
-type ExtendedUseRecord<T, E> = UseRecord<{ [K in keyof T]: Atom<T[K]> } & E>;
+type ExtendedUseRecord<T, E> = UseRecord<{ [K in keyof T]: SimpleWritableAtom<T[K]> } & E>;
 
 export type SimpleWritableAtom<T> = WritableAtom<T, [T], void>;
 

From 134ec07fe87213dfba94d00a26dd72836077a901 Mon Sep 17 00:00:00 2001
From: Joe Anderson <joe@mousetrapped.co.uk>
Date: Thu, 21 Dec 2023 15:04:09 +0000
Subject: [PATCH 09/20] Linter fixes

---
 packages/jotai-x/src/createAtomStore.ts | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/packages/jotai-x/src/createAtomStore.ts b/packages/jotai-x/src/createAtomStore.ts
index 1e21fa1..3a9764f 100644
--- a/packages/jotai-x/src/createAtomStore.ts
+++ b/packages/jotai-x/src/createAtomStore.ts
@@ -31,7 +31,9 @@ type SetRecord<O> = {
     : never;
 };
 
-type ExtendedSetRecord<T, E> = SetRecord<{ [K in keyof T]: SimpleWritableAtom<T[K]> } & E>;
+type ExtendedSetRecord<T, E> = SetRecord<
+  { [K in keyof T]: SimpleWritableAtom<T[K]> } & E
+>;
 
 type UseRecord<O> = {
   [K in keyof O]: O[K] extends WritableAtom<infer V, infer A, infer R>
@@ -39,7 +41,9 @@ type UseRecord<O> = {
     : never;
 };
 
-type ExtendedUseRecord<T, E> = UseRecord<{ [K in keyof T]: SimpleWritableAtom<T[K]> } & E>;
+type ExtendedUseRecord<T, E> = UseRecord<
+  { [K in keyof T]: SimpleWritableAtom<T[K]> } & E
+>;
 
 export type SimpleWritableAtom<T> = WritableAtom<T, [T], void>;
 

From 617d04fab4311bf0ad90454f3997e19590643ecc Mon Sep 17 00:00:00 2001
From: Joe Anderson <joe@mousetrapped.co.uk>
Date: Thu, 21 Dec 2023 17:41:19 +0000
Subject: [PATCH 10/20] Let consumer pass atoms in `initialState`

---
 packages/jotai-x/src/createAtomProvider.tsx   |   6 +-
 packages/jotai-x/src/createAtomStore.spec.tsx |  15 +--
 packages/jotai-x/src/createAtomStore.ts       | 126 ++++++++++++------
 packages/jotai-x/src/useHydrateStore.ts       |   6 +-
 4 files changed, 95 insertions(+), 58 deletions(-)

diff --git a/packages/jotai-x/src/createAtomProvider.tsx b/packages/jotai-x/src/createAtomProvider.tsx
index 1f95e6c..efcc6d8 100644
--- a/packages/jotai-x/src/createAtomProvider.tsx
+++ b/packages/jotai-x/src/createAtomProvider.tsx
@@ -9,7 +9,7 @@ import React, {
 import { createStore } from 'jotai/vanilla';
 
 import { AtomProvider, AtomProviderProps } from './atomProvider';
-import { JotaiStore, WritableAtomRecord } from './createAtomStore';
+import { JotaiStore, SimpleWritableAtomRecord } from './createAtomStore';
 import { useHydrateStore, useSyncStore } from './useHydrateStore';
 
 const getFullyQualifiedScope = (storeName: string, scope: string) => {
@@ -62,7 +62,7 @@ export const HydrateAtoms = <T extends object>({
   atoms,
   ...props
 }: Omit<ProviderProps<T>, 'scope'> & {
-  atoms: WritableAtomRecord<T>;
+  atoms: SimpleWritableAtomRecord<T>;
 }) => {
   useHydrateStore(atoms, { ...initialValues, ...props } as any, {
     store,
@@ -81,7 +81,7 @@ export const HydrateAtoms = <T extends object>({
  */
 export const createAtomProvider = <T extends object, N extends string = ''>(
   storeScope: N,
-  atoms: WritableAtomRecord<T>,
+  atoms: SimpleWritableAtomRecord<T>,
   options: { effect?: FC } = {}
 ) => {
   const Effect = options.effect;
diff --git a/packages/jotai-x/src/createAtomStore.spec.tsx b/packages/jotai-x/src/createAtomStore.spec.tsx
index 50d8795..3312718 100644
--- a/packages/jotai-x/src/createAtomStore.spec.tsx
+++ b/packages/jotai-x/src/createAtomStore.spec.tsx
@@ -342,14 +342,14 @@ describe('createAtomStore', () => {
       expect(result.current).toBe('Jane is 98 years old');
     });
 
-    it('includes extended atom in set hooks', () => {
+    it('does not include read-only extended atom in set hooks', () => {
       const { result } = renderHook(() => Object.keys(useUserStore().set));
-      expect(result.current).toContain('bio');
+      expect(result.current).not.toContain('bio');
     });
 
-    it('includes extended atom in use hooks', () => {
+    it('does not include read-only extended atom in use hooks', () => {
       const { result } = renderHook(() => Object.keys(useUserStore().use));
-      expect(result.current).toContain('bio');
+      expect(result.current).not.toContain('bio');
     });
 
     it('computes extended atom based on current state', () => {
@@ -363,7 +363,7 @@ describe('createAtomStore', () => {
     });
   });
 
-  describe('custom createAtom function', () => {
+  describe('passing atoms as part of initial state', () => {
     type CustomAtom<T> = PrimitiveAtom<T> & {
       isCustomAtom: true;
     };
@@ -375,15 +375,14 @@ describe('createAtomStore', () => {
 
     const { customStore } = createAtomStore(
       {
-        x: 5,
+        x: createCustomAtom(1),
       },
       {
         name: 'custom' as const,
-        createAtom: createCustomAtom,
       }
     );
 
-    it('uses custom createAtom function', () => {
+    it('uses passed atom', () => {
       const myAtom = customStore.atom.x as CustomAtom<number>;
       expect(myAtom.isCustomAtom).toBe(true);
     });
diff --git a/packages/jotai-x/src/createAtomStore.ts b/packages/jotai-x/src/createAtomStore.ts
index 3a9764f..8d5f61d 100644
--- a/packages/jotai-x/src/createAtomStore.ts
+++ b/packages/jotai-x/src/createAtomStore.ts
@@ -23,32 +23,34 @@ type GetRecord<O> = {
     : never;
 };
 
-type ExtendedGetRecord<T, E> = GetRecord<{ [K in keyof T]: Atom<T[K]> } & E>;
-
 type SetRecord<O> = {
   [K in keyof O]: O[K] extends WritableAtom<infer _V, infer A, infer R>
     ? (options?: UseAtomOptionsOrScope) => (...args: A) => R
     : never;
 };
 
-type ExtendedSetRecord<T, E> = SetRecord<
-  { [K in keyof T]: SimpleWritableAtom<T[K]> } & E
->;
-
 type UseRecord<O> = {
   [K in keyof O]: O[K] extends WritableAtom<infer V, infer A, infer R>
     ? (options?: UseAtomOptionsOrScope) => [V, (...args: A) => R]
     : never;
 };
 
-type ExtendedUseRecord<T, E> = UseRecord<
-  { [K in keyof T]: SimpleWritableAtom<T[K]> } & E
->;
+type StoreAtomsWithoutExtend<T> = {
+  [K in keyof T]: T[K] extends Atom<any> ? T[K] : SimpleWritableAtom<T[K]>;
+};
+
+type StoreAtoms<T, E> = StoreAtomsWithoutExtend<T> & E;
+
+type FilterWritableAtoms<T> = {
+  [K in keyof T]-?: T[K] extends WritableAtom<any, any, any> ? T[K] : never;
+};
+
+type WritableStoreAtoms<T, E> = FilterWritableAtoms<StoreAtoms<T, E>>;
 
 export type SimpleWritableAtom<T> = WritableAtom<T, [T], void>;
 
-export type WritableAtomRecord<O> = {
-  [K in keyof O]: SimpleWritableAtom<O[K]>;
+export type SimpleWritableAtomRecord<T> = {
+  [K in keyof T]: SimpleWritableAtom<T[K]>;
 };
 
 export type AtomRecord<O> = {
@@ -74,7 +76,7 @@ export type StoreApi<
   E extends AtomRecord<object>,
   N extends string = '',
 > = {
-  atom: WritableAtomRecord<T> & E;
+  atom: StoreAtoms<T, E>;
   name: N;
 };
 
@@ -91,9 +93,9 @@ type UseAtomFn = <V, A extends unknown[], R>(
 ) => [V, (...args: A) => R];
 
 export type UseStoreApi<T, E> = (options?: UseAtomOptionsOrScope) => {
-  get: ExtendedGetRecord<T, E> & { atom: GetAtomFn };
-  set: ExtendedSetRecord<T, E> & { atom: SetAtomFn };
-  use: ExtendedUseRecord<T, E> & { atom: UseAtomFn };
+  get: GetRecord<StoreAtoms<T, E>> & { atom: GetAtomFn };
+  set: SetRecord<WritableStoreAtoms<T, E>> & { atom: SetAtomFn };
+  use: UseRecord<WritableStoreAtoms<T, E>> & { atom: UseAtomFn };
   store: (options?: UseAtomOptionsOrScope) => JotaiStore | undefined;
 };
 
@@ -120,6 +122,12 @@ const getStoreIndex = (name = '') =>
 const getUseStoreIndex = (name = '') =>
   `use${capitalizeFirstLetter(name)}Store`;
 
+const isAtom = (possibleAtom: unknown): boolean =>
+  !!possibleAtom &&
+  typeof possibleAtom === 'object' &&
+  'read' in possibleAtom &&
+  typeof possibleAtom.read === 'function';
+
 const withDefaultOptions = <T extends object>(
   fnRecord: T,
   defaultOptions: UseAtomOptions
@@ -147,8 +155,7 @@ export interface CreateAtomStoreOptions<
   name: N;
   delay?: UseAtomOptions['delay'];
   effect?: FC;
-  extend?: (primitiveAtoms: WritableAtomRecord<T>) => E;
-  createAtom?: <V>(value: V) => SimpleWritableAtom<V>;
+  extend?: (atomsWithoutExtend: StoreAtomsWithoutExtend<T>) => E;
 }
 
 /**
@@ -169,22 +176,53 @@ export const createAtomStore = <
   N extends string = '',
 >(
   initialState: T,
-  {
-    name,
-    delay: delayRoot,
-    effect,
-    extend,
-    createAtom = atom,
-  }: CreateAtomStoreOptions<T, E, N>
+  { name, delay: delayRoot, effect, extend }: CreateAtomStoreOptions<T, E, N>
 ): AtomStoreApi<T, E, N> => {
+  type MyStoreAtoms = StoreAtoms<T, E>;
+  type MyWritableStoreAtoms = WritableStoreAtoms<T, E>;
+  type MyStoreAtomsWithoutExtend = StoreAtomsWithoutExtend<T>;
+  type MyWritableStoreAtomsWithoutExtend =
+    FilterWritableAtoms<MyStoreAtomsWithoutExtend>;
+
   const providerIndex = getProviderIndex(name) as NameProvider<N>;
   const useStoreIndex = getUseStoreIndex(name) as UseNameStore<N>;
   const storeIndex = getStoreIndex(name) as NameStore<N>;
 
-  const getAtoms = {} as ExtendedGetRecord<T, E>;
-  const setAtoms = {} as ExtendedSetRecord<T, E>;
-  const useAtoms = {} as ExtendedUseRecord<T, E>;
-  const primitiveAtoms = {} as WritableAtomRecord<T>;
+  const atomsWithoutExtend = {} as MyStoreAtomsWithoutExtend;
+  const writableAtomsWithoutExtend = {} as MyWritableStoreAtomsWithoutExtend;
+  const atomIsWritable = {} as Record<keyof MyStoreAtoms, boolean>;
+
+  for (const [key, atomOrValue] of Object.entries(initialState)) {
+    const atomConfig: Atom<unknown> = isAtom(atomOrValue)
+      ? atomOrValue
+      : atom(atomOrValue);
+    atomsWithoutExtend[key as keyof MyStoreAtomsWithoutExtend] =
+      atomConfig as any;
+
+    const writable = 'write' in atomConfig;
+    atomIsWritable[key as keyof MyStoreAtoms] = writable;
+
+    if (writable) {
+      writableAtomsWithoutExtend[
+        key as keyof MyWritableStoreAtomsWithoutExtend
+      ] = atomConfig as any;
+    }
+  }
+
+  const atoms = { ...atomsWithoutExtend } as MyStoreAtoms;
+
+  if (extend) {
+    const extendedAtoms = extend(atomsWithoutExtend);
+
+    for (const [key, atomConfig] of Object.entries(extendedAtoms)) {
+      atoms[key as keyof MyStoreAtoms] = atomConfig;
+      atomIsWritable[key as keyof MyStoreAtoms] = 'write' in atomConfig;
+    }
+  }
+
+  const getAtoms = {} as GetRecord<MyStoreAtoms>;
+  const setAtoms = {} as SetRecord<MyWritableStoreAtoms>;
+  const useAtoms = {} as UseRecord<MyWritableStoreAtoms>;
 
   const useStore = (optionsOrScope: UseAtomOptionsOrScope = {}) => {
     const { scope, store } = convertScopeShorthand(optionsOrScope);
@@ -209,31 +247,31 @@ export const createAtomStore = <
     return useAtom(atomConfig, { store, delay });
   };
 
-  for (const key of Object.keys(initialState)) {
-    (primitiveAtoms as any)[key] = createAtom(initialState[key as keyof T]);
-  }
-
-  const atoms = {
-    ...primitiveAtoms,
-    ...(extend ? extend(primitiveAtoms) : {}),
-  } as WritableAtomRecord<T> & E;
-
   for (const key of Object.keys(atoms)) {
-    const atomConfig = atoms[key as keyof T & keyof E];
+    const atomConfig = atoms[key as keyof MyStoreAtoms];
+    const isWritable: boolean = atomIsWritable[key as keyof MyStoreAtoms];
 
     (getAtoms as any)[key] = (optionsOrScope: UseAtomOptionsOrScope = {}) =>
       useAtomValueWithStore(atomConfig, optionsOrScope);
 
-    (setAtoms as any)[key] = (optionsOrScope: UseAtomOptionsOrScope = {}) =>
-      useSetAtomWithStore(atomConfig, optionsOrScope);
-
-    (useAtoms as any)[key] = (optionsOrScope: UseAtomOptionsOrScope = {}) =>
-      useAtomWithStore(atomConfig, optionsOrScope);
+    if (isWritable) {
+      (setAtoms as any)[key] = (optionsOrScope: UseAtomOptionsOrScope = {}) =>
+        useSetAtomWithStore(
+          atomConfig as WritableAtom<any, any, any>,
+          optionsOrScope
+        );
+
+      (useAtoms as any)[key] = (optionsOrScope: UseAtomOptionsOrScope = {}) =>
+        useAtomWithStore(
+          atomConfig as WritableAtom<any, any, any>,
+          optionsOrScope
+        );
+    }
   }
 
   const Provider: FC<ProviderProps<T>> = createAtomProvider(
     name,
-    primitiveAtoms,
+    writableAtomsWithoutExtend,
     { effect }
   );
 
diff --git a/packages/jotai-x/src/useHydrateStore.ts b/packages/jotai-x/src/useHydrateStore.ts
index 5ecbd8f..bb45d45 100644
--- a/packages/jotai-x/src/useHydrateStore.ts
+++ b/packages/jotai-x/src/useHydrateStore.ts
@@ -3,16 +3,16 @@ import { useSetAtom } from 'jotai';
 import { useHydrateAtoms } from 'jotai/utils';
 
 import type {
+  SimpleWritableAtomRecord,
   UseHydrateAtoms,
   UseSyncAtoms,
-  WritableAtomRecord,
 } from './createAtomStore';
 
 /**
  * Hydrate atoms with initial values for SSR.
  */
 export const useHydrateStore = (
-  atoms: WritableAtomRecord<any>,
+  atoms: SimpleWritableAtomRecord<any>,
   initialValues: Parameters<UseHydrateAtoms<any>>[0],
   options: Parameters<UseHydrateAtoms<any>>[1] = {}
 ) => {
@@ -38,7 +38,7 @@ export const useHydrateStore = (
  * Update atoms with new values on changes.
  */
 export const useSyncStore = (
-  atoms: WritableAtomRecord<any>,
+  atoms: SimpleWritableAtomRecord<any>,
   values: any,
   { store }: Parameters<UseSyncAtoms<any>>[1] = {}
 ) => {

From 1754a2c496446ac3762c167fcf6b42d068edfe19 Mon Sep 17 00:00:00 2001
From: Joe Anderson <joe@mousetrapped.co.uk>
Date: Thu, 21 Dec 2023 18:13:46 +0000
Subject: [PATCH 11/20] Add placeholder changeset to enable release

---
 .changeset/fluffy-llamas-drop.md | 5 +++++
 1 file changed, 5 insertions(+)
 create mode 100644 .changeset/fluffy-llamas-drop.md

diff --git a/.changeset/fluffy-llamas-drop.md b/.changeset/fluffy-llamas-drop.md
new file mode 100644
index 0000000..8d43c92
--- /dev/null
+++ b/.changeset/fluffy-llamas-drop.md
@@ -0,0 +1,5 @@
+---
+'jotai-x': minor
+---
+
+TODO: Write changeset

From 69052bafff2370fcfc834c03689b72950a475843 Mon Sep 17 00:00:00 2001
From: Joe Anderson <joe@mousetrapped.co.uk>
Date: Thu, 21 Dec 2023 18:18:31 +0000
Subject: [PATCH 12/20] Cherry-pick release:next script

---
 package.json | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/package.json b/package.json
index 76d710a..5754635 100644
--- a/package.json
+++ b/package.json
@@ -18,6 +18,7 @@
     "typecheck": "yarn g:typecheck",
     "typedoc": "npx typedoc --options scripts/typedoc.json",
     "release": "yarn build && yarn changeset publish",
+    "release:next": "yarn build && yarn changeset publish --tag next",
     "g:brl": "turbo --filter \"./packages/**\" brl --no-daemon",
     "g:build": "turbo --filter \"./packages/**\" build --no-daemon",
     "g:build:watch": "yarn build:watch",
@@ -25,7 +26,6 @@
     "g:clean": "yarn clean:turbo && turbo --filter \"./packages/**\" clean --no-daemon",
     "g:lint": "turbo --filter \"./packages/**\" lint --no-daemon",
     "g:lint:fix": "turbo lint:fix --no-daemon",
-    "g:release:next": "yarn yarn build && yarn changeset publish --tag next",
     "g:test": "turbo --filter \"./packages/**\" test --no-daemon",
     "g:test:watch": "turbo --filter \"./packages/**\" test:watch --no-daemon",
     "g:test:cov": "yarn g:test --coverage",

From f89840c2d6c43c3e65d1634ffec6633ea47e1937 Mon Sep 17 00:00:00 2001
From: Joe Anderson <joe@mousetrapped.co.uk>
Date: Thu, 21 Dec 2023 18:46:43 +0000
Subject: [PATCH 13/20] Fix: Should not warn if store is undefined in getter

---
 packages/jotai-x/src/createAtomStore.ts | 9 ++++++---
 1 file changed, 6 insertions(+), 3 deletions(-)

diff --git a/packages/jotai-x/src/createAtomStore.ts b/packages/jotai-x/src/createAtomStore.ts
index 8d5f61d..7dcfb84 100644
--- a/packages/jotai-x/src/createAtomStore.ts
+++ b/packages/jotai-x/src/createAtomStore.ts
@@ -224,14 +224,17 @@ export const createAtomStore = <
   const setAtoms = {} as SetRecord<MyWritableStoreAtoms>;
   const useAtoms = {} as UseRecord<MyWritableStoreAtoms>;
 
-  const useStore = (optionsOrScope: UseAtomOptionsOrScope = {}) => {
+  const useStore = (
+    optionsOrScope: UseAtomOptionsOrScope = {},
+    warnIfUndefined = true
+  ) => {
     const { scope, store } = convertScopeShorthand(optionsOrScope);
-    const contextStore = useAtomStore(name, scope);
+    const contextStore = useAtomStore(name, scope, warnIfUndefined);
     return store ?? contextStore;
   };
 
   const useAtomValueWithStore: GetAtomFn = (atomConfig, optionsOrScope) => {
-    const store = useStore(optionsOrScope);
+    const store = useStore(optionsOrScope, false);
     const { delay = delayRoot } = convertScopeShorthand(optionsOrScope);
     return useAtomValue(atomConfig, { store, delay });
   };

From 901e4baa176b9415a5258a525511a3fbb5a46815 Mon Sep 17 00:00:00 2001
From: Joe Anderson <joe@mousetrapped.co.uk>
Date: Fri, 22 Dec 2023 09:03:05 +0000
Subject: [PATCH 14/20] Update changeset and docs

---
 .changeset/fluffy-llamas-dance.md | 13 ++++++++++
 .changeset/fluffy-llamas-drop.md  |  5 ----
 README.md                         | 43 +++++++++++++++++++++++++++----
 packages/jotai-x/README.md        | 43 +++++++++++++++++++++++++++----
 4 files changed, 89 insertions(+), 15 deletions(-)
 create mode 100644 .changeset/fluffy-llamas-dance.md
 delete mode 100644 .changeset/fluffy-llamas-drop.md

diff --git a/.changeset/fluffy-llamas-dance.md b/.changeset/fluffy-llamas-dance.md
new file mode 100644
index 0000000..a1b1f1c
--- /dev/null
+++ b/.changeset/fluffy-llamas-dance.md
@@ -0,0 +1,13 @@
+---
+'jotai-x': major
+---
+
+- Atoms can now be passed in the `initialState` argument to `createAtomStore`
+- Added an `extend` option to `createAtomStore` that lets you add derived atoms to the store
+- New accessors on `UseStoreApi`
+  - `useMyStore().store()` returns the `JotaiStore` for the current context, or undefined if no store exists
+  - `useMyStore().{get,set,use}.atom(someAtom)` accesses `someAtom` through the store
+- Remove exports for some internal types
+  - `GetRecord`
+  - `SetRecord`
+  - `UseRecord`
diff --git a/.changeset/fluffy-llamas-drop.md b/.changeset/fluffy-llamas-drop.md
deleted file mode 100644
index 8d43c92..0000000
--- a/.changeset/fluffy-llamas-drop.md
+++ /dev/null
@@ -1,5 +0,0 @@
----
-'jotai-x': minor
----
-
-TODO: Write changeset
diff --git a/README.md b/README.md
index 6079108..12f6c4e 100644
--- a/README.md
+++ b/README.md
@@ -54,26 +54,25 @@ createAtomStore<T extends object>(initialState: T, options?: CreateAtomStoreOpti
 The **`options`** object can include several properties to customize the behavior of your store:
 
 - **`name`**: A string representing the name of the store, which can be helpful for debugging or when working with multiple stores.
-- **`store`**: Allows specifying a [Jotai store](https://jotai.org/docs/core/store) if you want to use a custom one. Optional.
 - **`delay`**: If you need to introduce a delay in state updates, you can specify it here. Optional.
 - **`effect`**: A React component that can be used to run effects inside the provider. Optional.
+- **`extend`**: Extend the store with derived atoms based on the store state. Optional.
 
 #### Return Value
 
 The **`createAtomStore`** function returns an object (**`AtomStoreApi`**) containing the following properties and methods for interacting with the store:
 
 - **`use<Name>Store`**: 
-  - A function that returns the following objects: **`get`**, **`set`**, and **`use`**, where values are hooks for each state defined in the store. 
+  - A function that returns the following objects: **`get`**, **`set`**, **`use`** and **`store`**, where values are hooks for each state defined in the store.
   - **`get`**: Hooks for accessing a state within a component,  ensuring re-rendering when the state changes. See [useAtomValue](https://jotai.org/docs/core/use-atom#useatomvalue).
   - **`set`**: Hooks for setting a state within a component. See [useSetAtom](https://jotai.org/docs/core/use-atom#usesetatom).
   - **`use`**: Hooks for accessing and setting a state within a component, ensuring re-rendering when the state changes. See [useAtom](https://jotai.org/docs/core/use-atom).
+  - **`store`**: A hook to access the [JotaiStore](https://jotai.org/docs/core/store) for the current context.
   - Example: `const [element, setElement] = useElementStore().use.element()`
 - **`<Name>Provider`**:
   - The API includes dynamically generated provider components for each defined store. This allows  scoped state management within your application. More information in the next section.
 - **`<name>Store`**:
-  - Advanced API you generally don't need.
-  - **`atom`**: A hook for accessing state within a component,  ensuring re-rendering when the state changes. See [atom](https://jotai.org/docs/core/atom).
-  - **`extend`**: Extends the store with additional atoms. 
+  - **`atom`**: Access the atoms used by the store, including derived atoms defined using `extend`. See [atom](https://jotai.org/docs/core/atom).
   
 ### **Provider-Based Store Hydration and Synchronization**
 
@@ -86,6 +85,40 @@ The **`createAtomStore`** function returns an object (**`AtomStoreApi`**) contai
 
 JotaiX creates scoped providers, enabling more granular control over different segments of state within your application. `createAtomStore` sets up a context for each store, which can be scoped using the **`scope`** prop. This is particularly beneficial in complex applications where nested providers are needed.
 
+### Derived Atoms
+
+There are two ways of creating derived atoms from your JotaiX store.
+
+#### Derived Atoms Using `extend`
+
+Atoms defined using the `extend` option are made available in the same places as other values in the store.
+
+```ts
+const { useUserStore } = createAtomStore({
+  username: 'Alice',
+}, {
+  name: 'user',
+  extend: (atoms) => ({
+    intro: atom((get) => `My name is ${get(atoms.username)}`),
+  }),
+});
+
+const intro = useAppStore().get.intro();
+```
+
+#### Externally Defined Derived Atoms
+
+Derived atoms can also be defined externally by accessing the store's atoms through the `<name>Store` API. Externally defined atoms can be accessed through the store using the special `use<Name>Store().{get,set,use}.atom` hooks.
+
+```ts
+const { userStore, useUserStore } = createAtomStore({
+  username: 'Alice',
+}, { name: 'user' });
+
+const introAtom = atom((get) => `My name is ${get(userStore.atom.username)}`);
+const intro = useUserStore().get.atom(introAtom);
+```
+
 ### Example Usage
 
 #### 1. Create a store
diff --git a/packages/jotai-x/README.md b/packages/jotai-x/README.md
index 6079108..12f6c4e 100644
--- a/packages/jotai-x/README.md
+++ b/packages/jotai-x/README.md
@@ -54,26 +54,25 @@ createAtomStore<T extends object>(initialState: T, options?: CreateAtomStoreOpti
 The **`options`** object can include several properties to customize the behavior of your store:
 
 - **`name`**: A string representing the name of the store, which can be helpful for debugging or when working with multiple stores.
-- **`store`**: Allows specifying a [Jotai store](https://jotai.org/docs/core/store) if you want to use a custom one. Optional.
 - **`delay`**: If you need to introduce a delay in state updates, you can specify it here. Optional.
 - **`effect`**: A React component that can be used to run effects inside the provider. Optional.
+- **`extend`**: Extend the store with derived atoms based on the store state. Optional.
 
 #### Return Value
 
 The **`createAtomStore`** function returns an object (**`AtomStoreApi`**) containing the following properties and methods for interacting with the store:
 
 - **`use<Name>Store`**: 
-  - A function that returns the following objects: **`get`**, **`set`**, and **`use`**, where values are hooks for each state defined in the store. 
+  - A function that returns the following objects: **`get`**, **`set`**, **`use`** and **`store`**, where values are hooks for each state defined in the store.
   - **`get`**: Hooks for accessing a state within a component,  ensuring re-rendering when the state changes. See [useAtomValue](https://jotai.org/docs/core/use-atom#useatomvalue).
   - **`set`**: Hooks for setting a state within a component. See [useSetAtom](https://jotai.org/docs/core/use-atom#usesetatom).
   - **`use`**: Hooks for accessing and setting a state within a component, ensuring re-rendering when the state changes. See [useAtom](https://jotai.org/docs/core/use-atom).
+  - **`store`**: A hook to access the [JotaiStore](https://jotai.org/docs/core/store) for the current context.
   - Example: `const [element, setElement] = useElementStore().use.element()`
 - **`<Name>Provider`**:
   - The API includes dynamically generated provider components for each defined store. This allows  scoped state management within your application. More information in the next section.
 - **`<name>Store`**:
-  - Advanced API you generally don't need.
-  - **`atom`**: A hook for accessing state within a component,  ensuring re-rendering when the state changes. See [atom](https://jotai.org/docs/core/atom).
-  - **`extend`**: Extends the store with additional atoms. 
+  - **`atom`**: Access the atoms used by the store, including derived atoms defined using `extend`. See [atom](https://jotai.org/docs/core/atom).
   
 ### **Provider-Based Store Hydration and Synchronization**
 
@@ -86,6 +85,40 @@ The **`createAtomStore`** function returns an object (**`AtomStoreApi`**) contai
 
 JotaiX creates scoped providers, enabling more granular control over different segments of state within your application. `createAtomStore` sets up a context for each store, which can be scoped using the **`scope`** prop. This is particularly beneficial in complex applications where nested providers are needed.
 
+### Derived Atoms
+
+There are two ways of creating derived atoms from your JotaiX store.
+
+#### Derived Atoms Using `extend`
+
+Atoms defined using the `extend` option are made available in the same places as other values in the store.
+
+```ts
+const { useUserStore } = createAtomStore({
+  username: 'Alice',
+}, {
+  name: 'user',
+  extend: (atoms) => ({
+    intro: atom((get) => `My name is ${get(atoms.username)}`),
+  }),
+});
+
+const intro = useAppStore().get.intro();
+```
+
+#### Externally Defined Derived Atoms
+
+Derived atoms can also be defined externally by accessing the store's atoms through the `<name>Store` API. Externally defined atoms can be accessed through the store using the special `use<Name>Store().{get,set,use}.atom` hooks.
+
+```ts
+const { userStore, useUserStore } = createAtomStore({
+  username: 'Alice',
+}, { name: 'user' });
+
+const introAtom = atom((get) => `My name is ${get(userStore.atom.username)}`);
+const intro = useUserStore().get.atom(introAtom);
+```
+
 ### Example Usage
 
 #### 1. Create a store

From 6ac8e319484f50334c9ebc171c09f2332024977f Mon Sep 17 00:00:00 2001
From: Ziad Beyens <zbeyens@udecode.io>
Date: Fri, 22 Dec 2023 13:07:30 +0100
Subject: [PATCH 15/20] Update fluffy-llamas-dance.md

---
 .changeset/fluffy-llamas-dance.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.changeset/fluffy-llamas-dance.md b/.changeset/fluffy-llamas-dance.md
index a1b1f1c..a4f4746 100644
--- a/.changeset/fluffy-llamas-dance.md
+++ b/.changeset/fluffy-llamas-dance.md
@@ -1,5 +1,5 @@
 ---
-'jotai-x': major
+'jotai-x': minor
 ---
 
 - Atoms can now be passed in the `initialState` argument to `createAtomStore`

From 5b28a6832515a85e390649daa24f6e126ba5ac53 Mon Sep 17 00:00:00 2001
From: Ziad Beyens <zbeyens@udecode.io>
Date: Fri, 22 Dec 2023 13:11:07 +0100
Subject: [PATCH 16/20] Update fluffy-llamas-dance.md

---
 .changeset/fluffy-llamas-dance.md | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/.changeset/fluffy-llamas-dance.md b/.changeset/fluffy-llamas-dance.md
index a4f4746..13b6b45 100644
--- a/.changeset/fluffy-llamas-dance.md
+++ b/.changeset/fluffy-llamas-dance.md
@@ -2,12 +2,12 @@
 'jotai-x': minor
 ---
 
-- Atoms can now be passed in the `initialState` argument to `createAtomStore`
+- Atoms other than `atom` can now be passed in the `initialState` argument to `createAtomStore`. Primitive values use `atom` by default
 - Added an `extend` option to `createAtomStore` that lets you add derived atoms to the store
 - New accessors on `UseStoreApi`
   - `useMyStore().store()` returns the `JotaiStore` for the current context, or undefined if no store exists
   - `useMyStore().{get,set,use}.atom(someAtom)` accesses `someAtom` through the store
-- Remove exports for some internal types
+- Types: remove exports for some internal types
   - `GetRecord`
   - `SetRecord`
   - `UseRecord`

From 3357a37f3efe5e3e4df38045408d2dda598bf0f8 Mon Sep 17 00:00:00 2001
From: Ziad Beyens <zbeyens@udecode.io>
Date: Fri, 22 Dec 2023 13:11:26 +0100
Subject: [PATCH 17/20] Update README.md

---
 README.md | 5 -----
 1 file changed, 5 deletions(-)

diff --git a/README.md b/README.md
index 12f6c4e..5f0a691 100644
--- a/README.md
+++ b/README.md
@@ -228,11 +228,6 @@ const Component = () => {
 
 ## Contributing
 
-### Roadmap
-
-- [ ] Support other atoms like `atomWithStorage`
-- [ ] Improve `extend` API to be more modular.
-
 ### Ideas and discussions
 
 [Discussions](https://github.com/udecode/jotai-x/discussions) is the best

From 2b5238fc8a639faf279226f702ead944dd73ee9c Mon Sep 17 00:00:00 2001
From: Ziad Beyens <zbeyens@udecode.io>
Date: Fri, 22 Dec 2023 13:11:44 +0100
Subject: [PATCH 18/20] Update README.md

---
 packages/jotai-x/README.md | 5 -----
 1 file changed, 5 deletions(-)

diff --git a/packages/jotai-x/README.md b/packages/jotai-x/README.md
index 12f6c4e..5f0a691 100644
--- a/packages/jotai-x/README.md
+++ b/packages/jotai-x/README.md
@@ -228,11 +228,6 @@ const Component = () => {
 
 ## Contributing
 
-### Roadmap
-
-- [ ] Support other atoms like `atomWithStorage`
-- [ ] Improve `extend` API to be more modular.
-
 ### Ideas and discussions
 
 [Discussions](https://github.com/udecode/jotai-x/discussions) is the best

From 19b63d9a81a466d63808017d9402018797356708 Mon Sep 17 00:00:00 2001
From: Joe Anderson <joe@mousetrapped.co.uk>
Date: Fri, 22 Dec 2023 13:55:13 +0000
Subject: [PATCH 19/20] Hide `{ fn: ... }` workaround behind the scenes

---
 packages/jotai-x/src/createAtomStore.spec.tsx | 118 +++++++++++++++++-
 packages/jotai-x/src/createAtomStore.ts       |  26 +++-
 packages/jotai-x/src/useHydrateStore.ts       |  11 +-
 3 files changed, 142 insertions(+), 13 deletions(-)

diff --git a/packages/jotai-x/src/createAtomStore.spec.tsx b/packages/jotai-x/src/createAtomStore.spec.tsx
index 3312718..b519f32 100644
--- a/packages/jotai-x/src/createAtomStore.spec.tsx
+++ b/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', () => {
diff --git a/packages/jotai-x/src/createAtomStore.ts b/packages/jotai-x/src/createAtomStore.ts
index 7dcfb84..89e590b 100644
--- a/packages/jotai-x/src/createAtomStore.ts
+++ b/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;
 
diff --git a/packages/jotai-x/src/useHydrateStore.ts b/packages/jotai-x/src/useHydrateStore.ts
index bb45d45..8af7ffe 100644
--- a/packages/jotai-x/src/useHydrateStore.ts
+++ b/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]);
   }

From 1ace216a41a427a2779dad9df852b08a79c3be8b Mon Sep 17 00:00:00 2001
From: Joe Anderson <joe@mousetrapped.co.uk>
Date: Fri, 22 Dec 2023 14:28:56 +0000
Subject: [PATCH 20/20] Refactor `atomWithFn` into util file

---
 packages/jotai-x/src/atomWithFn.ts      | 32 +++++++++++++++++++++++++
 packages/jotai-x/src/createAtomStore.ts | 27 ++-------------------
 packages/jotai-x/src/index.ts           |  1 +
 3 files changed, 35 insertions(+), 25 deletions(-)
 create mode 100644 packages/jotai-x/src/atomWithFn.ts

diff --git a/packages/jotai-x/src/atomWithFn.ts b/packages/jotai-x/src/atomWithFn.ts
new file mode 100644
index 0000000..5c759a1
--- /dev/null
+++ b/packages/jotai-x/src/atomWithFn.ts
@@ -0,0 +1,32 @@
+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;
+
+/**
+ * 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.
+ */
+export const atomWithFn = <T>(initialValue: T): WritableAtom<T, [T], void> => {
+  const baseAtom = atom(wrapFn(initialValue));
+
+  return atom(
+    (get) => unwrapFn(get(baseAtom)) as T,
+    (_get, set, value) => set(baseAtom, wrapFn(value))
+  );
+};
diff --git a/packages/jotai-x/src/createAtomStore.ts b/packages/jotai-x/src/createAtomStore.ts
index 89e590b..8f5f769 100644
--- a/packages/jotai-x/src/createAtomStore.ts
+++ b/packages/jotai-x/src/createAtomStore.ts
@@ -1,6 +1,7 @@
-import { atom, useAtom, useAtomValue, useSetAtom } from 'jotai';
+import { useAtom, useAtomValue, useSetAtom } from 'jotai';
 import { useHydrateAtoms } from 'jotai/utils';
 
+import { atomWithFn } from './atomWithFn';
 import { createAtomProvider, useAtomStore } from './createAtomProvider';
 
 import type { ProviderProps } from './createAtomProvider';
@@ -17,30 +18,6 @@ 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
diff --git a/packages/jotai-x/src/index.ts b/packages/jotai-x/src/index.ts
index f63fe17..b9d1112 100644
--- a/packages/jotai-x/src/index.ts
+++ b/packages/jotai-x/src/index.ts
@@ -3,6 +3,7 @@
  */
 
 export * from './atomProvider';
+export * from './atomWithFn';
 export * from './createAtomProvider';
 export * from './createAtomStore';
 export * from './useHydrateStore';