From 26b1c9ca0ece63d4ea6731d75b83fbe253022485 Mon Sep 17 00:00:00 2001
From: Daniel Roe <daniel@roe.dev>
Date: Mon, 10 Oct 2022 16:48:23 +0100
Subject: [PATCH] feat(nuxt): parse html to treeshake client-only components
 (#7527)

---
 packages/nuxt/package.json                 |  1 +
 packages/nuxt/src/components/tree-shake.ts | 49 ++++++++++++++++------
 yarn.lock                                  |  8 ++++
 3 files changed, 45 insertions(+), 13 deletions(-)

diff --git a/packages/nuxt/package.json b/packages/nuxt/package.json
index afc03527aa7..ff063a267f1 100644
--- a/packages/nuxt/package.json
+++ b/packages/nuxt/package.json
@@ -67,6 +67,7 @@
     "scule": "^0.3.2",
     "strip-literal": "^0.4.2",
     "ufo": "^0.8.5",
+    "ultrahtml": "^0.1.1",
     "unctx": "^2.0.2",
     "unenv": "^0.6.2",
     "unimport": "^0.6.8",
diff --git a/packages/nuxt/src/components/tree-shake.ts b/packages/nuxt/src/components/tree-shake.ts
index 23d869ddde9..263fb489750 100644
--- a/packages/nuxt/src/components/tree-shake.ts
+++ b/packages/nuxt/src/components/tree-shake.ts
@@ -1,6 +1,7 @@
 import { pathToFileURL } from 'node:url'
 import { parseURL } from 'ufo'
 import MagicString from 'magic-string'
+import { parse, walk, ELEMENT_NODE, Node } from 'ultrahtml'
 import { createUnplugin } from 'unplugin'
 import type { Component } from '@nuxt/schema'
 
@@ -9,8 +10,10 @@ interface TreeShakeTemplatePluginOptions {
   getComponents (): Component[]
 }
 
+const PLACEHOLDER_RE = /^(v-slot|#)(fallback|placeholder)/
+
 export const TreeShakeTemplatePlugin = createUnplugin((options: TreeShakeTemplatePluginOptions) => {
-  const regexpMap = new WeakMap<Component[], RegExp>()
+  const regexpMap = new WeakMap<Component[], [RegExp, string[]]>()
   return {
     name: 'nuxt:tree-shake-template',
     enforce: 'pre',
@@ -18,28 +21,48 @@ export const TreeShakeTemplatePlugin = createUnplugin((options: TreeShakeTemplat
       const { pathname } = parseURL(decodeURIComponent(pathToFileURL(id).href))
       return pathname.endsWith('.vue')
     },
-    transform (code, id) {
+    async transform (code, id) {
+      const template = code.match(/<template>([\s\S]*)<\/template>/)
+      if (!template) { return }
+
       const components = options.getComponents()
 
       if (!regexpMap.has(components)) {
         const clientOnlyComponents = components
           .filter(c => c.mode === 'client' && !components.some(other => other.mode !== 'client' && other.pascalName === c.pascalName))
-          .map(c => `${c.pascalName}|${c.kebabName}`)
-          .concat('ClientOnly|client-only')
-          .map(component => `<(${component})(| [^>]*)>[\\s\\S]*?<\\/(${component})>`)
+          .flatMap(c => [c.pascalName, c.kebabName])
+          .concat(['ClientOnly', 'client-only'])
+        const tags = clientOnlyComponents
+          .map(component => `<(${component})[^>]*>[\\s\\S]*?<\\/(${component})>`)
 
-        regexpMap.set(components, new RegExp(`(${clientOnlyComponents.join('|')})`, 'g'))
+        regexpMap.set(components, [new RegExp(`(${tags.join('|')})`, 'g'), clientOnlyComponents])
       }
 
-      const COMPONENTS_RE = regexpMap.get(components)!
+      const [COMPONENTS_RE, clientOnlyComponents] = regexpMap.get(components)!
+      if (!COMPONENTS_RE.test(code)) { return }
+
       const s = new MagicString(code)
 
-      // Do not render client-only slots on SSR, but preserve attributes and fallback/placeholder slots
-      s.replace(COMPONENTS_RE, r => r.replace(/<([^>]*[^/])\/?>[\s\S]*$/, (chunk: string, el: string) => {
-        const fallback = chunk.match(/<template[^>]*(#|v-slot:)(fallback|placeholder)[^>]*>[\s\S]*?<\/template>/)?.[0] || ''
-        const tag = el.split(' ').shift()
-        return `<${el}>${fallback}</${tag}>`
-      }))
+      const ast = parse(template[0])
+      await walk(ast, (node) => {
+        if (node.type !== ELEMENT_NODE || !clientOnlyComponents.includes(node.name) || !node.children?.length) {
+          return
+        }
+
+        const fallback = node.children.find(
+          (n: Node) => n.name === 'template' &&
+            Object.entries(n.attributes as Record<string, string>)?.flat().some(attr => PLACEHOLDER_RE.test(attr))
+        )
+
+        try {
+          // Replace node content
+          const text = fallback ? code.slice(template.index + fallback.loc[0].start, template.index + fallback.loc[fallback.loc.length - 1].end) : ''
+          s.overwrite(template.index + node.loc[0].end, template.index + node.loc[node.loc.length - 1].start, text)
+        } catch (err) {
+          // This may fail if we have a nested client-only component and are trying
+          // to replace some text that has already been replaced
+        }
+      })
 
       if (s.hasChanged()) {
         return {
diff --git a/yarn.lock b/yarn.lock
index 253f7e85087..6154a08f77f 100644
--- a/yarn.lock
+++ b/yarn.lock
@@ -10928,6 +10928,7 @@ __metadata:
     scule: ^0.3.2
     strip-literal: ^0.4.2
     ufo: ^0.8.5
+    ultrahtml: ^0.1.1
     unbuild: latest
     unctx: ^2.0.2
     unenv: ^0.6.2
@@ -13992,6 +13993,13 @@ __metadata:
   languageName: node
   linkType: hard
 
+"ultrahtml@npm:^0.1.1":
+  version: 0.1.1
+  resolution: "ultrahtml@npm:0.1.1"
+  checksum: a074b6d41e942b0ae3a4ea7372043bee32f5bad08009278adff04a96d784b316e68085d5f44b52c5c03ac0bf1d2a142090e58e555f953d9c69eb9639fd2f70d5
+  languageName: node
+  linkType: hard
+
 "unbox-primitive@npm:^1.0.2":
   version: 1.0.2
   resolution: "unbox-primitive@npm:1.0.2"