feat(nuxt): automatic HMR code (vite only) (#2954)

Co-authored-by: Eduardo San Martin Morote <posva13@gmail.com>

Author: BobbieGoede

packages/nuxt/playground/pages/index.vue

@@ -16,8 +16,11 @@ if (import.meta.server) {
 
 <template>
   <div>
-    <p>Count: {{ counter.$state.count }}</p>
+    <p>Count: {{ counter.count }} x 2 = {{ counter.double }}</p>
     <button @click="counter.increment()">+</button>
+    <pre>{{ counter.$state }}</pre>
+
+    <hr />
 
     <p>Layer store: {{ layerStore.count }}</p>
   </div>

packages/nuxt/playground/stores/counter.ts

@@ -17,10 +17,6 @@ export const useCounter = defineStore('counter', {
     },
   },
   getters: {
-    getCount: (state) => state.count,
+    double: (state) => state.count * 2,
   },
 })
-
-if (import.meta.hot) {
-  import.meta.hot.accept(acceptHMRUpdate(useCounter, import.meta.hot))
-}

packages/nuxt/playground/stores/nested/some-store.ts

@@ -2,7 +2,3 @@ export const useSomeStoreStore = defineStore('some-store', () => {
   // console.log('I was defined within a nested store directory')
   return {}
 })
-
-if (import.meta.hot) {
-  import.meta.hot.accept(acceptHMRUpdate(useSomeStoreStore, import.meta.hot))
-}

packages/nuxt/playground/stores/with-skip-hydrate.ts

@@ -8,9 +8,3 @@ export const useWithSkipHydrateStore = defineStore('with-skip-hydrate', () => {
   )
   return { skipped }
 })
-
-if (import.meta.hot) {
-  import.meta.hot.accept(
-    acceptHMRUpdate(useWithSkipHydrateStore, import.meta.hot)
-  )
-}

packages/nuxt/src/auto-hmr-plugin.ts

@@ -0,0 +1,64 @@
+import type { VariableDeclarator } from 'estree'
+import type { Nuxt } from 'nuxt/schema'
+import type { Plugin } from 'vite'
+
+function getStoreDeclaration(nodes?: VariableDeclarator[]) {
+  return nodes?.find(
+    (x) =>
+      x.init?.type === 'CallExpression' &&
+      x.init.callee.type === 'Identifier' &&
+      x.init.callee.name === 'defineStore'
+  )
+}
+
+function nameFromDeclaration(node?: VariableDeclarator) {
+  return node?.id.type === 'Identifier' && node.id.name
+}
+
+export function autoRegisterHMRPlugin(rootDir: string) {
+  return {
+    name: 'pinia:auto-hmr-registration',
+
+    transform(code, id) {
+      if (id.startsWith('\x00')) return
+      if (!id.startsWith(rootDir)) return
+      if (!code.includes('defineStore') || code.includes('acceptHMRUpdate')) {
+        return
+      }
+
+      const ast = this.parse(code)
+
+      // walk top-level nodes
+      for (const n of ast.body) {
+        if (
+          n.type === 'VariableDeclaration' ||
+          n.type === 'ExportNamedDeclaration'
+        ) {
+          // find export or variable declaration that uses `defineStore`
+          const storeDeclaration = getStoreDeclaration(
+            n.type === 'VariableDeclaration'
+              ? n.declarations
+              : n.declaration?.type === 'VariableDeclaration'
+                ? n.declaration?.declarations
+                : undefined
+          )
+
+          // retrieve the variable name
+          const storeName = nameFromDeclaration(storeDeclaration)
+          if (storeName) {
+            // append HMR code
+            return {
+              code: [
+                `import { acceptHMRUpdate } from 'pinia'`,
+                code,
+                'if (import.meta.hot) {',
+                `  import.meta.hot.accept(acceptHMRUpdate(${storeName}, import.meta.hot))`,
+                '}',
+              ].join('\n'),
+            }
+          }
+        }
+      }
+    },
+  } satisfies Plugin
+}

packages/nuxt/src/module.ts

@@ -8,9 +8,11 @@ import {
   createResolver,
   addImportsDir,
   getLayerDirectories,
+  addVitePlugin,
 } from '@nuxt/kit'
 import type { NuxtModule } from '@nuxt/schema'
 import { fileURLToPath } from 'node:url'
+import { autoRegisterHMRPlugin } from './auto-hmr-plugin'
 
 export interface ModuleOptions {
   /**
@@ -82,6 +84,11 @@ const module: NuxtModule<ModuleOptions> = defineNuxtModule<ModuleOptions>({
         }
       }
     }
+
+    // Register automatic hmr code plugin - dev mode only
+    if (nuxt.options.dev) {
+      addVitePlugin(autoRegisterHMRPlugin(resolve(nuxt.options.rootDir)))
+    }
   },
 })