From c44660252914f2d2a0721cfa68c345f0a435d55b Mon Sep 17 00:00:00 2001
From: Daniel Roe <daniel@roe.dev>
Date: Sun, 28 Jan 2024 21:25:42 +0000
Subject: [PATCH] fix(vite): extract styles for shared chunks (#25455)

---
 packages/vite/src/plugins/ssr-styles.ts       | 40 +++++++++++--------
 packages/vite/src/utils/index.ts              |  2 +
 test/basic.test.ts                            | 15 +++++--
 .../components/ServerOnlyComponent.server.vue |  1 +
 .../components/ServerOnlyComponentChild.vue   | 11 +++++
 .../basic/components/SharedComponent.vue      |  9 +++++
 test/fixtures/basic/pages/styles.vue          |  1 +
 test/fixtures/basic/pages/vueuse-head.vue     |  2 +
 8 files changed, 60 insertions(+), 21 deletions(-)
 create mode 100644 test/fixtures/basic/components/ServerOnlyComponentChild.vue
 create mode 100644 test/fixtures/basic/components/SharedComponent.vue

diff --git a/packages/vite/src/plugins/ssr-styles.ts b/packages/vite/src/plugins/ssr-styles.ts
index cb779f689b..46a6e5345f 100644
--- a/packages/vite/src/plugins/ssr-styles.ts
+++ b/packages/vite/src/plugins/ssr-styles.ts
@@ -8,7 +8,7 @@ import type { Component } from '@nuxt/schema'
 import MagicString from 'magic-string'
 import { findStaticImports } from 'mlly'
 
-import { isCSS } from '../utils'
+import { isCSS, isVue } from '../utils'
 
 interface SSRStylePluginOptions {
   srcDir: string
@@ -107,25 +107,31 @@ export function ssrStylesPlugin (options: SSRStylePluginOptions): Plugin {
       })
     },
     renderChunk (_code, chunk) {
-      if (!chunk.facadeModuleId) { return null }
-
-      // 'Teleport' CSS chunks that made it into the bundle on the client side
-      // to be inlined on server rendering
-      if (options.mode === 'client') {
-        options.clientCSSMap[chunk.facadeModuleId] ||= new Set()
-        for (const id of chunk.moduleIds) {
-          if (isCSS(id)) {
-            options.clientCSSMap[chunk.facadeModuleId].add(id)
-          }
-        }
-        return
+      const isEntry = chunk.facadeModuleId === options.entry
+      if (isEntry) {
+        options.clientCSSMap[chunk.facadeModuleId!] ||= new Set()
       }
+      for (const moduleId of [chunk.facadeModuleId, ...chunk.moduleIds].filter(Boolean) as string[]) {
+        // 'Teleport' CSS chunks that made it into the bundle on the client side
+        // to be inlined on server rendering
+        if (options.mode === 'client') {
+          options.clientCSSMap[moduleId] ||= new Set()
+          if (isCSS(moduleId)) {
+            // Vue files can (also) be their own entrypoints as they are tracked separately
+            if (isVue(moduleId)) {
+              options.clientCSSMap[moduleId].add(moduleId)
+            }
+            // This is required to track CSS in entry chunk
+            if (isEntry) {
+              options.clientCSSMap[chunk.facadeModuleId!].add(moduleId)
+            }
+          }
+          continue
+        }
 
-      const id = relativeToSrcDir(chunk.facadeModuleId)
-      for (const file in chunk.modules) {
-        const relativePath = relativeToSrcDir(file)
+        const relativePath = relativeToSrcDir(moduleId)
         if (relativePath in cssMap) {
-          cssMap[relativePath].inBundle = cssMap[relativePath].inBundle ?? !!id
+          cssMap[relativePath].inBundle = cssMap[relativePath].inBundle ?? ((isVue(moduleId) && relativeToSrcDir(moduleId)) || isEntry)
         }
       }
 
diff --git a/packages/vite/src/utils/index.ts b/packages/vite/src/utils/index.ts
index 7ef4bdf192..ef1262cffb 100644
--- a/packages/vite/src/utils/index.ts
+++ b/packages/vite/src/utils/index.ts
@@ -1,5 +1,7 @@
 import { hash } from 'ohash'
 
+export { isVue } from '../../../nuxt/src/core/utils/plugins'
+
 export function uniq<T> (arr: T[]): T[] {
   return Array.from(new Set(arr))
 }
diff --git a/test/basic.test.ts b/test/basic.test.ts
index 44f7045564..fc6dd4d9f5 100644
--- a/test/basic.test.ts
+++ b/test/basic.test.ts
@@ -112,7 +112,7 @@ describe('pages', () => {
     // should apply attributes to client-only components
     expect(html).toContain('<div style="color:red;" class="client-only"></div>')
     // should render server-only components
-    expect(html.replace(/ data-island-uid="[^"]*"/, '')).toContain('<div class="server-only" style="background-color:gray;"> server-only component </div>')
+    expect(html.replace(/ data-island-uid="[^"]*"/, '')).toContain('<div class="server-only" style="background-color:gray;"> server-only component <div> server-only component child (non-server-only) </div></div>')
     // should register global components automatically
     expect(html).toContain('global component registered automatically')
     expect(html).toContain('global component via suffix')
@@ -1382,6 +1382,8 @@ describe.skipIf(isDev() || isWebpack)('inlining component styles', () => {
     '{--assets:"assets"}', // <script>
     '{--postcss:"postcss"}', // <style lang=postcss>
     '{--scoped:"scoped"}', // <style lang=css>
+    '{--shared-component:"shared-component"}', // styles in a chunk shared between pages
+    '{--server-only-child:"server-only-child"}', // child of a server-only component
     '{--server-only:"server-only"}' // server-only component not in client build
     // TODO: ideally both client/server components would have inlined css when used
     // '{--client-only:"client-only"}', // client-only component not in server build
@@ -1392,7 +1394,7 @@ describe.skipIf(isDev() || isWebpack)('inlining component styles', () => {
   it('should inline styles', async () => {
     const html = await $fetch('/styles')
     for (const style of inlinedCSS) {
-      expect(html).toContain(style)
+      expect.soft(html).toContain(style)
     }
   })
 
@@ -1403,7 +1405,7 @@ describe.skipIf(isDev() || isWebpack)('inlining component styles', () => {
     ]
     const html = await $fetch('/route-rules/spa')
     for (const style of globalCSS) {
-      expect(html).toContain(style)
+      expect.soft(html).toContain(style)
     }
   })
 
@@ -1414,7 +1416,7 @@ describe.skipIf(isDev() || isWebpack)('inlining component styles', () => {
     expect(files.map(m => m.replace(/\.\w+(\.\w+)$/, '$1'))).toContain('css-only-asset.svg')
   })
 
-  it('should not include inlined CSS in generated CSS file', async () => {
+  it('should not include inlined CSS in generated CSS file', async ()  => {
     const html: string = await $fetch('/styles')
     const cssFiles = new Set([...html.matchAll(/<link [^>]*href="([^"]*\.css)">/g)].map(m => m[1]))
     let css = ''
@@ -1950,6 +1952,11 @@ describe('component islands', () => {
       expect(result.head).toMatchInlineSnapshot(`
         {
           "link": [
+            {
+              "href": "/_nuxt/components/SharedComponent.vue?vue&type=style&index=0&scoped=3ee84738&lang.css",
+              "key": "island-link",
+              "rel": "stylesheet",
+            },
             {
               "href": "/_nuxt/components/islands/PureComponent.vue?vue&type=style&index=0&scoped=c0c0cf89&lang.css",
               "key": "island-link",
diff --git a/test/fixtures/basic/components/ServerOnlyComponent.server.vue b/test/fixtures/basic/components/ServerOnlyComponent.server.vue
index a775c99524..ff3b40a1ba 100644
--- a/test/fixtures/basic/components/ServerOnlyComponent.server.vue
+++ b/test/fixtures/basic/components/ServerOnlyComponent.server.vue
@@ -5,6 +5,7 @@ prerenderRoutes(['/some/url/from/server-only/component'])
 <template>
   <div>
     server-only component
+    <ServerOnlyComponentChild />
   </div>
 </template>
 
diff --git a/test/fixtures/basic/components/ServerOnlyComponentChild.vue b/test/fixtures/basic/components/ServerOnlyComponentChild.vue
new file mode 100644
index 0000000000..5283fcd948
--- /dev/null
+++ b/test/fixtures/basic/components/ServerOnlyComponentChild.vue
@@ -0,0 +1,11 @@
+<template>
+  <div>
+    server-only component child (non-server-only)
+  </div>
+</template>
+
+<style>
+:root {
+  --server-only-child: 'server-only-child';
+}
+</style>
diff --git a/test/fixtures/basic/components/SharedComponent.vue b/test/fixtures/basic/components/SharedComponent.vue
new file mode 100644
index 0000000000..ebf2e90c3f
--- /dev/null
+++ b/test/fixtures/basic/components/SharedComponent.vue
@@ -0,0 +1,9 @@
+<template>
+  <span class="shared-component" />
+</template>
+
+<style scoped>
+.shared-component {
+  --shared-component: 'shared-component';
+}
+</style>
diff --git a/test/fixtures/basic/pages/styles.vue b/test/fixtures/basic/pages/styles.vue
index 5fc55fafc9..2aadb5da0e 100644
--- a/test/fixtures/basic/pages/styles.vue
+++ b/test/fixtures/basic/pages/styles.vue
@@ -3,6 +3,7 @@
     <ClientOnlyScript />
     <FunctionalComponent />
     <ServerOnlyComponent />
+    <SharedComponent />
   </div>
 </template>
 
diff --git a/test/fixtures/basic/pages/vueuse-head.vue b/test/fixtures/basic/pages/vueuse-head.vue
index ebadc6ec98..232bc37806 100644
--- a/test/fixtures/basic/pages/vueuse-head.vue
+++ b/test/fixtures/basic/pages/vueuse-head.vue
@@ -21,5 +21,7 @@ useLegacyVueUseHead()
 <template>
   <div>
     <h1>VueUse head polyfill test</h1>
+    <!-- This component is only here to make it a shared chunk for test in `styles.vue` -->
+    <SharedComponent />
   </div>
 </template>