fix(nuxt): use parser to treeshake <ClientOnly>

packages/nuxt/package.json

@@ -70,7 +70,6 @@
     "scule": "^1.0.0",
     "strip-literal": "^1.0.1",
     "ufo": "^1.0.1",
-    "ultrahtml": "^1.2.0",
     "unctx": "^2.1.1",
     "unenv": "^1.1.0",
     "unhead": "^1.0.21",

packages/nuxt/src/components/module.ts

@@ -192,35 +192,35 @@ export default defineNuxtModule<ComponentsOptions>({
       const mode = isClient ? 'client' : 'server'
 
       config.plugins = config.plugins || []
-      config.plugins.push(loaderPlugin.vite({
-        sourcemap: nuxt.options.sourcemap[mode],
-        getComponents,
-        mode,
-        experimentalComponentIslands: nuxt.options.experimental.componentIslands
-      }))
       if (nuxt.options.experimental.treeshakeClientOnly && isServer) {
         config.plugins.push(TreeShakeTemplatePlugin.vite({
           sourcemap: nuxt.options.sourcemap[mode],
           getComponents
         }))
       }
+      config.plugins.push(loaderPlugin.vite({
+        sourcemap: nuxt.options.sourcemap[mode],
+        getComponents,
+        mode,
+        experimentalComponentIslands: nuxt.options.experimental.componentIslands
+      }))
     })
     nuxt.hook('webpack:config', (configs) => {
       configs.forEach((config) => {
         const mode = config.name === 'client' ? 'client' : 'server'
         config.plugins = config.plugins || []
-        config.plugins.push(loaderPlugin.webpack({
-          sourcemap: nuxt.options.sourcemap[mode],
-          getComponents,
-          mode,
-          experimentalComponentIslands: nuxt.options.experimental.componentIslands
-        }))
         if (nuxt.options.experimental.treeshakeClientOnly && mode === 'server') {
           config.plugins.push(TreeShakeTemplatePlugin.webpack({
             sourcemap: nuxt.options.sourcemap[mode],
             getComponents
           }))
         }
+        config.plugins.push(loaderPlugin.webpack({
+          sourcemap: nuxt.options.sourcemap[mode],
+          getComponents,
+          mode,
+          experimentalComponentIslands: nuxt.options.experimental.componentIslands
+        }))
       })
     })
   }

packages/nuxt/src/components/tree-shake.ts

@@ -1,9 +1,10 @@
 import { pathToFileURL } from 'node:url'
 import { parseURL } from 'ufo'
 import MagicString from 'magic-string'
-import type { Node } from 'ultrahtml'
-import { parse, walk, ELEMENT_NODE } from 'ultrahtml'
+import { walk } from 'estree-walker'
+import type { CallExpression, Property, Identifier, ImportDeclaration, MemberExpression, Literal, ReturnStatement, VariableDeclaration, ObjectExpression, Node } from 'estree'
 import { createUnplugin } from 'unplugin'
+import escapeStringRegexp from 'escape-string-regexp'
 import type { Component } from '@nuxt/schema'
 import { resolve } from 'pathe'
 import { distDir } from '../dirs'
@@ -13,57 +14,115 @@ interface TreeShakeTemplatePluginOptions {
   getComponents (): Component[]
 }
 
-const PLACEHOLDER_RE = /^(v-slot|#)(fallback|placeholder)/
+type AcornNode<N> = N & { start: number, end: number }
+
+const SSR_RENDER_RE = /ssrRenderComponent/
+const PLACEHOLDER_EXACT_RE = /^(fallback|placeholder)$/
 
 export const TreeShakeTemplatePlugin = createUnplugin((options: TreeShakeTemplatePluginOptions) => {
-  const regexpMap = new WeakMap<Component[], [RegExp, string[]]>()
+  const regexpMap = new WeakMap<Component[], [RegExp, RegExp, string[]]>()
   return {
     name: 'nuxt:tree-shake-template',
-    enforce: 'pre',
+    enforce: 'post',
     transformInclude (id) {
       const { pathname } = parseURL(decodeURIComponent(pathToFileURL(id).href))
       return pathname.endsWith('.vue')
     },
-    async transform (code, id) {
-      const template = code.match(/<template>([\s\S]*)<\/template>/)
-      if (!template) { return }
-
+    transform (code, id) {
       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 && other.filePath !== resolve(distDir, 'app/components/server-placeholder')))
-          .flatMap(c => [c.pascalName, c.kebabName])
-          .concat(['ClientOnly', 'client-only'])
-        const tags = clientOnlyComponents
-          .map(component => `<(${component})[^>]*>[\\s\\S]*?<\\/(${component})>`)
+          .flatMap(c => [c.pascalName, c.kebabName.replaceAll('-', '_')])
+          .concat(['ClientOnly', 'client_only'])
 
-        regexpMap.set(components, [new RegExp(`(${tags.join('|')})`, 'g'), clientOnlyComponents])
+        regexpMap.set(components, [new RegExp(`(${clientOnlyComponents.join('|')})`), new RegExp(`^(${clientOnlyComponents.map(c => `(?:(?:_unref\\()?(?:_component_)?(?:Lazy|lazy_)?${c}\\)?)`).join('|')})$`), clientOnlyComponents])
       }
 
-      const [COMPONENTS_RE, clientOnlyComponents] = regexpMap.get(components)!
+      const s = new MagicString(code)
+      const importDeclarations: AcornNode<ImportDeclaration>[] = []
+
+      const [COMPONENTS_RE, COMPONENTS_IDENTIFIERS_RE] = regexpMap.get(components)!
       if (!COMPONENTS_RE.test(code)) { return }
 
-      const s = new MagicString(code)
+      walk(this.parse(code, { sourceType: 'module', ecmaVersion: 'latest' }) as Node, {
+        enter: (_node) => {
+          const node = _node as AcornNode<CallExpression | ImportDeclaration>
+          if (node.type === 'ImportDeclaration') {
+            importDeclarations.push(node)
+          } else if (
+            node.type === 'CallExpression' &&
+            node.callee.type === 'Identifier' &&
+            SSR_RENDER_RE.test(node.callee.name)
+          ) {
+            const [componentCall, _, children] = node.arguments
+            if (componentCall.type === 'Identifier' || componentCall.type === 'MemberExpression' || componentCall.type === 'CallExpression') {
+              const componentName = getComponentName(node)
+              const isClientComponent = COMPONENTS_IDENTIFIERS_RE.test(componentName)
+              const isClientOnlyComponent = /^(?:_unref\()?(?:_component_)?(?:Lazy|lazy_)?(?:client_only|ClientOnly\)?)$/.test(componentName)
+              if (isClientComponent && children?.type === 'ObjectExpression') {
+                const slotsToRemove = isClientOnlyComponent ? children.properties.filter(prop => prop.type === 'Property' && prop.key.type === 'Identifier' && !PLACEHOLDER_EXACT_RE.test(prop.key.name)) as AcornNode<Property>[] : children.properties as AcornNode<Property>[]
 
-      const ast = parse(template[0])
-      await walk(ast, (node) => {
-        if (node.type !== ELEMENT_NODE || !clientOnlyComponents.includes(node.name) || !node.children?.length) {
-          return
-        }
+                for (const slot of slotsToRemove) {
+                  const componentsSet = new Set<string>()
+                  s.remove(slot.start, slot.end + 1)
+                  const removedCode = `({${code.slice(slot.start, slot.end + 1)}})`
+                  const currentCode = s.toString()
+                  walk(this.parse(removedCode, { sourceType: 'module', ecmaVersion: 'latest' }) as Node, {
+                    enter: (_node) => {
+                      const node = _node as AcornNode<CallExpression>
+                      if (node.type === 'CallExpression' && node.callee.type === 'Identifier' && SSR_RENDER_RE.test(node.callee.name)) {
+                        const componentNode = node.arguments[0]
+
+                        if (componentNode.type === 'CallExpression') {
+                          const identifier = componentNode.arguments[0] as Identifier
+                          if (!isRenderedInCode(currentCode, removedCode.slice((componentNode as AcornNode<CallExpression>).start, (componentNode as AcornNode<CallExpression>).end))) { componentsSet.add(identifier.name) }
+                        } else if (componentNode.type === 'Identifier' && !isRenderedInCode(currentCode, componentNode.name)) {
+                          componentsSet.add(componentNode.name)
+                        } else if (componentNode.type === 'MemberExpression') {
+                          // expect componentNode to be a memberExpression (mostly used in dev with $setup[])
+                          const { start, end } = componentNode as AcornNode<MemberExpression>
+                          if (!isRenderedInCode(currentCode, removedCode.slice(start, end))) {
+                            componentsSet.add(((componentNode as MemberExpression).property as Literal).value as string)
+                            // remove the component from the return statement of `setup()`
+                            walk(this.parse(code, { sourceType: 'module', ecmaVersion: 'latest' }) as Node, {
+                              enter: (node) => {
+                                removeFromSetupReturnStatement(s, node as Property, ((componentNode as MemberExpression).property as Literal).value as string)
+                              }
+                            })
+                          }
+                        }
+                      }
+                    }
+                  })
+                  const componentsToRemove = [...componentsSet]
+                  for (const componentName of componentsToRemove) {
+                    let removed = false
+                    // remove const _component_ = resolveComponent...
+                    const VAR_RE = new RegExp(`(?:const|let|var) ${componentName} = ([^;\\n]*);?`)
+                    s.replace(VAR_RE, () => {
+                      removed = true
+                      return ''
+                    })
+                    if (!removed) {
+                      // remove direct import
+                      const declaration = findImportDeclaration(importDeclarations, componentName)
+                      if (declaration) {
+                        if (declaration.specifiers.length > 1) {
+                          const componentSpecifier = declaration.specifiers.find(s => s.local.name === componentName) as AcornNode<Identifier> | undefined
 
-        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 (componentSpecifier) { s.remove(componentSpecifier.start, componentSpecifier.end + 1) }
+                        } else {
+                          s.remove(declaration.start, declaration.end)
+                        }
+                      }
+                    }
+                  }
+                }
+              }
+            }
+          }
         }
       })
 
@@ -78,3 +137,63 @@ export const TreeShakeTemplatePlugin = createUnplugin((options: TreeShakeTemplat
     }
   }
 })
+
+/**
+ * find and return the importDeclaration that contain the import specifier
+ *
+ * @param {AcornNode<ImportDeclaration>[]} declarations - list of import declarations
+ * @param {string} importName - name of the import
+ */
+function findImportDeclaration (declarations: AcornNode<ImportDeclaration>[], importName: string): AcornNode<ImportDeclaration> | undefined {
+  const declaration = declarations.find((d) => {
+    const specifier = d.specifiers.find(s => s.local.name === importName)
+    if (specifier) { return true }
+    return false
+  })
+
+  return declaration
+}
+
+/**
+ * test if the name argument is used to render a component in the code
+ *
+ * @param code code to test
+ * @param name component name
+ */
+function isRenderedInCode (code: string, name: string) {
+  return new RegExp(`ssrRenderComponent\\(${escapeStringRegexp(name)}`).test(code)
+}
+
+/**
+ * retrieve the component identifier being used on ssrRender callExpression
+ *
+ * @param {CallExpression} ssrRenderNode - ssrRender callExpression
+ */
+function getComponentName (ssrRenderNode: AcornNode<CallExpression>): string {
+  const componentCall = ssrRenderNode.arguments[0] as Identifier | MemberExpression | CallExpression
+
+  if (componentCall.type === 'Identifier') {
+    return componentCall.name
+  } else if (componentCall.type === 'MemberExpression') {
+    return (componentCall.property as Literal).value as string
+  }
+  return (componentCall.arguments[0] as Identifier).name
+}
+
+/**
+ * remove a key from the return statement of the setup function
+ */
+function removeFromSetupReturnStatement (s: MagicString, node: Property, name: string) {
+  if (node.type === 'Property' && node.key.type === 'Identifier' && node.key.name === 'setup' && node.value.type === 'FunctionExpression') {
+    const returnStatement = node.value.body.body.find(n => n.type === 'ReturnStatement') as ReturnStatement | undefined
+    if (returnStatement?.argument?.type === 'Identifier') {
+      const returnIdentifier = returnStatement.argument.name
+      const returnedDeclaration = node.value.body.body.find(n => n.type === 'VariableDeclaration' && (n.declarations[0].id as Identifier).name === returnIdentifier) as AcornNode<VariableDeclaration>
+      const componentProperty = (returnedDeclaration?.declarations[0].init as ObjectExpression)?.properties.find(p => ((p as Property).key as Identifier).name === name) as AcornNode<Property>
+      if (componentProperty) { s.remove(componentProperty.start, componentProperty.end + 1) }
+    } else if (returnStatement?.argument?.type === 'ObjectExpression') {
+      const componentProperty = returnStatement.argument?.properties.find(p => ((p as Property).key as Identifier).name === name) as AcornNode<Property>
+      if (componentProperty) { s.remove(componentProperty.start, componentProperty.end + 1) }
+    }
+  }
+}

packages/nuxt/test/fixture/components/client/WithClientOnlySetup.vue

@@ -0,0 +1,38 @@
+<template>
+  <div>
+    <div>
+      <Glob />
+    </div>
+    {{ hello }}
+    <div class="not-client">
+      Hello
+    </div>
+    <ClientOnly>
+      <HelloWorld />
+      <Glob />
+      <SomeGlob />
+    </ClientOnly>
+    <ClientOnly>
+      <div class="should-be-treeshaken">
+        this should not be visible
+      </div>
+      <ClientImport />
+      <Treeshaken />
+      <ResolvedImport />
+    </ClientOnly>
+  </div>
+</template>
+
+<script setup>
+import { Treeshaken } from 'somepath'
+import HelloWorld from '../HelloWorld.vue'
+import { Glob, ClientImport } from '#components'
+
+const hello = 'world'
+</script>
+
+<style scoped>
+    .not-client {
+        color: "red";
+    }
+</style>

packages/nuxt/test/scan-components.test.ts

@@ -147,6 +147,18 @@ const expectedComponents = [
     prefetch: false,
     preload: false
   },
+  {
+    chunkName: 'components/client-with-client-only-setup',
+    export: 'default',
+    global: undefined,
+    island: undefined,
+    kebabName: 'client-with-client-only-setup',
+    mode: 'all',
+    pascalName: 'ClientWithClientOnlySetup',
+    prefetch: false,
+    preload: false,
+    shortPath: 'components/client/WithClientOnlySetup.vue'
+  },
   {
     mode: 'server',
     pascalName: 'ParentFolder',