feat(utils): atomWithLazy for lazily initialized primitive atoms. (#2465)

* feat(utils): atomWithLazy for lazily initialized primitive atoms.

* tweak test for older ts versions

* add `unstable_is` hack for use in `jotai-store`

* add documentation for atomWithLazy

* update lazy docs

* remove semicolons in lazy docs

* add more elegant implementation for atomWithLazy

* clear type declarations

---------

Co-authored-by: Daishi Kato <[email protected]>

Author: iwoplaza

docs/utilities/lazy.mdx

@@ -0,0 +1,96 @@
+---
+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 (
+    <Router>
+      <Route path="/" component={Home} />
+      <Route path="/edit" component={ImageEditor} />
+    </Router>
+  )
+}
+```
+
+## 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'
+
+// passing the initialization function
+const imageDataAtom = atomWithLazy(initializeExpensiveImage)
+
+function Home() {
+  ...
+}
+
+function ImageEditor() {
+  // only gets initialized if user goes to `/edit`.
+  const [imageData, setImageData] = useAtom(imageDataAtom);
+  ...
+}
+
+function App() {
+  return (
+    <Router>
+      <Route path="/" component={Home} />
+      <Route path="/edit" component={ImageEditor} />
+    </Router>
+  )
+}
+```
+
+### 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
+```

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'

src/vanilla/utils/atomWithLazy.ts

@@ -0,0 +1,12 @@
+import { PrimitiveAtom, atom } from '../../vanilla.ts'
+
+export function atomWithLazy<Value>(
+  makeInitial: () => Value,
+): PrimitiveAtom<Value> {
+  return {
+    ...atom(undefined as unknown as Value),
+    get init() {
+      return makeInitial()
+    },
+  }
+}

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: number) => prev * prev)
+
+  expect(store.get(anAtom)).toEqual(16)
+})