From 50f8d0fe545ec895e57b072e9e231c0521c6fc83 Mon Sep 17 00:00:00 2001
From: Eduardo San Martin Morote <posva13@gmail.com>
Date: Thu, 15 Jul 2021 11:40:03 +0200
Subject: [PATCH] feat(hmr): working version

---
 playground/src/App.vue           |  1 +
 playground/src/main.ts           |  5 +-
 playground/src/stores/counter.ts |  9 ++--
 src/devtools/plugin.ts           | 65 +++++++++++++++---------
 src/store.ts                     | 85 ++++++++++++++++----------------
 5 files changed, 92 insertions(+), 73 deletions(-)

diff --git a/playground/src/App.vue b/playground/src/App.vue
index 3edb3052..07db57d7 100644
--- a/playground/src/App.vue
+++ b/playground/src/App.vue
@@ -1,5 +1,6 @@
 <template>
   <button @click="n++">Increment {{ n }}</button>
+  <button @click="counter.changeMe()"><code>counter.changeMe()</code></button>
   <pre>{{ counter.$state }}</pre>
   <!-- <button @click="counter.newOne()">Click me</button> -->
   <TestStore />
diff --git a/playground/src/main.ts b/playground/src/main.ts
index f31805e7..bfb1990d 100644
--- a/playground/src/main.ts
+++ b/playground/src/main.ts
@@ -1,9 +1,6 @@
 import { createApp } from 'vue'
 import App from './App.vue'
-import { createPinia, StoreDefinition } from '../../src'
-import { ha } from './test'
-
-console.log({ ha })
+import { createPinia } from '../../src'
 
 const pinia = createPinia()
 
diff --git a/playground/src/stores/counter.ts b/playground/src/stores/counter.ts
index e9008527..4941ba03 100644
--- a/playground/src/stores/counter.ts
+++ b/playground/src/stores/counter.ts
@@ -6,7 +6,7 @@ export const useCounter = defineStore({
   id: 'counter',
 
   state: () => ({
-    n: 5,
+    n: 8,
     incrementedTimes: 0,
     decrementedTimes: 0,
     numbers: [] as number[],
@@ -25,8 +25,8 @@ export const useCounter = defineStore({
       this.n += amount
     },
 
-    newOne() {
-      console.log('neeeew')
+    changeMe() {
+      console.log('change me to test HMR')
     },
 
     async fail() {
@@ -55,7 +55,7 @@ export const useCounter = defineStore({
           //   state.decrementedTimes++
           // })
         } else {
-          this.n--
+          this.n -= 1
         }
         await delay(interval)
       }
@@ -104,7 +104,6 @@ if (import.meta.hot) {
           // TODO: replace the useCounter var???
           return
         }
-        console.log('patching')
         useStore(pinia, existingStore)
         // remove the existing store from the cache to force a new one
         // pinia._s.delete(id)
diff --git a/src/devtools/plugin.ts b/src/devtools/plugin.ts
index 17e21c0e..526a0b27 100644
--- a/src/devtools/plugin.ts
+++ b/src/devtools/plugin.ts
@@ -359,33 +359,24 @@ let runningActionId = 0
 let activeAction: number | undefined
 
 /**
- * pinia.use(devtoolsPlugin)
+ * Patches a store to enable action grouping in devtools by wrapping the store with a Proxy that is passed as the context of all actions, allowing us to set `runningAction` on each access and effectively associating any state mutation to the action.
+ *
+ * @param store - store to patch
+ * @param actionNames - list of actionst to patch
  */
-export function devtoolsPlugin<
-  Id extends string = string,
-  S extends StateTree = StateTree,
-  G extends GettersTree<S> = GettersTree<S>,
-  A /* extends ActionsTree */ = ActionsTree
->({ app, store, options, pinia }: PiniaPluginContext<Id, S, G, A>) {
-  // HMR module
-  if (store.$id.startsWith('__hot:')) {
-    return
-  }
+function patchActionForGrouping(store: Store, actionNames: string[]) {
   // original actions of the store as they are given by pinia. We are going to override them
-  const actions = Object.keys(options.actions).reduce(
-    (storeActions, actionName) => {
-      // @ts-expect-error
-      // use toRaw to avoid tracking #541
-      storeActions[actionName] = toRaw(store)[actionName]
-      return storeActions
-    },
-    {} as ActionsTree
-  )
+  const actions = actionNames.reduce((storeActions, actionName) => {
+    // use toRaw to avoid tracking #541
+    // @ts-expect-error
+    storeActions[actionName] = toRaw(store)[actionName]
+    return storeActions
+  }, {} as ActionsTree)
 
   for (const actionName in actions) {
     // @ts-expect-error
     store[actionName] = function () {
-      setActivePinia(pinia)
+      setActivePinia(store._p)
       // the running action id is incremented in a before action hook
       const _actionId = runningActionId
       const trackedStore = new Proxy(store, {
@@ -404,8 +395,38 @@ export function devtoolsPlugin<
       )
     }
   }
+}
 
-  // TODO: replace existing one for HMR?
+/**
+ * pinia.use(devtoolsPlugin)
+ */
+export function devtoolsPlugin<
+  Id extends string = string,
+  S extends StateTree = StateTree,
+  G extends GettersTree<S> = GettersTree<S>,
+  A /* extends ActionsTree */ = ActionsTree
+>({ app, store, options }: PiniaPluginContext<Id, S, G, A>) {
+  // HMR module
+  if (store.$id.startsWith('__hot:')) {
+    return
+  }
+
+  patchActionForGrouping(
+    // @ts-expect-error: can cast the store...
+    store,
+    Object.keys(options.actions)
+  )
+
+  const originalHotUpdate = store.hotUpdate
+
+  toRaw(store).hotUpdate = function (newStore) {
+    originalHotUpdate.apply(this, arguments as any)
+    patchActionForGrouping(
+      // @ts-expect-error: can cast the store...
+      store,
+      Object.keys(toRaw(newStore)._hmrPayload.actions)
+    )
+  }
 
   addDevtools(
     app,
diff --git a/src/store.ts b/src/store.ts
index 568f1b97..9e1fa9fa 100644
--- a/src/store.ts
+++ b/src/store.ts
@@ -437,42 +437,9 @@ function createSetupStore<
     set: (state) => (pinia.state.value[$id] = state),
   })
 
-  // apply all plugins
-  pinia._p.forEach((extender) => {
-    if (__DEV__ && IS_CLIENT) {
-      const extensions = extender({
-        // @ts-expect-error: conflict between A and ActionsTree
-        store,
-        app: pinia._a,
-        pinia,
-        // @ts-expect-error
-        options: optionsForPlugin,
-      })
-      Object.keys(extensions || {}).forEach((key) =>
-        store._customProperties.add(key)
-      )
-      assign(store, extensions)
-    } else {
-      assign(
-        store,
-        extender({
-          // @ts-expect-error: conflict between A and ActionsTree
-          store,
-          app: pinia._a,
-          pinia,
-          // @ts-expect-error
-          options: optionsForPlugin,
-        })
-      )
-    }
-  })
-
-  if (initialState) {
-    ;(options.hydrate || innerPatch)(store, initialState)
-  }
-
+  // add the hotUpdate before plugins to allow them to override it
   if (__DEV__) {
-    store.hotUpdate = (newStore) => {
+    store.hotUpdate = markRaw((newStore) => {
       newStore._hmrPayload.state.forEach((stateKey) => {
         if (!(stateKey in store.$state)) {
           console.log('setting new key', stateKey)
@@ -518,7 +485,41 @@ function createSetupStore<
       }
 
       // TODO: remove old actions and getters
+    })
+  }
+
+  // apply all plugins
+  pinia._p.forEach((extender) => {
+    if (__DEV__ && IS_CLIENT) {
+      const extensions = extender({
+        // @ts-expect-error: conflict between A and ActionsTree
+        store,
+        app: pinia._a,
+        pinia,
+        // @ts-expect-error
+        options: optionsForPlugin,
+      })
+      Object.keys(extensions || {}).forEach((key) =>
+        store._customProperties.add(key)
+      )
+      assign(store, extensions)
+    } else {
+      assign(
+        store,
+        extender({
+          // @ts-expect-error: conflict between A and ActionsTree
+          store,
+          app: pinia._a,
+          pinia,
+          // @ts-expect-error
+          options: optionsForPlugin,
+        })
+      )
     }
+  })
+
+  if (initialState) {
+    ;(options.hydrate || innerPatch)(store, initialState)
   }
 
   isListening = true
@@ -651,9 +652,6 @@ export function defineSetupStore<Id extends string, SS>(
       // cleanup the things
       delete pinia.state.value[hotId]
       pinia._s.delete(hotId)
-
-      // TODO: add the patched store to devtools again to override its previous version
-      // addDevtools(pinia._a, hot)
     }
 
     // save stores in instances to access them devtools
@@ -731,13 +729,16 @@ export function defineStore<
       // cleanup the things
       delete pinia.state.value[hotId]
       pinia._s.delete(hotId)
-
-      // TODO: add the patched store to devtools again to override its previous version
-      // addDevtools(pinia._a, hot)
     }
 
     // save stores in instances to access them devtools
-    if (__DEV__ && IS_CLIENT && currentInstance && currentInstance.proxy) {
+    if (
+      __DEV__ &&
+      IS_CLIENT &&
+      currentInstance &&
+      currentInstance.proxy &&
+      !hot
+    ) {
       const vm = currentInstance.proxy
       const cache = '_pStores' in vm ? vm._pStores! : (vm._pStores = {})
       // @ts-expect-error: still can't cast Store with generics to Store
-- 
2.47.2