From d23a4d3805eb0727f8b7c03cb8e36c0bfb7873e2 Mon Sep 17 00:00:00 2001
From: Eduardo San Martin Morote <posva13@gmail.com>
Date: Mon, 12 Jul 2021 18:25:13 +0200
Subject: [PATCH] feat(dx): HMR wip

---
 playground/src/App.vue           |  10 +-
 playground/src/main.ts           |  54 +++++-
 playground/src/stores/counter.ts |  63 ++++++-
 playground/tsconfig.json         |  14 ++
 src/createPinia.ts               |   8 +-
 src/devtools/plugin.ts           |   6 +
 src/store.ts                     | 273 ++++++++++++++++++++++++-------
 src/types.ts                     |  28 +++-
 8 files changed, 384 insertions(+), 72 deletions(-)
 create mode 100644 playground/tsconfig.json

diff --git a/playground/src/App.vue b/playground/src/App.vue
index f151f0a9..3edb3052 100644
--- a/playground/src/App.vue
+++ b/playground/src/App.vue
@@ -1,9 +1,18 @@
 <template>
+  <button @click="n++">Increment {{ n }}</button>
+  <pre>{{ counter.$state }}</pre>
+  <!-- <button @click="counter.newOne()">Click me</button> -->
   <TestStore />
 </template>
 
 <script lang="ts" setup>
+import { ref } from 'vue'
 import TestStore from './components/TestStore.vue'
+import { useCounter } from './stores/counter'
+
+const counter = useCounter()
+
+const n = ref(4)
 </script>
 
 <style>
@@ -11,7 +20,6 @@ import TestStore from './components/TestStore.vue'
   font-family: Avenir, Helvetica, Arial, sans-serif;
   -webkit-font-smoothing: antialiased;
   -moz-osx-font-smoothing: grayscale;
-  text-align: center;
   color: #2c3e50;
   margin-top: 60px;
 }
diff --git a/playground/src/main.ts b/playground/src/main.ts
index 1d3ea99e..f31805e7 100644
--- a/playground/src/main.ts
+++ b/playground/src/main.ts
@@ -1,5 +1,55 @@
 import { createApp } from 'vue'
 import App from './App.vue'
-import { createPinia } from '../../src'
+import { createPinia, StoreDefinition } from '../../src'
+import { ha } from './test'
 
-createApp(App).use(createPinia()).mount('#app')
+console.log({ ha })
+
+const pinia = createPinia()
+
+if (import.meta.hot) {
+  import.meta.hot.data.pinia = pinia
+  console.log('set', import.meta.hot.data)
+  //   const isUseStore = (fn: any): fn is StoreDefinition => {
+  //     return typeof fn === 'function' && typeof fn.$id === 'string'
+  //   }
+
+  //   // import.meta.hot.accept(
+  //   //   './stores/counter.ts',
+  //   //   (newStore: Record<string, unknown>) => {
+  //   //     console.log('haha', newStore)
+  //   //   }
+  //   // )
+
+  //   import.meta.hot.accept('./test.ts', (newTest) => {
+  //     console.log('test updated', newTest)
+  //   })
+
+  //   const stores = import.meta.glob('./stores/*.ts')
+
+  //   for (const storeId in stores) {
+  //     console.log('configuring HMR for', storeId)
+  //     const oldUseStore = await stores[storeId]()
+  //     console.log('got', oldUseStore)
+  //     import.meta.hot!.accept(storeId, (newStore: Record<string, unknown>) => {
+  //       console.log('Accepting update for', storeId)
+  //       for (const exportName in newStore) {
+  //         const useStore = newStore[exportName]
+  //         if (isUseStore(useStore) && pinia._s.has(useStore.$id)) {
+  //           const id = useStore.$id
+  //           const existingStore = pinia._s.get(id)!
+  //           // remove the existing store from the cache to force a new one
+  //           pinia._s.delete(id)
+  //           // this adds any new state to pinia and then runs the `hydrate` function
+  //           // which, by default, will reuse the existing state in pinia
+  //           const newStore = useStore(pinia)
+  //           // pinia._s.set(id, existingStore)
+  //         }
+  //       }
+  //     })
+  //   }
+}
+
+// TODO: HMR for plugins
+
+createApp(App).use(pinia).mount('#app')
diff --git a/playground/src/stores/counter.ts b/playground/src/stores/counter.ts
index c4827136..e9008527 100644
--- a/playground/src/stores/counter.ts
+++ b/playground/src/stores/counter.ts
@@ -1,4 +1,4 @@
-import { defineStore } from '../../../src'
+import { defineStore, Pinia, Store, StoreDefinition } from '../../../src'
 
 const delay = (t: number) => new Promise((r) => setTimeout(r, t))
 
@@ -6,7 +6,7 @@ export const useCounter = defineStore({
   id: 'counter',
 
   state: () => ({
-    n: 0,
+    n: 5,
     incrementedTimes: 0,
     decrementedTimes: 0,
     numbers: [] as number[],
@@ -25,6 +25,10 @@ export const useCounter = defineStore({
       this.n += amount
     },
 
+    newOne() {
+      console.log('neeeew')
+    },
+
     async fail() {
       const n = this.n
       await delay(1000)
@@ -58,3 +62,58 @@ export const useCounter = defineStore({
     },
   },
 })
+
+if (import.meta.hot) {
+  const isUseStore = (fn: any): fn is StoreDefinition => {
+    return typeof fn === 'function' && typeof fn.$id === 'string'
+  }
+
+  const oldUseStore = useCounter
+  import.meta.hot.accept((newStore) => {
+    if (!import.meta.hot) throw new Error('import.meta.hot disappeared')
+
+    const pinia: Pinia | undefined =
+      import.meta.hot.data.pinia || oldUseStore._pinia
+
+    if (!pinia) {
+      console.warn(`Missing the pinia instance for "${oldUseStore.$id}".`)
+      return import.meta.hot.invalidate()
+    }
+
+    // preserve the pinia instance across loads
+    import.meta.hot.data.pinia = pinia
+
+    console.log('got data', newStore)
+    for (const exportName in newStore) {
+      const useStore = newStore[exportName]
+      console.log('checking for', exportName)
+      if (isUseStore(useStore) && pinia._s.has(useStore.$id)) {
+        console.log('Accepting update for', useStore.$id)
+        const id = useStore.$id
+
+        if (id !== oldUseStore.$id) {
+          console.warn(
+            `The id of the store changed from "${oldUseStore.$id}" to "${id}". Reloading.`
+          )
+          return import.meta.hot!.invalidate()
+        }
+
+        const existingStore: Store = pinia._s.get(id)!
+        if (!existingStore) {
+          console.log(`skipping hmr because store doesn't exist yet`)
+          // 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)
+        // this adds any new state to pinia and then runs the `hydrate` function
+        // which, by default, will reuse the existing state in pinia
+        // const newStore = useStore(pinia)
+        // console.log('going there', newStore._hmrPayload)
+        // pinia._s.set(id, existingStore)
+      }
+    }
+  })
+}
diff --git a/playground/tsconfig.json b/playground/tsconfig.json
new file mode 100644
index 00000000..5b6c24df
--- /dev/null
+++ b/playground/tsconfig.json
@@ -0,0 +1,14 @@
+{
+  "compilerOptions": {
+    "target": "esnext",
+    "module": "esnext",
+    "moduleResolution": "node",
+    "strict": true,
+    "jsx": "preserve",
+    "sourceMap": true,
+    "resolveJsonModule": true,
+    "esModuleInterop": true,
+    "lib": ["esnext", "dom"]
+  },
+  "include": ["src/**/*.ts", "src/**/*.d.ts", "src/**/*.tsx", "src/**/*.vue"]
+}
diff --git a/src/createPinia.ts b/src/createPinia.ts
index aaaee6fd..e0ecb77c 100644
--- a/src/createPinia.ts
+++ b/src/createPinia.ts
@@ -18,14 +18,13 @@ export function createPinia(): Pinia {
   // if there is anything like it with Vue 3 SSR
   const state = scope.run(() => ref<Record<string, StateTree>>({}))!
 
-  let localApp: App | undefined
   let _p: Pinia['_p'] = []
   // plugins added before calling app.use(pinia)
   const toBeInstalled: PiniaStorePlugin[] = []
 
   const pinia: Pinia = markRaw({
     install(app: App) {
-      pinia._a = localApp = app
+      pinia._a = app
       app.provide(piniaSymbol, pinia)
       app.config.globalProperties.$pinia = pinia
       if (IS_CLIENT) {
@@ -37,7 +36,7 @@ export function createPinia(): Pinia {
     },
 
     use(plugin) {
-      if (!localApp) {
+      if (!this._a) {
         toBeInstalled.push(plugin)
       } else {
         _p.push(plugin)
@@ -47,7 +46,8 @@ export function createPinia(): Pinia {
 
     _p,
     // it's actually undefined here
-    _a: localApp!,
+    // @ts-expect-error
+    _a: null,
     _e: scope,
     _s: new Map<string, Store>(),
     state,
diff --git a/src/devtools/plugin.ts b/src/devtools/plugin.ts
index ade52b4e..17e21c0e 100644
--- a/src/devtools/plugin.ts
+++ b/src/devtools/plugin.ts
@@ -367,6 +367,10 @@ export function devtoolsPlugin<
   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
+  }
   // 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) => {
@@ -401,6 +405,8 @@ export function devtoolsPlugin<
     }
   }
 
+  // TODO: replace existing one for HMR?
+
   addDevtools(
     app,
     // @ts-expect-error: FIXME: if possible...
diff --git a/src/store.ts b/src/store.ts
index c8e2307f..568f1b97 100644
--- a/src/store.ts
+++ b/src/store.ts
@@ -16,6 +16,8 @@ import {
   ComputedRef,
   toRef,
   toRefs,
+  Ref,
+  ref,
 } from 'vue'
 import {
   StateTree,
@@ -43,6 +45,7 @@ import {
   activePinia,
 } from './rootStore'
 import { IS_CLIENT } from './env'
+import { createPinia } from './createPinia'
 
 function innerPatch<T extends StateTree>(
   target: T,
@@ -79,18 +82,26 @@ function createOptionsStore<
   S extends StateTree,
   G extends GettersTree<S>,
   A extends ActionsTree
->(options: DefineStoreOptions<Id, S, G, A>, pinia: Pinia): Store<Id, S, G, A> {
+>(
+  options: DefineStoreOptions<Id, S, G, A>,
+  pinia: Pinia,
+  hot?: boolean
+): Store<Id, S, G, A> {
   const { id, state, actions, getters } = options
   function $reset() {
     pinia.state.value[id] = state ? state() : {}
   }
 
+  const initialState: StateTree | undefined = pinia.state.value[id]
+
   function setup() {
-    $reset()
+    if (!initialState) {
+      $reset()
+    }
     // pinia.state.value[id] = state ? state() : {}
 
     return assign(
-      toRefs(pinia.state.value[id]),
+      initialState || toRefs(pinia.state.value[id]),
       actions,
       Object.keys(getters || {}).reduce((computedGetters, name) => {
         computedGetters[name] = computed(() => {
@@ -103,7 +114,9 @@ function createOptionsStore<
     )
   }
 
-  const store = createSetupStore(id, setup, options)
+  const store = createSetupStore(id, setup, options, hot)
+
+  // TODO: HMR should also replace getters here
 
   store.$reset = $reset
 
@@ -123,7 +136,8 @@ function createSetupStore<
   setup: () => SS,
   options:
     | DefineSetupStoreOptions<Id, S, G, A>
-    | DefineStoreOptions<Id, S, G, A> = {}
+    | DefineStoreOptions<Id, S, G, A> = {},
+  hot?: boolean
 ): Store<Id, S, G, A> {
   const pinia = getActivePinia()
   let scope!: EffectScope
@@ -185,22 +199,25 @@ function createSetupStore<
     return scope.run(() => {
       const store = setup()
 
-      watch(
-        () => pinia.state.value[$id] as UnwrapRef<S>,
-        (state, oldState) => {
-          if (isListening) {
-            triggerSubscriptions(
-              {
-                storeId: $id,
-                type: MutationType.direct,
-                events: debuggerEvents as DebuggerEvent,
-              },
-              state
-            )
-          }
-        },
-        $subscribeOptions
-      )!
+      // skip setting up the watcher on HMR
+      if (!__DEV__ || !hot) {
+        watch(
+          () => pinia.state.value[$id] as UnwrapRef<S>,
+          (state, oldState) => {
+            if (isListening) {
+              triggerSubscriptions(
+                {
+                  storeId: $id,
+                  type: MutationType.direct,
+                  events: debuggerEvents as DebuggerEvent,
+                },
+                state
+              )
+            }
+          },
+          $subscribeOptions
+        )!
+      }
 
       return store
     })
@@ -294,6 +311,58 @@ function createSetupStore<
     }
   }
 
+  /**
+   * Wraps an action to handle subscriptions
+   *
+   * @param name - name of the action
+   * @param action - action to wrap
+   * @returns a wrapped action to handle subscriptions
+   */
+  function wrapAction(name: string, action: _Method) {
+    return function (this: any) {
+      setActivePinia(pinia)
+      const args = Array.from(arguments)
+
+      let afterCallback: (resolvedReturn: any) => void = noop
+      let onErrorCallback: (error: unknown) => void = noop
+      function after(callback: typeof afterCallback) {
+        afterCallback = callback
+      }
+      function onError(callback: typeof onErrorCallback) {
+        onErrorCallback = callback
+      }
+
+      actionSubscriptions.forEach((callback) => {
+        // @ts-expect-error
+        callback({
+          args,
+          name,
+          store,
+          after,
+          onError,
+        })
+      })
+
+      let ret: any
+      try {
+        ret = action.apply(this || store, args)
+        Promise.resolve(ret).then(afterCallback).catch(onErrorCallback)
+      } catch (error) {
+        onErrorCallback(error)
+        throw error
+      }
+
+      return ret
+    }
+  }
+
+  // TODO: PURE to tree shake?
+  const _hmrPayload = markRaw({
+    actions: {} as Record<string, any>,
+    getters: {} as Record<string, Ref>,
+    state: [] as string[],
+  })
+
   // overwrite existing actions to support $onAction
   for (const key in setupStore) {
     const prop = setupStore[key]
@@ -304,54 +373,35 @@ function createSetupStore<
         // mark it as a piece of state to be serialized
         pinia.state.value[$id][key] = toRef(setupStore as any, key)
       }
+      if (__DEV__) {
+        _hmrPayload.state.push(key)
+      }
       // action
     } else if (typeof prop === 'function') {
       // @ts-expect-error: we are overriding the function
-      setupStore[key] = function () {
-        setActivePinia(pinia)
-        const args = Array.from(arguments)
-
-        let afterCallback: (resolvedReturn: any) => void = noop
-        let onErrorCallback: (error: unknown) => void = noop
-        function after(callback: typeof afterCallback) {
-          afterCallback = callback
-        }
-        function onError(callback: typeof onErrorCallback) {
-          onErrorCallback = callback
-        }
+      // we avoid wrapping if this a hot module replacement store
+      setupStore[key] = __DEV__ && hot ? prop : wrapAction(key, prop)
 
-        actionSubscriptions.forEach((callback) => {
-          // @ts-expect-error
-          callback({
-            args,
-            name: key,
-            store,
-            after,
-            onError,
-          })
-        })
-
-        let ret: any
-        try {
-          ret = prop.apply(this || store, args)
-          Promise.resolve(ret).then(afterCallback).catch(onErrorCallback)
-        } catch (error) {
-          onErrorCallback(error)
-          throw error
-        }
-
-        return ret
+      if (__DEV__) {
+        _hmrPayload.actions[key] = prop
       }
+
       // list actions so they can be used in plugins
       // @ts-expect-error
-      optionsForPlugin.actions[key] = prop
-    } else if (__DEV__ && IS_CLIENT) {
+      optionsForPlugin.actions[key] = setupStore[key] // TODO: check this change from `prop` is correct
+    } else if (__DEV__) {
       // add getters for devtools
       if (isComputed(prop)) {
-        const getters: string[] =
-          // @ts-expect-error: it should be on the store
-          setupStore._getters || (setupStore._getters = markRaw([]))
-        getters.push(key)
+        _hmrPayload.getters[key] = buildState
+          ? // @ts-expect-error
+            options.getters[key]
+          : prop
+        if (IS_CLIENT) {
+          const getters: string[] =
+            // @ts-expect-error: it should be on the store
+            setupStore._getters || (setupStore._getters = markRaw([]))
+          getters.push(key)
+        }
       }
     }
   }
@@ -371,6 +421,7 @@ function createSetupStore<
         ? // devtools custom properties
           {
             _customProperties: markRaw(new Set<string>()),
+            _hmrPayload,
           }
         : {},
       partialStore,
@@ -420,6 +471,56 @@ function createSetupStore<
     ;(options.hydrate || innerPatch)(store, initialState)
   }
 
+  if (__DEV__) {
+    store.hotUpdate = (newStore) => {
+      newStore._hmrPayload.state.forEach((stateKey) => {
+        if (!(stateKey in store.$state)) {
+          console.log('setting new key', stateKey)
+          // @ts-expect-error
+          // transfer the ref
+          store.$state[stateKey] = newStore.$state[stateKey]
+        }
+      })
+
+      // 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]
+        }
+      })
+
+      for (const actionName in newStore._hmrPayload.actions) {
+        const action: _Method =
+          // @ts-expect-error
+          newStore[actionName]
+
+        // @ts-expect-error: new key
+        store[actionName] =
+          // new line forced for TS
+          wrapAction(actionName, action)
+      }
+
+      for (const getterName in newStore._hmrPayload.getters) {
+        const getter: _Method = newStore._hmrPayload.getters[getterName]
+
+        // @ts-expect-error
+        store[getterName] =
+          // ---
+          buildState
+            ? // special handling of options api
+              computed(() => {
+                setActivePinia(pinia)
+                return getter.call(store, store)
+              })
+            : getter
+      }
+
+      // TODO: remove old actions and getters
+    }
+  }
+
   isListening = true
   return store
 }
@@ -500,7 +601,8 @@ export function defineSetupStore<Id extends string, SS>(
   _ExtractActionsFromSetupStore<SS>
 > {
   function useStore(
-    pinia?: Pinia | null
+    pinia?: Pinia | null,
+    hot?: Store
   ): Store<
     Id,
     _ExtractStateFromSetupStore<SS>,
@@ -519,6 +621,10 @@ export function defineSetupStore<Id extends string, SS>(
 
     if (!pinia._s.has(id)) {
       pinia._s.set(id, createSetupStore(id, storeSetup, options))
+      if (__DEV__) {
+        // @ts-expect-error: not the right inferred type
+        useStore._pinia = pinia
+      }
     }
 
     const store: Store<
@@ -533,6 +639,23 @@ export function defineSetupStore<Id extends string, SS>(
       _ExtractActionsFromSetupStore<SS>
     >
 
+    if (__DEV__ && hot) {
+      const hotId = '__hot:' + id
+      const newStore = createSetupStore(hotId, storeSetup, options, true)
+      hot.hotUpdate(newStore as any)
+
+      // for state that exists in newStore, try to copy from old state
+
+      // actions
+
+      // 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) {
       const vm = currentInstance.proxy
@@ -563,7 +686,7 @@ export function defineStore<
 >(options: DefineStoreOptions<Id, S, G, A>): StoreDefinition<Id, S, G, A> {
   const { id } = options
 
-  function useStore(pinia?: Pinia | null) {
+  function useStore(pinia?: Pinia | null, hot?: Store) {
     const currentInstance = getCurrentInstance()
     pinia =
       // in test mode, ignore the argument provided as we can always retrieve a
@@ -583,10 +706,36 @@ export function defineStore<
           pinia
         )
       )
+
+      if (__DEV__) {
+        // @ts-expect-error: not the right inferred type
+        useStore._pinia = pinia
+      }
     }
 
     const store: Store<Id, S, G, A> = pinia._s.get(id)! as Store<Id, S, G, A>
 
+    if (__DEV__ && hot) {
+      const hotId = '__hot:' + id
+      const newStore = createOptionsStore(
+        assign({}, options, { id: hotId }) as any,
+        pinia,
+        true
+      )
+      hot.hotUpdate(newStore as any)
+
+      // for state that exists in newStore, try to copy from old state
+
+      // actions
+
+      // 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) {
       const vm = currentInstance.proxy
diff --git a/src/types.ts b/src/types.ts
index 90163db2..2d269b95 100644
--- a/src/types.ts
+++ b/src/types.ts
@@ -251,6 +251,24 @@ export interface StoreWithState<
    */
   _customProperties: Set<string>
 
+  /**
+   * Handles a HMR replacement of this store. Dev Only.
+   *
+   * @internal
+   */
+  hotUpdate(useStore: Store<Id, S, G, A>): void
+
+  /**
+   * Payload of the hmr update. Dev only.
+   *
+   * @internal
+   */
+  _hmrPayload: {
+    state: string[]
+    actions: ActionsTree
+    getters: ActionsTree
+  }
+
   /**
    * Applies a state patch to current state. Allows passing nested values
    *
@@ -405,13 +423,21 @@ export interface StoreDefinition<
    * Returns a store, creates it if necessary.
    *
    * @param pinia - Pinia instance to retrieve the store
+   * @param hot - dev only hot module replacement
    */
-  (pinia?: Pinia | null | undefined): Store<Id, S, G, A>
+  (pinia?: Pinia | null | undefined, hot?: Store): Store<Id, S, G, A>
 
   /**
    * Id of the store. Used by map helpers.
    */
   $id: Id
+
+  /**
+   * Dev only pinia for HMR.
+   *
+   * @internal
+   */
+  _pinia?: Pinia
 }
 
 /**
-- 
2.47.2