From 2998d53272856d64587d5e880f7dca41775356f0 Mon Sep 17 00:00:00 2001
From: Eduardo San Martin Morote <posva13@gmail.com>
Date: Tue, 14 Jan 2020 20:20:39 +0100
Subject: [PATCH] feat: handle SSR state hydration

---
 src/index.ts |  7 ++++++-
 src/store.ts | 48 ++++++++++++++++++++++++++++++++++++++----------
 2 files changed, 44 insertions(+), 11 deletions(-)

diff --git a/src/index.ts b/src/index.ts
index 4d62a5f8..0640b0ca 100644
--- a/src/index.ts
+++ b/src/index.ts
@@ -1,2 +1,7 @@
-export { createStore, CombinedStore, setActiveReq } from './store'
+export {
+  createStore,
+  CombinedStore,
+  setActiveReq,
+  setStateProvider,
+} from './store'
 export { StateTree, StoreGetter } from './types'
diff --git a/src/store.ts b/src/store.ts
index a0c2f7b6..e0b40c96 100644
--- a/src/store.ts
+++ b/src/store.ts
@@ -146,11 +146,15 @@ export function buildStore<
   return store
 }
 
+type NonNullObject = Record<any, any>
+
 /**
  * setActiveReq must be called to handle SSR at the top of functions like `fetch`, `setup`, `serverPrefetch` and others
  */
-export let activeReq: object = {}
-export const setActiveReq = (req: object) => (activeReq = req)
+export let activeReq: NonNullObject = {}
+export const setActiveReq = (req: NonNullObject | undefined) =>
+  req && (activeReq = req)
+
 export const getActiveReq = () => activeReq
 
 /**
@@ -160,10 +164,37 @@ export const getActiveReq = () => activeReq
  */
 
 const storesMap = new WeakMap<
-  object,
+  NonNullObject,
   Record<string, CombinedStore<any, any, any>>
 >()
 
+/**
+ * A state provider allows to set how states are stored for hydration. e.g. setting a property on a context, getting a property from window
+ */
+interface StateProvider {
+  get(): Record<string, StateTree>
+  set(store: CombinedStore<any, any, any>): any
+}
+
+/**
+ * Map of initial states used for hydration
+ */
+export const stateProviders = new WeakMap<NonNullObject, StateProvider>()
+
+export function setStateProvider(stateProvider: StateProvider) {
+  stateProviders.set(getActiveReq(), stateProvider)
+}
+
+function getInitialState(id: string): StateTree | undefined {
+  const provider = stateProviders.get(getActiveReq())
+  return provider && provider.get()[id]
+}
+
+function setInitialState(store: CombinedStore<any, any, any>): void {
+  const provider = stateProviders.get(getActiveReq())
+  if (provider) provider.set(store)
+}
+
 /**
  * Creates a `useStore` function that retrieves the store instance
  * @param id id of the store we are creating
@@ -175,23 +206,20 @@ export function createStore<
   S extends StateTree,
   G extends Record<string, StoreGetter<S>>
 >(id: Id, buildState: () => S, getters: G = {} as G) {
-  let store: CombinedStore<Id, S, G> | undefined
-
-  return function useStore(
-    initialStates: Record<Id, S> = {} as Record<Id, S>
-  ): CombinedStore<Id, S, G> {
+  return function useStore(): CombinedStore<Id, S, G> {
     const req = getActiveReq()
     let stores = storesMap.get(req)
     if (!stores) storesMap.set(req, (stores = {}))
 
-    store = stores[id]
+    let store = stores[id]
     if (!store) {
       stores[id] = store = buildStore(
         id,
         buildState,
         getters,
-        initialStates[id]
+        getInitialState(id)
       )
+      setInitialState(store)
       if (isClient) useStoreDevtools(store)
     }
 
-- 
2.47.2