feat: getInitialState (#2277)

* feat(react): implement getServerState by closing over the defaultState

serverState will be used by react on the first client render; this should avoid hydration mismatches when combined with the persist middleware, which can change the state between the SSR and the first CSR

* define getServerState in vanilla.ts

* feat: implement getServerResult in persist middleware

this avoids hydration errors when state is restored from e.g. localstorage synchronously

* feat: capture initialState for getServerState in react

this avoids hydration mismatches when updates happen to the store state between ssr and csr

* refactor: revert changes to oldImpl

* fix: make selector default to identity function

if we default to `api.getState`, we will always read the client snapshot if there is no selector passed. An identity function returns its argument, which is either the snapshot (api.getState) or the server snapshot (api.getServerState)

* define getInitialState in vanilla

* revert WithReact

* fix them

* fix test

* oops, fix another test too

* forgot to use identity

* test: add a test for hydration errors

* fix(readme): imply getInitialState is a public api

---------

Co-authored-by: daishi <daishi@axlight.com>

readme.md

@@ -228,7 +228,7 @@ Zustand core can be imported and used without the React dependency. The only dif
 import { createStore } from 'zustand/vanilla'
 
 const store = createStore((set) => ...)
-const { getState, setState, subscribe } = store
+const { getState, setState, subscribe, getInitialState } = store
 
 export default store
 ```

src/middleware/persist.ts

@@ -425,6 +425,8 @@ const newImpl: PersistImpl = (config, baseOptions) => (set, get, api) => {
     api,
   )
 
+  api.getInitialState = () => configResult
+
   // a workaround to solve the issue of not storing rehydrated state in sync storage
   // the set(state) value would be later overridden with initial state by create()
   // to avoid this, we merge the state from localStorage into the initial state.

src/react.ts

@@ -28,6 +28,8 @@ type WithReact<S extends ReadonlyStoreApi<unknown>> = S & {
 
 let didWarnAboutEqualityFn = false
 
+const identity = <T>(arg: T): T => arg
+
 export function useStore<S extends WithReact<StoreApi<unknown>>>(
   api: S,
 ): ExtractState<S>
@@ -49,7 +51,7 @@ export function useStore<S extends WithReact<StoreApi<unknown>>, U>(
 
 export function useStore<TState, StateSlice>(
   api: WithReact<StoreApi<TState>>,
-  selector: (state: TState) => StateSlice = api.getState as any,
+  selector: (state: TState) => StateSlice = identity as any,
   equalityFn?: (a: StateSlice, b: StateSlice) => boolean,
 ) {
   if (
@@ -65,7 +67,7 @@ export function useStore<TState, StateSlice>(
   const slice = useSyncExternalStoreWithSelector(
     api.subscribe,
     api.getState,
-    api.getServerState || api.getState,
+    api.getServerState || api.getInitialState,
     selector,
     equalityFn,
   )

src/traditional.ts

@@ -26,6 +26,8 @@ type WithReact<S extends ReadonlyStoreApi<unknown>> = S & {
   getServerState?: () => ExtractState<S>
 }
 
+const identity = <T>(arg: T): T => arg
+
 export function useStoreWithEqualityFn<S extends WithReact<StoreApi<unknown>>>(
   api: S,
 ): ExtractState<S>
@@ -41,13 +43,13 @@ export function useStoreWithEqualityFn<
 
 export function useStoreWithEqualityFn<TState, StateSlice>(
   api: WithReact<StoreApi<TState>>,
-  selector: (state: TState) => StateSlice = api.getState as any,
+  selector: (state: TState) => StateSlice = identity as any,
   equalityFn?: (a: StateSlice, b: StateSlice) => boolean,
 ) {
   const slice = useSyncExternalStoreWithSelector(
     api.subscribe,
     api.getState,
-    api.getServerState || api.getState,
+    api.getServerState || api.getInitialState,
     selector,
     equalityFn,
   )

src/vanilla.ts

@@ -8,6 +8,7 @@ type SetStateInternal<T> = {
 export interface StoreApi<T> {
   setState: SetStateInternal<T>
   getState: () => T
+  getInitialState: () => T
   subscribe: (listener: (state: T, prevState: T) => void) => () => void
   /**
    * @deprecated Use `unsubscribe` returned by `subscribe`
@@ -82,6 +83,9 @@ const createStoreImpl: CreateStoreImpl = (createState) => {
 
   const getState: StoreApi<TState>['getState'] = () => state
 
+  const getInitialState: StoreApi<TState>['getInitialState'] = () =>
+    initialState
+
   const subscribe: StoreApi<TState>['subscribe'] = (listener) => {
     listeners.add(listener)
     // Unsubscribe
@@ -97,8 +101,8 @@ const createStoreImpl: CreateStoreImpl = (createState) => {
     listeners.clear()
   }
 
-  const api = { setState, getState, subscribe, destroy }
-  state = createState(setState, getState, api)
+  const api = { setState, getState, getInitialState, subscribe, destroy }
+  const initialState = (state = createState(setState, getState, api))
   return api as any
 }
 

tests/basic.test.tsx

@@ -31,6 +31,7 @@ it('creates a store hook and api object', () => {
         [Function],
         {
           "destroy": [Function],
+          "getInitialState": [Function],
           "getState": [Function],
           "setState": [Function],
           "subscribe": [Function],

tests/ssr.test.tsx

@@ -63,6 +63,55 @@ describe.skipIf(!React.version.startsWith('18'))(
         )
       })
 
+      const bearCountText = await screen.findByText('bears: 1')
+      expect(bearCountText).not.toBeNull()
+      document.body.removeChild(container)
+    })
+    it('should not have hydration errors', async () => {
+      const useStore = create(() => ({
+        bears: 0,
+      }))
+
+      const { hydrateRoot } =
+        await vi.importActual<typeof import('react-dom/client')>(
+          'react-dom/client',
+        )
+
+      const Component = () => {
+        const bears = useStore((state) => state.bears)
+        return <div>bears: {bears}</div>
+      }
+
+      const markup = renderToString(
+        <React.Suspense fallback={<div>Loading...</div>}>
+          <Component />
+        </React.Suspense>,
+      )
+
+      const container = document.createElement('div')
+      document.body.appendChild(container)
+      container.innerHTML = markup
+
+      expect(container.textContent).toContain('bears: 0')
+
+      const consoleMock = vi.spyOn(console, 'error')
+
+      const hydratePromise = act(async () => {
+        hydrateRoot(
+          container,
+          <React.Suspense fallback={<div>Loading...</div>}>
+            <Component />
+          </React.Suspense>,
+        )
+      })
+
+      // set state during hydration
+      useStore.setState({ bears: 1 })
+
+      await hydratePromise
+
+      expect(consoleMock).toHaveBeenCalledTimes(0)
+
       const bearCountText = await screen.findByText('bears: 1')
       expect(bearCountText).not.toBeNull()
       document.body.removeChild(container)

tests/vanilla/basic.test.ts

@@ -17,12 +17,13 @@ it('create a store', () => {
     return { value: null }
   })
   expect({ params, result }).toMatchInlineSnapshot(`
-   {
+    {
       "params": [
         [Function],
         [Function],
         {
           "destroy": [Function],
+          "getInitialState": [Function],
           "getState": [Function],
           "setState": [Function],
           "subscribe": [Function],
@@ -30,6 +31,7 @@ it('create a store', () => {
       ],
       "result": {
         "destroy": [Function],
+        "getInitialState": [Function],
         "getState": [Function],
         "setState": [Function],
         "subscribe": [Function],