updated zustand example (#24884)

..to use new zustand/context module from zustand 3.5.0. Also fixed code for merging states on client-side navigation.


## Documentation / Examples

- [x] Make sure the linting passes

examples/with-zustand/README.md

@@ -10,7 +10,7 @@ To illustrate SSG and SSR, go to `/ssg` and `/ssr`, those pages are using Next.j
 
 The trick here for supporting universal Zustand is to separate the cases for the client and the server. When we are on the server we want to create a new store every time, otherwise different users data will be mixed up. If we are in the client we want to use always the same store. That's what we accomplish on `store.js`.
 
-All components have access to the Zustand store using `useStore`.
+All components have access to the Zustand store using `useStore()` returned from zustand's `createContext()` function.
 
 On the server side every request initializes a new store, because otherwise different user data can be mixed up. On the client side the same store is used, even between page changes.
 

examples/with-zustand/components/clock.js

@@ -1,4 +1,4 @@
-import { useStore } from '../lib/zustandProvider'
+import { useStore } from '../lib/store'
 import shallow from 'zustand/shallow'
 
 const useClock = () => {

examples/with-zustand/components/counter.js

@@ -1,4 +1,4 @@
-import { useStore } from '../lib/zustandProvider'
+import { useStore } from '../lib/store'
 import shallow from 'zustand/shallow'
 const useCounter = () => {
   const { count, increment, decrement, reset } = useStore(

examples/with-zustand/components/page.js

@@ -2,7 +2,7 @@ import useInterval from '../lib/useInterval'
 import Clock from './clock'
 import Counter from './counter'
 import Nav from './nav'
-import { useStore } from '../lib/zustandProvider'
+import { useStore } from '../lib/store'
 
 export default function Page() {
   const { tick } = useStore()

examples/with-zustand/lib/store.js

@@ -1,5 +1,6 @@
-import { useMemo } from 'react'
+import { useLayoutEffect } from 'react'
 import create from 'zustand'
+import createContext from 'zustand/context'
 
 let store
 
@@ -9,7 +10,9 @@ const initialState = {
   count: 0,
 }
 
-function initStore(preloadedState = initialState) {
+export const { Provider, useStore } = createContext(initialState)
+
+export const initializeStore = (preloadedState = {}) => {
   return create((set, get) => ({
     ...initialState,
     ...preloadedState,
@@ -37,31 +40,30 @@ function initStore(preloadedState = initialState) {
   }))
 }
 
-export const initializeStore = (preloadedState) => {
-  let _store = store ?? initStore(preloadedState)
+export function useHydrate(initialState) {
+  let _store = store ?? initializeStore(initialState)
+
+  // For SSR & SSG, always use a new store.
+  if (typeof window !== 'undefined') {
+    // For CSR, always re-use same store.
+    if (!store) {
+      store = _store
+    }
 
-  // After navigating to a page with an initial Zustand state, merge that state
-  // with the current state in the store, and create a new store
-  if (preloadedState && store) {
-    _store = initStore({
-      ...store.getState(),
-      ...preloadedState,
-    })
-    // Reset the current store
-    store = undefined
+    // And if initialState changes, then merge states in the next render cycle.
+    //
+    // eslint complaining "React Hooks must be called in the exact same order in every component render"
+    // is ignorable as this code runs in the same order in a given environment (CSR/SSR/SSG)
+    // eslint-disable-next-line react-hooks/rules-of-hooks
+    useLayoutEffect(() => {
+      if (initialState && store) {
+        store.setState({
+          ...store.getState(),
+          ...initialState,
+        })
+      }
+    }, [initialState])
   }
 
-  // For SSG and SSR always create a new store
-  if (typeof window === 'undefined') return _store
-  // Create the store once in the client
-  if (!store) store = _store
-
   return _store
 }
-
-export function useHydrate(initialState) {
-  const state =
-    typeof initialState === 'string' ? JSON.parse(initialState) : initialState
-  const store = useMemo(() => initializeStore(state), [state])
-  return store
-}

examples/with-zustand/package.json

@@ -10,7 +10,7 @@
     "next": "latest",
     "react": "^16.9.0",
     "react-dom": "^16.9.0",
-    "zustand": "3.1.3"
+    "zustand": "3.5.0"
   },
   "license": "MIT"
 }

examples/with-zustand/pages/_app.js

@@ -1,12 +1,10 @@
-import { StoreProvider } from '../lib/zustandProvider'
-import { useHydrate } from '../lib/store'
+import { useHydrate, Provider } from '../lib/store'
 
 export default function App({ Component, pageProps }) {
   const store = useHydrate(pageProps.initialZustandState)
-
   return (
-    <StoreProvider store={store}>
+    <Provider initialStore={store}>
       <Component {...pageProps} />
-    </StoreProvider>
+    </Provider>
   )
 }

examples/with-zustand/pages/ssr.js

@@ -14,6 +14,10 @@ export function getServerSideProps() {
   zustandStore.getState().tick(Date.now(), false)
 
   return {
-    props: { initialZustandState: JSON.stringify(zustandStore.getState()) },
+    props: {
+      // the "stringify and then parse again" piece is required as next.js
+      // isn't able to serialize it to JSON properly
+      initialZustandState: JSON.parse(JSON.stringify(zustandStore.getState())),
+    },
   }
 }