feat(utils/useHydrateAtoms) - Optionally rehydrate with force hydrate

docs/utilities/ssr.mdx

@@ -48,6 +48,20 @@ useHydrateAtoms(new Map([[count, 42]]))
 ```
 
 Atoms can only be hydrated once per store. Therefore, if the initial value used is changed during rerenders, it won't update the atom value.
+If there is a unique need to re-hydrate a previously hydrated atom, pass the optional dangerouslyForceHydrate as true
+and note that it may behave wrongly in concurrent rendering.
+
+```js
+useHydrateAtoms(
+  [
+    [countAtom, 42],
+    [frameworkAtom, 'Next.js'],
+  ],
+  {
+    dangerouslyForceHydrate: true,
+  }
+)
+```
 
 If there's a need to hydrate in multiple stores, use multiple `useHydrateAtoms` hooks to achieve that.
 

src/react/utils/useHydrateAtoms.ts

@@ -2,7 +2,9 @@ import { useStore } from '../../react.ts'
 import type { WritableAtom } from '../../vanilla.ts'
 
 type Store = ReturnType<typeof useStore>
-type Options = Parameters<typeof useStore>[0]
+type Options = Parameters<typeof useStore>[0] & {
+  dangerouslyForceHydrate?: boolean
+}
 type AnyWritableAtom = WritableAtom<unknown, any[], any>
 type AtomMap<A = AnyWritableAtom, V = unknown> = Map<A, V>
 type AtomTuple<A = AnyWritableAtom, V = unknown> = readonly [A, V]
@@ -36,7 +38,7 @@ export function useHydrateAtoms<T extends Iterable<AtomTuple>>(
 
   const hydratedSet = getHydratedSet(store)
   for (const [atom, value] of values) {
-    if (!hydratedSet.has(atom)) {
+    if (!hydratedSet.has(atom) || options?.dangerouslyForceHydrate) {
       hydratedSet.add(atom)
       store.set(atom, value)
     }

tests/react/utils/useHydrateAtoms.test.tsx

@@ -240,3 +240,78 @@ it('useHydrateAtoms should respect onMount', async () => {
   await findByText('count: 42')
   expect(onMountFn).toHaveBeenCalledTimes(1)
 })
+
+it.only('passing dangerouslyForceHydrate to useHydrateAtoms will re-hydrated atoms', async () => {
+  const countAtom = atom(0)
+  const statusAtom = atom('fulfilled')
+
+  const Counter = ({
+    initialCount,
+    initialStatus,
+    dangerouslyForceHydrate = false,
+  }: {
+    initialCount: number
+    initialStatus: string
+    dangerouslyForceHydrate?: boolean
+  }) => {
+    useHydrateAtoms(
+      [
+        [countAtom, initialCount],
+        [statusAtom, initialStatus],
+      ],
+      {
+        dangerouslyForceHydrate,
+      }
+    )
+    const [countValue, setCount] = useAtom(countAtom)
+    const [statusValue, setStatus] = useAtom(statusAtom)
+
+    return (
+      <>
+        <div>count: {countValue}</div>
+        <button onClick={() => setCount((count) => count + 1)}>dispatch</button>
+        <div>status: {statusValue}</div>
+        <button
+          onClick={() =>
+            setStatus((status) =>
+              status === 'fulfilled' ? 'rejected' : 'fulfilled'
+            )
+          }>
+          update
+        </button>
+      </>
+    )
+  }
+  const { findByText, getByText, rerender } = render(
+    <StrictMode>
+      <Counter initialCount={42} initialStatus="rejected" />
+    </StrictMode>
+  )
+
+  await findByText('count: 42')
+  await findByText('status: rejected')
+  fireEvent.click(getByText('dispatch'))
+  fireEvent.click(getByText('update'))
+  await findByText('count: 43')
+  await findByText('status: fulfilled')
+
+  rerender(
+    <StrictMode>
+      <Counter initialCount={65} initialStatus="rejected" />
+    </StrictMode>
+  )
+  await findByText('count: 43')
+  await findByText('status: fulfilled')
+
+  rerender(
+    <StrictMode>
+      <Counter
+        initialCount={11}
+        initialStatus="rejected"
+        dangerouslyForceHydrate={true}
+      />
+    </StrictMode>
+  )
+  await findByText('count: 11')
+  await findByText('status: rejected')
+})