new design - wip

examples/zustand/store.ts

@@ -1,17 +1,21 @@
 export { useStore }
 
-import { createUseStore, server } from 'vike-react-zustand'
-import { PageContext } from 'vike/types'
-import { create } from 'zustand'
+import { create, server } from 'vike-react-zustand'
 
 interface Store {
   counter: number
   setCounter: (value: number) => void
   serverEnv: string
 }
 
-const useStore = createUseStore((pageContext: PageContext) =>
-  create<Store>()((set, get) => ({
+const useStore = create<Store>(
+  (
+    set,
+    get
+  /* TODO
+  pageContext
+  */
+  ) => ({
     counter: Math.floor(10000 * Math.random()),
     setCounter(value) {
       set({ counter: value })
@@ -22,5 +26,5 @@ const useStore = createUseStore((pageContext: PageContext) =>
     ...server(() => ({
       serverEnv: process.env.SOME_ENV!
     }))
-  }))
+  })
 )

packages/vike-react-zustand/src/index.tsx

@@ -1,21 +1,20 @@
-export { createUseStore, server }
+export { create, server }
 
 import { useContext } from 'react'
-import { PageContext } from 'vike/types'
-import { useStore as useZustandStore } from 'zustand'
 import { getContext, setCreateStore } from './renderer/context.js'
-import { ExtractState, StoreApi } from './types.js'
+import { create as create_ } from 'zustand'
 
-function createUseStore<S extends StoreApi>(createStore: (pageContext: PageContext) => S) {
-  setCreateStore(createStore)
+function create(createStore: any): any {
+  setCreateStore(() => {
+    return create_(createStore)
+  })
 
-  function useStore(): ExtractState<S>
-  function useStore<TSelection>(selector: (state: ExtractState<S>) => TSelection): TSelection
-  function useStore<TSelection>(selector?: (state: ExtractState<S>) => TSelection) {
-    const zustandContext = getContext<S>()
+  function useStore(...args: any[]) {
+    const zustandContext = getContext()
     const store = useContext(zustandContext)
     if (!store) throw new Error('Store is missing the provider')
-    return useZustandStore(store, selector ?? store.getState)
+    // @ts-ignore
+    return store(...args)
   }
 
   return useStore

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?.(pageContext), [createStore])
+  const store = useMemo(() => createStore?.(), [createStore])
   if (!store) {
     // Is that the best thing to do?
     return children

packages/vike-react-zustand/src/renderer/context.ts

@@ -1,14 +1,17 @@
 import React, { createContext } from 'react'
-import { CreateStore, StoreApi } from '../types.js'
 import { getGlobalObject } from '../utils.js'
+import type { create } from 'zustand'
+
+type StoreAndHook = ReturnType<typeof create>
+type CreateStore = () => StoreAndHook & { __hydrated__?: true }
 
 const globalObject = getGlobalObject('VikeReactZustandContext.ts', {
   createStore: undefined as CreateStore | undefined,
-  context: createContext<StoreApi | undefined>(undefined)
+  context: createContext<StoreAndHook | undefined>(undefined)
 })
 
 export const getCreateStore = () => globalObject.createStore
-export const getContext = <S extends StoreApi>() => globalObject.context as unknown as React.Context<S | undefined>
+export const getContext = () => globalObject.context as unknown as React.Context<StoreAndHook | undefined>
 
 export const setCreateStore = (createStore_: CreateStore) => {
   globalObject.createStore = createStore_