wip

vikejs · Dec 28, 2023 · 47123f4 · 47123f4
1 parent 492e93a
commit 47123f4
Show file tree

Hide file tree

Showing 4 changed files with 37 additions and 28 deletions.
diff --git a/examples/zustand/store.ts b/examples/zustand/store.ts
@@ -8,23 +8,27 @@ interface Store {
   serverEnv: string
 }
 
-const useStore = create<Store>()(
-  (
-    set,
-    get
-    /* TODO
+const useStore = withPageContext(
+  // pageContextProviderFn
+  (pageContext) =>
+    create<Store>(
+      (
+        set,
+        get
+        /* TODO
   pageContext
   */
-  ) => ({
-    counter: Math.floor(10000 * Math.random()),
-    setCounter(value) {
-      set({ counter: value })
-    },
+      ) => ({
+        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!
-    }))
-  })
+        // 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
@@ -1,9 +1,10 @@
-export { create, server }
+export { create, server, 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 {assert} from './utils.js'
 
 type Create = {
   <T, Mos extends [StoreMutatorIdentifier, unknown][] = []>(
@@ -39,7 +40,12 @@ const create: Create = ((createState: any) => {
 function createImpl(createStore: any): any {
   // @ts-ignore
   setCreateStore((pageContext: any) => {
-    return create_(createStore)
+    if( createStore._withPageContext_  ) {
+      return create_(createStore(pageContext))
+    } else {
+      assert(false)
+      // return create_(createStore)
+    }
   })
 
   function useStore(...args: any[]) {
@@ -60,11 +66,8 @@ function server<T extends Record<string, any>>(fn: () => T) {
   return {} as T
 }
 type StoreAndHook = ReturnType<typeof create>
-function withPageContext<S extends StoreAndHook>(storeCreatorCreatorFn: (pageContext: StoreAndHook) => S) {
-  //@ts-ignore
-  // createImpl._withPageContext_ = storeCreatorFn
-  // const storeCreatorFn = () => {
-  //   storeCreatorCreatorFn(pageContext)
-  //   return createImpl(storeCreatorFn)
-  // }
+function withPageContext<S extends StoreAndHook>(pageContextProviderFn: (pageContext: StoreAndHook) => S) {
+  const fn = (pageContext: any) => pageContextProviderFn(pageContext)
+  fn._withPageContext_ = true
+  return createImpl(fn)
 }
diff --git a/packages/vike-react-zustand/src/renderer/VikeReactZustandWrapper.tsx b/packages/vike-react-zustand/src/renderer/VikeReactZustandWrapper.tsx
@@ -13,7 +13,8 @@ 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])
+  assert(store)
   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
@@ -1,9 +1,10 @@
 import React, { createContext } from 'react'
 import { getGlobalObject } from '../utils.js'
 import type { create } from 'zustand'
+import type { PageContext } from 'vike/types'
 
 type StoreAndHook = ReturnType<typeof create>
-type CreateStore = () => StoreAndHook & { __hydrated__?: true }
+type CreateStore = (pageContext: PageContext) => StoreAndHook & { __hydrated__?: true }
 
 const globalObject = getGlobalObject('VikeReactZustandContext.ts', {
   createStore: undefined as CreateStore | undefined,
@@ -14,5 +15,5 @@ export const getCreateStore = () => globalObject.createStore
 export const getContext = () => globalObject.context as unknown as React.Context<StoreAndHook | undefined>
 
 export const setCreateStore = (createStore_: CreateStore) => {
-  globalObject.createStore = createStore_
+  globalObject.createStore = globalObject.createStore || createStore_
 }