Feature: useHydrateAtoms (#637)

* wip

* useHydrateAtoms wip

* Add first test

* More test cases

* Simplify implementation

* Update size-snapshot

* Revert to original implementation

* Optimize hook and add tests

* Fix types

* Fix types

* Use symbol property for hydrated check

* Allow hydration on atom once per store

* Add scope test

* Use new useAtom API for scope

* Update size-snapshot

* Fix isDevStore

* Fix typo - use weakset instead of set

.size-snapshot.json

@@ -1,8 +1,8 @@
 {
   "index.js": {
-    "bundled": 21267,
-    "minified": 10148,
-    "gzipped": 3367,
+    "bundled": 21477,
+    "minified": 10278,
+    "gzipped": 3412,
     "treeshaked": {
       "rollup": {
         "code": 14,
@@ -14,23 +14,23 @@
     }
   },
   "utils.js": {
-    "bundled": 16025,
-    "minified": 7685,
-    "gzipped": 2838,
+    "bundled": 15849,
+    "minified": 7524,
+    "gzipped": 2769,
     "treeshaked": {
       "rollup": {
         "code": 28,
         "import_statements": 28
       },
       "webpack": {
-        "code": 1305
+        "code": 1317
       }
     }
   },
   "devtools.js": {
-    "bundled": 19122,
-    "minified": 9537,
-    "gzipped": 3288,
+    "bundled": 19055,
+    "minified": 9478,
+    "gzipped": 3259,
     "treeshaked": {
       "rollup": {
         "code": 28,
@@ -42,9 +42,9 @@
     }
   },
   "immer.js": {
-    "bundled": 1637,
-    "minified": 856,
-    "gzipped": 418,
+    "bundled": 1615,
+    "minified": 844,
+    "gzipped": 407,
     "treeshaked": {
       "rollup": {
         "code": 42,
@@ -56,9 +56,9 @@
     }
   },
   "optics.js": {
-    "bundled": 1684,
-    "minified": 828,
-    "gzipped": 440,
+    "bundled": 1646,
+    "minified": 812,
+    "gzipped": 429,
     "treeshaked": {
       "rollup": {
         "code": 32,
@@ -70,9 +70,9 @@
     }
   },
   "query.js": {
-    "bundled": 5444,
-    "minified": 2279,
-    "gzipped": 710,
+    "bundled": 5106,
+    "minified": 2143,
+    "gzipped": 684,
     "treeshaked": {
       "rollup": {
         "code": 80,
@@ -84,9 +84,9 @@
     }
   },
   "xstate.js": {
-    "bundled": 3565,
-    "minified": 1470,
-    "gzipped": 684,
+    "bundled": 2977,
+    "minified": 1314,
+    "gzipped": 653,
     "treeshaked": {
       "rollup": {
         "code": 29,
@@ -98,9 +98,9 @@
     }
   },
   "valtio.js": {
-    "bundled": 1333,
-    "minified": 642,
-    "gzipped": 350,
+    "bundled": 1235,
+    "minified": 598,
+    "gzipped": 335,
     "treeshaked": {
       "rollup": {
         "code": 37,
@@ -112,9 +112,9 @@
     }
   },
   "zustand.js": {
-    "bundled": 647,
-    "minified": 306,
-    "gzipped": 204,
+    "bundled": 549,
+    "minified": 262,
+    "gzipped": 188,
     "treeshaked": {
       "rollup": {
         "code": 14,
@@ -126,9 +126,9 @@
     }
   },
   "redux.js": {
-    "bundled": 516,
-    "minified": 248,
-    "gzipped": 180,
+    "bundled": 458,
+    "minified": 220,
+    "gzipped": 167,
     "treeshaked": {
       "rollup": {
         "code": 14,
@@ -140,9 +140,9 @@
     }
   },
   "urql.js": {
-    "bundled": 4539,
-    "minified": 2369,
-    "gzipped": 856,
+    "bundled": 4241,
+    "minified": 2259,
+    "gzipped": 837,
     "treeshaked": {
       "rollup": {
         "code": 170,

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
 }
 
@@ -84,5 +86,5 @@ export const getStoreContext = (scope?: Scope) => {
 }
 
 export const isDevStore = (store: Store): store is StoreForDevelopment => {
-  return store.length > 3
+  return store.length > 4
 }

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]) => {
       restore(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,37 @@
+import { useContext } from 'react'
+import { SECRET_INTERNAL_getStoreContext as getStoreContext } from 'jotai'
+import type { Atom, Scope } from '../core/atom'
+import type { Store } from '../core/contexts'
+
+const hydratedMap: WeakMap<Store, WeakSet<Atom<unknown>>> = new WeakMap()
+
+export function useHydrateAtoms(
+  values: Iterable<readonly [Atom<unknown>, unknown]>,
+  scope?: Scope
+) {
+  const StoreContext = getStoreContext(scope)
+  const store = useContext(StoreContext)
+  const restoreAtoms = store[3]
+
+  const hydratedSet = getHydratedSet(store)
+  const tuplesToRestore = []
+  for (const tuple of values) {
+    const atom = tuple[0]
+    if (!hydratedSet.has(atom)) {
+      hydratedSet.add(atom)
+      tuplesToRestore.push(tuple)
+    }
+  }
+  if (tuplesToRestore.length) {
+    restoreAtoms(tuplesToRestore)
+  }
+}
+
+function getHydratedSet(store: Store) {
+  let hydratedSet = hydratedMap.get(store)
+  if (!hydratedSet) {
+    hydratedSet = new WeakSet()
+    hydratedMap.set(store, hydratedSet)
+  }
+  return hydratedSet
+}