From 3453549254312cb317df527449928e788251c239 Mon Sep 17 00:00:00 2001
From: Eduardo San Martin Morote <posva13@gmail.com>
Date: Fri, 9 Jul 2021 11:44:43 +0200
Subject: [PATCH] refactor(setup): correctly extract the state

---
 src/index.ts |  2 +-
 src/store.ts | 32 +++++++++++---------------------
 2 files changed, 12 insertions(+), 22 deletions(-)

diff --git a/src/index.ts b/src/index.ts
index 09813b1b..b6913215 100644
--- a/src/index.ts
+++ b/src/index.ts
@@ -2,7 +2,7 @@ export { setActivePinia } from './rootStore'
 export { createPinia } from './createPinia'
 export type { Pinia, PiniaStorePlugin, PiniaPluginContext } from './rootStore'
 
-export { defineStore } from './store'
+export { defineStore, defineSetupStore } from './store'
 
 export type {
   StateTree,
diff --git a/src/store.ts b/src/store.ts
index cef3dc9e..aa28b93a 100644
--- a/src/store.ts
+++ b/src/store.ts
@@ -16,7 +16,6 @@ import {
   effectScope,
   onScopeDispose,
   EffectScope,
-  getCurrentScope,
   onUnmounted,
   ComputedRef,
   toRef,
@@ -400,22 +399,6 @@ function buildStoreToUse<
   return store
 }
 
-type _SetupOfStateActions<
-  S extends StateTree,
-  A extends Record<any, _Method>
-> = (StateTree extends S ? {} : S) & (Record<any, _Method> extends A ? {} : A)
-
-function stateBuilder<S extends StateTree>(
-  builder: (initial?: S | undefined) => S
-) {
-  return {} as S
-}
-
-const a = stateBuilder<{ n: number; toggle: boolean }>((initial) => ({
-  n: initial?.n || 2,
-  toggle: true,
-}))
-
 export interface DefineSetupStoreOptions<
   Id extends string,
   S extends StateTree,
@@ -438,7 +421,6 @@ function createSetupStore<
 >(
   $id: Id,
   setup: () => SS,
-  initialState: S | undefined,
   {
     // @ts-expect-error
     hydrate = innerPatch,
@@ -473,6 +455,12 @@ function createSetupStore<
   let subscriptions: SubscriptionCallback<S>[] = markRaw([])
   let actionSubscriptions: StoreOnActionListener<Id, S, G, A>[] = markRaw([])
   let debuggerEvents: DebuggerEvent[] | DebuggerEvent
+  const initialState = pinia.state.value[$id] as S | undefined
+
+  if (!initialState) {
+    // should be set in Vue 2
+    pinia.state.value[$id] = {}
+  }
 
   const triggerSubscriptions: SubscriptionCallback<S> = (mutation, state) => {
     subscriptions.forEach((callback) => {
@@ -490,8 +478,6 @@ function createSetupStore<
     scope = effectScope()
     return scope.run(() => {
       const store = setup()
-      // TODO: extract state and set it to pinia.state.value[$id]
-      // pinia.state.value[$id] = initialState || buildState()
 
       watch(
         () => pinia.state.value[$id] as UnwrapRef<S>,
@@ -634,7 +620,7 @@ function createSetupStore<
       }
     } else if ((isRef(prop) && !isComputed(prop)) || isReactive(prop)) {
       // mark it as a piece of state to be serialized
-      pinia.state.value[$id] = toRef(setupStore as any, key)
+      pinia.state.value[$id][key] = toRef(setupStore as any, key)
     } else if (__DEV__ && IS_CLIENT) {
       // add getters for devtools
       if (isComputed(prop)) {
@@ -691,6 +677,10 @@ function createSetupStore<
     }
   })
 
+  if (initialState) {
+    hydrate(store, initialState)
+  }
+
   return store
 }
 
-- 
2.47.2