Add APIs to support working with custom and derived atoms

.changeset/fluffy-llamas-dance.md

@@ -0,0 +1,13 @@
+---
+'jotai-x': minor
+---
+
+- 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
+- Types: remove exports for some internal types
+  - `GetRecord`
+  - `SetRecord`
+  - `UseRecord`

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
@@ -195,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

package.json

@@ -18,14 +18,14 @@
     "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",
     "g:changeset": "changeset",
     "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",

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
@@ -195,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

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))
+  );
+};

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, 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: AtomRecord<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: AtomRecord<T>,
+  atoms: SimpleWritableAtomRecord<T>,
   options: { effect?: FC } = {}
 ) => {
   const Effect = options.effect;