Feature: useHydrateAtoms

.size-snapshot.json

@@ -1,8 +1,8 @@
 {
  "index.js": {
- "bundled": 21267,
- "minified": 10148,
- "gzipped": 3367,
+ "bundled": 21335,
+ "minified": 10169,
+ "gzipped": 3364,
  "treeshaked": {
  "rollup": {
  "code": 14,
@@ -14,22 +14,22 @@
  }
  },
  "utils.js": {
- "bundled": 16025,
- "minified": 7685,
- "gzipped": 2838,
+ "bundled": 16550,
+ "minified": 7958,
+ "gzipped": 2921,
  "treeshaked": {
  "rollup": {
  "code": 28,
  "import_statements": 28
  },
  "webpack": {
- "code": 1305
+ "code": 1314
  }
  }
  },
  "devtools.js": {
- "bundled": 19122,
- "minified": 9537,
+ "bundled": 19190,
+ "minified": 9558,
  "gzipped": 3288,
  "treeshaked": {
  "rollup": {

src/core/Provider.ts

@@ -81,7 +81,7 @@ export const subscribeDebugStore = (
 // We keep a reference to the atoms in Provider's registeredAtoms in dev mode,
 // so atoms aren't garbage collected by the WeakMap of mounted atoms
 const useDebugState = (store: StoreForDevelopment) => {
- const debugMutableSource = store[3]
+ const debugMutableSource = store[4]
  const [state, atoms] = useMutableSource(
  debugMutableSource,
  getDebugStateAndAtoms,

src/core/contexts.ts

@@ -14,7 +14,9 @@ const createStoreForProduction = (
  atom: WritableAtom<Value, Update>,
  update: Update
  ) => writeAtom(state, atom, update)
- return [stateMutableSource, updateAtom, commitCallback] as const
+ const restore = (values: Iterable<readonly [Atom<unknown>, unknown]>) =>
+ restoreAtoms(state, values)
+ return [stateMutableSource, updateAtom, commitCallback, restore] as const
 }
 
 const createStoreForDevelopment = (
@@ -54,8 +56,8 @@ const createStoreForDevelopment = (
  stateMutableSource,
  updateAtom,
  commitCallback,
- debugMutableSource,
  restore,
+ debugMutableSource,
  ] as const
 }
 

src/devtools/useAtomsSnapshot.ts

@@ -12,7 +12,7 @@ type AtomsSnapshot = Map<Atom<unknown>, unknown>
 
 export function useAtomsSnapshot(scope?: Scope): AtomsSnapshot {
  const StoreContext = getStoreContext(scope)
- const debugMutableSource = useContext(StoreContext)[3]
+ const debugMutableSource = useContext(StoreContext)[4]
 
  if (debugMutableSource === undefined) {
  throw Error('useAtomsSnapshot can only be used in dev mode.')

src/devtools/useGotoAtomsSnapshot.ts

@@ -11,7 +11,7 @@ export function useGotoAtomsSnapshot(scope?: Scope) {
  if (!isDevStore(store)) {
  throw new Error('useGotoAtomsSnapshot can only be used in dev mode.')
  }
- const restore = store[4]
+ const restore = store[3]
  return useCallback(
  (values: Parameters<typeof restore>[0]) => {
  for (const [atom] of values) {

src/utils.ts

@@ -18,3 +18,4 @@ export {
  createJSONStorage,
 } from './utils/atomWithStorage'
 export { atomWithObservable } from './utils/atomWithObservable'
+export { useHydrateAtoms } from './utils/useHydrateAtoms'

src/utils/useHydrateAtoms.ts

@@ -0,0 +1,25 @@
+import { useContext, useMemo } from 'react'
+import { SECRET_INTERNAL_getStoreContext as getStoreContext } from 'jotai'
+import type { Atom, Scope } from '../core/atom'
+
+export function useHydrateAtoms(
+ values: Iterable<readonly [Atom<unknown> & { hydrated?: Symbol }, unknown]>,
+ scope?: Scope
+) {
+ const StoreContext = getStoreContext(scope)
+ const restoreAtoms = useContext(StoreContext)[3]
+
+ useMemo(() => {
+ const tuplesToRestore = []
+ for (const tuple of values) {
+ const atom = tuple[0]
+ if (atom.hydrated !== hydratedSymbol) {
+ tuplesToRestore.push(tuple)
+ atom.hydrated = hydratedSymbol
+ }
+ }
+ restoreAtoms(tuplesToRestore)
+ }, [values, restoreAtoms])
+}
+
+const hydratedSymbol = Symbol()

tests/utils/useHydrateAtoms.test.tsx

@@ -0,0 +1,216 @@
+import { FC, useRef } from 'react'
+import { fireEvent, render } from '@testing-library/react'
+import { atom, useAtom } from '../../src/index'
+import { useHydrateAtoms } from '../../src/utils'
+import { getTestProvider } from '../testUtils'
+
+const Provider = getTestProvider()
+
+it('useHydrateAtoms should only hydrate on first render', async () => {
+ const countAtom = atom(0)
+
+ const Counter: FC<{ initialCount: number }> = ({ initialCount }) => {
+ useHydrateAtoms([[countAtom, initialCount]])
+ const [countValue, setCount] = useAtom(countAtom)
+
+ return (
+ <>
+ <div>count: {countValue}</div>
+ <button onClick={() => setCount((count) => count + 1)}>dispatch</button>
+ </>
+ )
+ }
+ const { findByText, getByText, rerender } = render(
+ <Provider>
+ <Counter initialCount={42} />
+ </Provider>
+ )
+
+ await findByText('count: 42')
+ fireEvent.click(getByText('dispatch'))
+ await findByText('count: 43')
+
+ rerender(
+ <Provider>
+ <Counter initialCount={65} />
+ </Provider>
+ )
+ await findByText('count: 43')
+})
+
+it('useHydrateAtoms should not trigger unnessesary rerenders', async () => {
+ const countAtom = atom(0)
+
+ const Counter: FC<{ initialCount: number }> = ({ initialCount }) => {
+ useHydrateAtoms([[countAtom, initialCount]])
+ const [countValue, setCount] = useAtom(countAtom)
+ const renderCount = useRef(0)
+ ++renderCount.current
+ return (
+ <>
+ <div>renders: {renderCount.current}</div>
+ <div>count: {countValue}</div>
+ <button onClick={() => setCount((count) => count + 1)}>dispatch</button>
+ </>
+ )
+ }
+
+ const { findByText, getByText } = render(
+ <Provider>
+ <Counter initialCount={42} />
+ </Provider>
+ )
+
+ await findByText('count: 42')
+ await findByText('renders: 1')
+ fireEvent.click(getByText('dispatch'))
+ await findByText('count: 43')
+ await findByText('renders: 2')
+})
+
+it('useHydrateAtoms should work with derived atoms', async () => {
+ const countAtom = atom(0)
+ const doubleAtom = atom((get) => get(countAtom) * 2)
+
+ const Counter: FC<{ initialCount: number }> = ({ initialCount }) => {
+ useHydrateAtoms([[countAtom, initialCount]])
+ const [countValue, setCount] = useAtom(countAtom)
+ const [doubleCount] = useAtom(doubleAtom)
+ return (
+ <>
+ <div>count: {countValue}</div>
+ <div>doubleCount: {doubleCount}</div>
+ <button onClick={() => setCount((count) => count + 1)}>dispatch</button>
+ </>
+ )
+ }
+
+ const { findByText, getByText } = render(
+ <Provider>
+ <Counter initialCount={42} />
+ </Provider>
+ )
+
+ await findByText('count: 42')
+ await findByText('doubleCount: 84')
+ fireEvent.click(getByText('dispatch'))
+ await findByText('count: 43')
+ await findByText('doubleCount: 86')
+})
+
+it('useHydrateAtoms can only restore an atom once', async () => {
+ const countAtom = atom(0)
+
+ const Counter: FC<{ initialCount: number }> = ({ initialCount }) => {
+ useHydrateAtoms([[countAtom, initialCount]])
+ const [countValue, setCount] = useAtom(countAtom)
+
+ return (
+ <>
+ <div>count: {countValue}</div>
+ <button onClick={() => setCount((count) => count + 1)}>dispatch</button>
+ </>
+ )
+ }
+ const Counter2: FC<{ count: number }> = ({ count }) => {
+ useHydrateAtoms([[countAtom, count]])
+ const [countValue, setCount] = useAtom(countAtom)
+
+ return (
+ <>
+ <div>count: {countValue}</div>
+ <button onClick={() => setCount((count) => count + 1)}>dispatch</button>
+ </>
+ )
+ }
+ const { findByText, getByText, rerender } = render(
+ <Provider>
+ <Counter initialCount={42} />
+ </Provider>
+ )
+
+ await findByText('count: 42')
+ fireEvent.click(getByText('dispatch'))
+ await findByText('count: 43')
+
+ rerender(
+ <Provider>
+ <Counter2 count={65} />
+ </Provider>
+ )
+
+ await findByText('count: 43')
+ fireEvent.click(getByText('dispatch'))
+ await findByText('count: 44')
+})
+
+it('useHydrateAtoms can only restore an atom once', async () => {
+ const countAtom = atom(0)
+
+ const Counter: FC<{ initialCount: number }> = ({ initialCount }) => {
+ useHydrateAtoms([[countAtom, initialCount]])
+ const [countValue, setCount] = useAtom(countAtom)
+
+ return (
+ <>
+ <div>count: {countValue}</div>
+ <button onClick={() => setCount((count) => count + 1)}>dispatch</button>
+ </>
+ )
+ }
+ const Counter2: FC<{ count: number }> = ({ count }) => {
+ useHydrateAtoms([[countAtom, count]])
+ const [countValue, setCount] = useAtom(countAtom)
+
+ return (
+ <>
+ <div>count: {countValue}</div>
+ <button onClick={() => setCount((count) => count + 1)}>dispatch</button>
+ </>
+ )
+ }
+ const { findByText, getByText, rerender } = render(
+ <Provider>
+ <Counter initialCount={42} />
+ </Provider>
+ )
+
+ await findByText('count: 42')
+ fireEvent.click(getByText('dispatch'))
+ await findByText('count: 43')
+
+ rerender(
+ <Provider>
+ <Counter2 count={65} />
+ </Provider>
+ )
+
+ await findByText('count: 43')
+ fireEvent.click(getByText('dispatch'))
+ await findByText('count: 44')
+})
+
+it('useHydrateAtoms should respect onMount', async () => {
+ const countAtom = atom(0)
+ const onMountFn = jest.fn()
+ countAtom.onMount = onMountFn
+
+ const Counter: FC<{ initialCount: number }> = ({ initialCount }) => {
+ useHydrateAtoms([[countAtom, initialCount]])
+ const [countValue] = useAtom(countAtom)
+
+ return (
+ <>
+ <div>count: {countValue}</div>
+ </>
+ )
+ }
+ const { findByText } = render(
+ <Provider>
+ <Counter initialCount={42} />
+ </Provider>
+ )
+
+ await findByText('count: 42')
+ expect(onMountFn).toBeCalledTimes(1)
+})