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

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,11 @@
+import { atom } from '../../vanilla.ts'
+import type { SetStateAction } from '../../vanilla.ts'
+
+export function atomWithLazy<Value>(makeInitial: () => Value) {
+  const wrappedAtom = atom(() => atom(makeInitial()))
+
+  return atom(
+    (get) => get(get(wrappedAtom)),
+    (get, set, value: SetStateAction<Value>) => set(get(wrappedAtom), value),
+  )
+}

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