From f79c762c438db1dfdb8a4ee02fdb4c64dc3851ff Mon Sep 17 00:00:00 2001
From: Eduardo San Martin Morote <posva13@gmail.com>
Date: Thu, 15 Jul 2021 17:57:16 +0200
Subject: [PATCH] feat(hmr): handle changes to setup store

---
 playground/src/views/CounterSetupStore.vue |  2 +-
 playground/src/views/CounterStore.vue      |  2 +-
 src/devtools/plugin.ts                     | 28 +++++++-----
 src/store.ts                               | 53 ++++++++++++++++------
 src/types.ts                               |  3 +-
 5 files changed, 59 insertions(+), 29 deletions(-)

diff --git a/playground/src/views/CounterSetupStore.vue b/playground/src/views/CounterSetupStore.vue
index b68ad5af..c06656b6 100644
--- a/playground/src/views/CounterSetupStore.vue
+++ b/playground/src/views/CounterSetupStore.vue
@@ -5,7 +5,7 @@
 
   <h2>Counter Store</h2>
 
-  <p>Counter :{{ counter.n }}</p>
+  <p>Counter :{{ counter.n }}. Double: {{ counter.double }}</p>
 
   <p>
     Increment the Store <br />
diff --git a/playground/src/views/CounterStore.vue b/playground/src/views/CounterStore.vue
index 60b5da96..5d0b14eb 100644
--- a/playground/src/views/CounterStore.vue
+++ b/playground/src/views/CounterStore.vue
@@ -5,7 +5,7 @@
 
   <h2>Counter Store</h2>
 
-  <p>Counter :{{ counter.n }}</p>
+  <p>Counter :{{ counter.n }}. Double: {{ counter.double }}</p>
 
   <p>
     Increment the Store <br />
diff --git a/src/devtools/plugin.ts b/src/devtools/plugin.ts
index 3f415609..e8ff0886 100644
--- a/src/devtools/plugin.ts
+++ b/src/devtools/plugin.ts
@@ -417,22 +417,26 @@ export function devtoolsPlugin<
     return
   }
 
-  patchActionForGrouping(
-    // @ts-expect-error: can cast the store...
-    store,
-    Object.keys(options.actions)
-  )
-
-  const originalHotUpdate = store.hotUpdate
-
-  // Upgrade the HMR to also update the new actions
-  toRaw(store).hotUpdate = function (newStore) {
-    originalHotUpdate.apply(this, arguments as any)
+  // only wrap actions in option-defined stores as this technique relies on
+  // wrapping the context of the action with a proxy
+  if ('id' in options) {
     patchActionForGrouping(
       // @ts-expect-error: can cast the store...
       store,
-      Object.keys(newStore._hmrPayload.actions)
+      Object.keys(options.actions)
     )
+
+    const originalHotUpdate = store.hotUpdate
+
+    // Upgrade the HMR to also update the new actions
+    toRaw(store).hotUpdate = function (newStore) {
+      originalHotUpdate.apply(this, arguments as any)
+      patchActionForGrouping(
+        // @ts-expect-error: can cast the store...
+        store,
+        Object.keys(newStore._hmrPayload.actions)
+      )
+    }
   }
 
   addStoreToDevtools(
diff --git a/src/store.ts b/src/store.ts
index 8ac71fb7..74811796 100644
--- a/src/store.ts
+++ b/src/store.ts
@@ -94,13 +94,19 @@ function createOptionsStore<
   const initialState: StateTree | undefined = pinia.state.value[id]
 
   function setup() {
-    if (!initialState) {
+    if (!initialState && (!__DEV__ || !hot)) {
       $reset()
     }
     // pinia.state.value[id] = state ? state() : {}
 
+    // avoid creating a state in pinia.state.value
+    const localState =
+      __DEV__ && hot
+        ? toRefs(ref(state ? state() : {}).value)
+        : initialState || toRefs(pinia.state.value[id])
+
     return assign(
-      initialState || toRefs(pinia.state.value[id]),
+      localState,
       actions,
       Object.keys(getters || {}).reduce((computedGetters, name) => {
         computedGetters[name] = computed(() => {
@@ -176,11 +182,13 @@ function createSetupStore<
   let debuggerEvents: DebuggerEvent[] | DebuggerEvent
   const initialState = pinia.state.value[$id] as UnwrapRef<S> | undefined
 
-  if (!initialState) {
+  if (!initialState && !hot) {
     // should be set in Vue 2
     pinia.state.value[$id] = {}
   }
 
+  const hotState = ref({} as S)
+
   const triggerSubscriptions: SubscriptionCallback<S> = (mutation, state) => {
     subscriptions.forEach((callback) => {
       callback(mutation, state)
@@ -310,7 +318,7 @@ function createSetupStore<
   }
 
   /**
-   * Wraps an action to handle subscriptions
+   * Wraps an action to handle subscriptions.
    *
    * @param name - name of the action
    * @param action - action to wrap
@@ -359,6 +367,7 @@ function createSetupStore<
     actions: {} as Record<string, any>,
     getters: {} as Record<string, Ref>,
     state: [] as string[],
+    hotState,
   })
 
   // overwrite existing actions to support $onAction
@@ -366,11 +375,14 @@ function createSetupStore<
     const prop = setupStore[key]
 
     if ((isRef(prop) && !isComputed(prop)) || isReactive(prop)) {
-      // createOptionStore already did this
-      if (!buildState) {
-        // mark it as a piece of state to be serialized
+      // mark it as a piece of state to be serialized
+      if (__DEV__ && hot) {
+        hotState.value[key] = toRef(setupStore as any, key)
+        // createOptionStore already did this
+      } else if (!buildState) {
         pinia.state.value[$id][key] = toRef(setupStore as any, key)
       }
+
       if (__DEV__) {
         _hmrPayload.state.push(key)
       }
@@ -431,28 +443,41 @@ function createSetupStore<
   // without linking the computed lifespan to wherever the store is first
   // created.
   Object.defineProperty(store, '$state', {
-    get: () => pinia.state.value[$id],
-    set: (state) => (pinia.state.value[$id] = state),
+    get: () => (__DEV__ && hot ? hotState.value : pinia.state.value[$id]),
+    set: (state) => {
+      if (__DEV__ && hot) {
+        throw new Error('cannot set hotState')
+      }
+      pinia.state.value[$id] = state
+    },
   })
 
   // add the hotUpdate before plugins to allow them to override it
   if (__DEV__) {
     store.hotUpdate = markRaw((newStore) => {
       newStore._hmrPayload.state.forEach((stateKey) => {
-        if (!(stateKey in store.$state)) {
-          console.log('setting new key', stateKey)
+        if (stateKey in store.$state) {
           // @ts-expect-error
           // transfer the ref
-          store.$state[stateKey] = newStore.$state[stateKey]
+          newStore.$state[stateKey] =
+            // ---
+            // @ts-expect-error
+            store.$state[stateKey]
+
+          // patch direct access properties to allow store.stateProperty to work as
+          // store.$state.stateProperty
+          // @ts-expect-error
+          store[stateKey] = toRef(newStore.$state, stateKey)
         }
       })
 
+      pinia.state.value[$id] = toRef(newStore._hmrPayload, 'hotState')
+
       // remove deleted keys
       Object.keys(store.$state).forEach((stateKey) => {
         if (!(stateKey in newStore.$state)) {
-          console.log('deleting old key', stateKey)
           // @ts-expect-error
-          delete store.$state[stateKey]
+          delete store[stateKey]
         }
       })
 
diff --git a/src/types.ts b/src/types.ts
index 45883dc5..6c665a71 100644
--- a/src/types.ts
+++ b/src/types.ts
@@ -1,4 +1,4 @@
-import { DebuggerEvent, UnwrapRef } from 'vue'
+import { DebuggerEvent, Ref, UnwrapRef } from 'vue'
 import { Pinia } from './rootStore'
 
 /**
@@ -265,6 +265,7 @@ export interface StoreWithState<
    */
   _hmrPayload: {
     state: string[]
+    hotState: Ref<S>
     actions: ActionsTree
     getters: ActionsTree
   }
-- 
2.47.2