From b52928900d95d76d6301b7272d41f7c6ce7bb7d5 Mon Sep 17 00:00:00 2001 From: Iwo Plaza Date: Mon, 25 Mar 2024 19:47:05 +0100 Subject: [PATCH 1/8] feat(utils): atomWithLazy for lazily initialized primitive atoms. --- src/vanilla/utils.ts | 1 + src/vanilla/utils/atomWithLazy.ts | 11 +++++++ tests/vanilla/utils/atomWithLazy.test.ts | 41 ++++++++++++++++++++++++ 3 files changed, 53 insertions(+) create mode 100644 src/vanilla/utils/atomWithLazy.ts create mode 100644 tests/vanilla/utils/atomWithLazy.test.ts diff --git a/src/vanilla/utils.ts b/src/vanilla/utils.ts index 290e01f61b..f24e5535a3 100644 --- a/src/vanilla/utils.ts +++ b/src/vanilla/utils.ts @@ -15,3 +15,4 @@ export { atomWithObservable } from './utils/atomWithObservable.ts' export { loadable } from './utils/loadable.ts' export { unwrap } from './utils/unwrap.ts' export { atomWithRefresh } from './utils/atomWithRefresh.ts' +export { atomWithLazy } from './utils/atomWithLazy.ts' diff --git a/src/vanilla/utils/atomWithLazy.ts b/src/vanilla/utils/atomWithLazy.ts new file mode 100644 index 0000000000..4c4bd43442 --- /dev/null +++ b/src/vanilla/utils/atomWithLazy.ts @@ -0,0 +1,11 @@ +import { atom } from '../../vanilla.ts' +import type { SetStateAction } from '../../vanilla.ts' + +export function atomWithLazy(makeInitial: () => Value) { + const wrappedAtom = atom(() => atom(makeInitial())) + + return atom( + (get) => get(get(wrappedAtom)), + (get, set, value: SetStateAction) => set(get(wrappedAtom), value), + ) +} diff --git a/tests/vanilla/utils/atomWithLazy.test.ts b/tests/vanilla/utils/atomWithLazy.test.ts new file mode 100644 index 0000000000..dcce5197f5 --- /dev/null +++ b/tests/vanilla/utils/atomWithLazy.test.ts @@ -0,0 +1,41 @@ +import { expect, it, vi } from 'vitest' +import { createStore } from 'jotai/vanilla' +import { atomWithLazy } from 'jotai/vanilla/utils' + +it('initializes on first store get', async () => { + const storeA = createStore() + const storeB = createStore() + + let externalState = 'first' + const initializer = vi.fn(() => externalState) + const anAtom = atomWithLazy(initializer) + + expect(initializer).not.toHaveBeenCalled() + expect(storeA.get(anAtom)).toEqual('first') + expect(initializer).toHaveBeenCalledOnce() + + externalState = 'second' + + expect(storeA.get(anAtom)).toEqual('first') + expect(initializer).toHaveBeenCalledOnce() + expect(storeB.get(anAtom)).toEqual('second') + expect(initializer).toHaveBeenCalledTimes(2) +}) + +it('is writable', async () => { + const store = createStore() + const anAtom = atomWithLazy(() => 0) + + store.set(anAtom, 123) + + expect(store.get(anAtom)).toEqual(123) +}) + +it('should work with a set state action', async () => { + const store = createStore() + const anAtom = atomWithLazy(() => 4) + + store.set(anAtom, (prev) => prev * prev) + + expect(store.get(anAtom)).toEqual(16) +}) From 85045536de02d12962eb9ea35339c441dba1b463 Mon Sep 17 00:00:00 2001 From: Iwo Plaza Date: Mon, 25 Mar 2024 19:52:22 +0100 Subject: [PATCH 2/8] tweak test for older ts versions --- tests/vanilla/utils/atomWithLazy.test.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/vanilla/utils/atomWithLazy.test.ts b/tests/vanilla/utils/atomWithLazy.test.ts index dcce5197f5..43d53b895e 100644 --- a/tests/vanilla/utils/atomWithLazy.test.ts +++ b/tests/vanilla/utils/atomWithLazy.test.ts @@ -35,7 +35,7 @@ it('should work with a set state action', async () => { const store = createStore() const anAtom = atomWithLazy(() => 4) - store.set(anAtom, (prev) => prev * prev) + store.set(anAtom, (prev: number) => prev * prev) expect(store.get(anAtom)).toEqual(16) }) From 193bc1fb6d09817b7d24881d8b47d9c2811ee929 Mon Sep 17 00:00:00 2001 From: Iwo Plaza Date: Wed, 27 Mar 2024 16:19:26 +0100 Subject: [PATCH 3/8] add `unstable_is` hack for use in `jotai-store` --- src/vanilla/utils/atomWithLazy.ts | 15 ++++++++++++--- 1 file changed, 12 insertions(+), 3 deletions(-) diff --git a/src/vanilla/utils/atomWithLazy.ts b/src/vanilla/utils/atomWithLazy.ts index 4c4bd43442..e5457a5cc1 100644 --- a/src/vanilla/utils/atomWithLazy.ts +++ b/src/vanilla/utils/atomWithLazy.ts @@ -1,11 +1,20 @@ import { atom } from '../../vanilla.ts' -import type { SetStateAction } from '../../vanilla.ts' +import type { Atom, PrimitiveAtom, SetStateAction } from '../../vanilla.ts' export function atomWithLazy(makeInitial: () => Value) { - const wrappedAtom = atom(() => atom(makeInitial())) + const wrappedAtom: Atom> & { init?: Atom } = + atom(() => atom(makeInitial())) - return atom( + const proxyAtom = atom( (get) => get(get(wrappedAtom)), (get, set, value: SetStateAction) => set(get(wrappedAtom), value), ) + + wrappedAtom.init = atom(undefined) + // when writing to wrappedAtom through proxyAtom, the store + // thinks we are actually storing the value of `proxyAtom`. + wrappedAtom.unstable_is = (a) => + a.unstable_is ? a.unstable_is(proxyAtom) : a === proxyAtom + + return proxyAtom } From 94c47dc702184f86309bae25179d2e94cad66181 Mon Sep 17 00:00:00 2001 From: Iwo Plaza Date: Wed, 27 Mar 2024 22:06:01 +0100 Subject: [PATCH 4/8] add documentation for atomWithLazy --- docs/utilities/lazy.mdx | 95 +++++++++++++++++++++++++++++++++++++++++ 1 file changed, 95 insertions(+) create mode 100644 docs/utilities/lazy.mdx diff --git a/docs/utilities/lazy.mdx b/docs/utilities/lazy.mdx new file mode 100644 index 0000000000..5a18b7fe71 --- /dev/null +++ b/docs/utilities/lazy.mdx @@ -0,0 +1,95 @@ +--- +title: Lazy +nav: 3.03 +keywords: lazy,initialize,init,loading +--- + +When defining primitive atoms, their initial value has to be bound at definition time. +If creating that initial value is computationally expensive, or the value is not accessible during definition, +it would be best to postpone the atom's initialization until its [first use in the store](#using-multiple-stores). + +```jsx +const imageDataAtom = atom(initializeExpensiveImage()) // 1) has to be computed here + +function Home() { + ... +} + +function ImageEditor() { + // 2) used only in this route + const [imageData, setImageData] = useAtom(imageDataAtom); + ... +} + +function App() { + return ( + + + + + ) +} +``` + +## atomWithLazy + +Ref: https://github.com/pmndrs/jotai/pull/2465 + +We can use `atomWithLazy` to create a primitive atom whose initial value will be computed at [first use in the store](#using-multiple-stores). +After initialization, it will behave like a regular primitive atom (can be written to). + +### Usage + +```jsx +import { atomWithLazy } from 'jotai/utils' + +const imageDataAtom = atomWithLazy(() => initializeExpensiveImage()) + +function Home() { + ... +} + +function ImageEditor() { + // only gets initialized if user goes to `/edit`. + const [imageData, setImageData] = useAtom(imageDataAtom); + ... +} + +function App() { + return ( + + + + + ) +} +``` + +### Using multiple stores + +Since each store is its separate universe, the initial value will be recreated exactly once per store (unless using something like `jotai-scope`, which fractures a store into smaller universes). + +```ts +type RGB = [number, number, number]; + +function randomRGB(): RGB { + ... +} + +const lift = (value: number) => ([r, g, b]: RGB) => { + return [r + value, g + value, b + value] +} + +const colorAtom = lazyAtom(randomRGB) + +let store = createStore(); + +console.log(store.get(colorAtom)) // [0, 36, 128] +store.set(colorAtom, lift(8)) +console.log(store.get(colorAtom)) // [8, 44, 136] + +// recreating store, sometimes done when logging out or resetting the app in some way +store = createStore(); + +console.log(store.get(colorAtom)) // [255, 12, 46] -- a new random color +``` From 2396718fd61c54620e2628744f30c81c58c62566 Mon Sep 17 00:00:00 2001 From: Iwo Plaza Date: Wed, 27 Mar 2024 22:08:27 +0100 Subject: [PATCH 5/8] update lazy docs --- docs/utilities/lazy.mdx | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/docs/utilities/lazy.mdx b/docs/utilities/lazy.mdx index 5a18b7fe71..b5f429879e 100644 --- a/docs/utilities/lazy.mdx +++ b/docs/utilities/lazy.mdx @@ -43,7 +43,8 @@ After initialization, it will behave like a regular primitive atom (can be writt ```jsx import { atomWithLazy } from 'jotai/utils' -const imageDataAtom = atomWithLazy(() => initializeExpensiveImage()) +// passing the initialization function +const imageDataAtom = atomWithLazy(initializeExpensiveImage) function Home() { ... From 90bc12d223a7709b740bc02a45b3fc8e83dea3a6 Mon Sep 17 00:00:00 2001 From: Iwo Plaza Date: Wed, 27 Mar 2024 22:12:46 +0100 Subject: [PATCH 6/8] remove semicolons in lazy docs --- docs/utilities/lazy.mdx | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/docs/utilities/lazy.mdx b/docs/utilities/lazy.mdx index b5f429879e..8ff33440d9 100644 --- a/docs/utilities/lazy.mdx +++ b/docs/utilities/lazy.mdx @@ -83,14 +83,14 @@ const lift = (value: number) => ([r, g, b]: RGB) => { const colorAtom = lazyAtom(randomRGB) -let store = createStore(); +let store = createStore() console.log(store.get(colorAtom)) // [0, 36, 128] store.set(colorAtom, lift(8)) console.log(store.get(colorAtom)) // [8, 44, 136] // recreating store, sometimes done when logging out or resetting the app in some way -store = createStore(); +store = createStore() console.log(store.get(colorAtom)) // [255, 12, 46] -- a new random color ``` From 240b13de864aa5de7bbc9d90f72980edd4b97ba1 Mon Sep 17 00:00:00 2001 From: Iwo Plaza Date: Fri, 29 Mar 2024 20:12:38 +0100 Subject: [PATCH 7/8] add more elegant implementation for atomWithLazy --- src/vanilla/utils/atomWithLazy.ts | 22 ++++++---------------- 1 file changed, 6 insertions(+), 16 deletions(-) diff --git a/src/vanilla/utils/atomWithLazy.ts b/src/vanilla/utils/atomWithLazy.ts index e5457a5cc1..234e447fd7 100644 --- a/src/vanilla/utils/atomWithLazy.ts +++ b/src/vanilla/utils/atomWithLazy.ts @@ -1,20 +1,10 @@ import { atom } from '../../vanilla.ts' -import type { Atom, PrimitiveAtom, SetStateAction } from '../../vanilla.ts' export function atomWithLazy(makeInitial: () => Value) { - const wrappedAtom: Atom> & { init?: Atom } = - atom(() => atom(makeInitial())) - - const proxyAtom = atom( - (get) => get(get(wrappedAtom)), - (get, set, value: SetStateAction) => set(get(wrappedAtom), value), - ) - - wrappedAtom.init = atom(undefined) - // when writing to wrappedAtom through proxyAtom, the store - // thinks we are actually storing the value of `proxyAtom`. - wrappedAtom.unstable_is = (a) => - a.unstable_is ? a.unstable_is(proxyAtom) : a === proxyAtom - - return proxyAtom + return { + ...atom(undefined as unknown as Value), + get init() { + return makeInitial() + }, + } } From 0249b0e96a5c350c074816d9b6027b6e55b77119 Mon Sep 17 00:00:00 2001 From: Iwo Plaza Date: Fri, 29 Mar 2024 20:32:10 +0100 Subject: [PATCH 8/8] clear type declarations --- src/vanilla/utils/atomWithLazy.ts | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/src/vanilla/utils/atomWithLazy.ts b/src/vanilla/utils/atomWithLazy.ts index 234e447fd7..208be7e547 100644 --- a/src/vanilla/utils/atomWithLazy.ts +++ b/src/vanilla/utils/atomWithLazy.ts @@ -1,6 +1,8 @@ -import { atom } from '../../vanilla.ts' +import { PrimitiveAtom, atom } from '../../vanilla.ts' -export function atomWithLazy(makeInitial: () => Value) { +export function atomWithLazy( + makeInitial: () => Value, +): PrimitiveAtom { return { ...atom(undefined as unknown as Value), get init() {