diff --git a/examples/zustand/store.ts b/examples/zustand/store.ts index da7cf051..0d03bdf7 100644 --- a/examples/zustand/store.ts +++ b/examples/zustand/store.ts @@ -1,30 +1,40 @@ export { useStore } -import { create, server, withPageContext } from 'vike-react-zustand' +import { create, serverOnly, withPageContext } from 'vike-react-zustand' interface Store { counter: number setCounter: (value: number) => void serverEnv: string + url: string } -const useStore = create()( - ( - set, - get - /* TODO - pageContext - */ - ) => ({ +const useStore = withPageContext((pageContext) => + create()((set, get) => ({ counter: Math.floor(10000 * Math.random()), setCounter(value) { set({ counter: value }) }, + url: pageContext.urlOriginal, // the callback only runs on the server, // the return value is passed to the client on the initial navigation - ...server(() => ({ + ...serverOnly(() => ({ serverEnv: process.env.SOME_ENV! })) - }) + })) ) + +// This works, too +// const useStore = create()((set, get) => ({ +// counter: Math.floor(10000 * Math.random()), +// setCounter(value) { +// set({ counter: value }) +// }, + +// // the callback only runs on the server, +// // the return value is passed to the client on the initial navigation +// ...server(() => ({ +// serverEnv: process.env.SOME_ENV! +// })) +// })) diff --git a/packages/vike-react-zustand/src/index.tsx b/packages/vike-react-zustand/src/index.tsx index 77eb67d7..019c7f14 100644 --- a/packages/vike-react-zustand/src/index.tsx +++ b/packages/vike-react-zustand/src/index.tsx @@ -1,9 +1,10 @@ -export { create, server } +export { create, serverOnly, withPageContext } import { useContext } from 'react' import { getContext, setCreateStore } from './renderer/context.js' import { create as create_ } from 'zustand' import type { StoreMutatorIdentifier, UseBoundStore, Mutate, StoreApi as ZustandStoreApi } from 'zustand' +import type { PageContext } from 'vike/types' type Create = { ( @@ -32,14 +33,25 @@ export type StateCreator< $$storeMutators?: Mos } -const create: Create = ((createState: any) => { - return createState ? createImpl(createState) : createImpl +const create: Create = ((storeCreatorFn: any) => { + return storeCreatorFn ? createImpl(storeCreatorFn) : createImpl }) as any -function createImpl(createStore: any): any { - // @ts-ignore +function createImpl(storeCreatorFn: any): any { setCreateStore((pageContext: any) => { - return create_(createStore) + // This is called only once per request + if (storeCreatorFn._withPageContext) { + // storeCreatorFn(pageContext) looks like this: + // (pageContext) => + // create((set, get) => ({ + // counter: 123 + // })) + // create calls createImpl a second time, and it returns useStore. + // but we need to pass the original storeCreatorFn(_storeCreatorFn) to create_ + return create_()(storeCreatorFn(pageContext)._storeCreatorFn) + } else { + return create_()(storeCreatorFn) + } }) function useStore(...args: any[]) { @@ -50,21 +62,19 @@ function createImpl(createStore: any): any { return store(...args) } + useStore._storeCreatorFn = storeCreatorFn return useStore } -function server>(fn: () => T) { +function serverOnly>(fn: () => T) { if (typeof window === 'undefined') { return fn() } return {} as T } -type StoreAndHook = ReturnType -function withPageContext(storeCreatorCreatorFn: (pageContext: StoreAndHook) => S) { - //@ts-ignore - // createImpl._withPageContext_ = storeCreatorFn - // const storeCreatorFn = () => { - // storeCreatorCreatorFn(pageContext) - // return createImpl(storeCreatorFn) - // } + +function withPageContext>(storeCreatorFn: (pageContext: PageContext) => S): S { + const wrappedStoreCreatorFn = (pageContext: any) => storeCreatorFn(pageContext) + wrappedStoreCreatorFn._withPageContext = true + return createImpl(wrappedStoreCreatorFn) } diff --git a/packages/vike-react-zustand/src/renderer/VikeReactZustandWrapper.tsx b/packages/vike-react-zustand/src/renderer/VikeReactZustandWrapper.tsx index 863af90d..c1271412 100644 --- a/packages/vike-react-zustand/src/renderer/VikeReactZustandWrapper.tsx +++ b/packages/vike-react-zustand/src/renderer/VikeReactZustandWrapper.tsx @@ -13,7 +13,7 @@ export default function VikeReactZustandWrapper({ pageContext, children }: VikeR const context = getContext() assert(context) const createStore = getCreateStore() - const store = useMemo(() => createStore?.(), [createStore]) + const store = useMemo(() => createStore?.(pageContext), [createStore]) if (!store) { // Is that the best thing to do? return children diff --git a/packages/vike-react-zustand/src/renderer/context.ts b/packages/vike-react-zustand/src/renderer/context.ts index 5b0a502b..ffcaed5f 100644 --- a/packages/vike-react-zustand/src/renderer/context.ts +++ b/packages/vike-react-zustand/src/renderer/context.ts @@ -1,18 +1,20 @@ +export { getContext, getCreateStore, setCreateStore } + import React, { createContext } from 'react' import { getGlobalObject } from '../utils.js' import type { create } from 'zustand' type StoreAndHook = ReturnType -type CreateStore = () => StoreAndHook & { __hydrated__?: true } +type CreateStore = (pageContext: any) => StoreAndHook & { __hydrated__?: true } const globalObject = getGlobalObject('VikeReactZustandContext.ts', { createStore: undefined as CreateStore | undefined, context: createContext(undefined) }) -export const getCreateStore = () => globalObject.createStore -export const getContext = () => globalObject.context as unknown as React.Context +const getContext = () => globalObject.context as unknown as React.Context -export const setCreateStore = (createStore_: CreateStore) => { +const getCreateStore = () => globalObject.createStore +const setCreateStore = (createStore_: CreateStore) => { globalObject.createStore = createStore_ }