New handling for `define:vars` scripts and styles (#3976)

* feat: new handling for `define:vars` scripts and styles

* fix: handle new script hoisting pattern

* refactor: compiler handles sourcemaps

* chore: update to handle is:inline define:vars

* chore: bump compiler to latest

* chore: update define:vars tests

* fix: output of `define:vars` is not object style

* fix: appease ts

* chore: remove unused file

* chore: revert unecessary refactors

* chore: prefer sync `defineScriptVars`

* chore: add changeset

Co-authored-by: Nate Moore <nate@astro.build>
Co-authored-by: Okiki Ojo <okikio.dev@gmail.com>

8 files changed, 93 insertions, 66 deletions

diff --git a/.changeset/sixty-drinks-search.md b/.changeset/sixty-drinks-search.md
new file mode 100644
index 000000000..3565b2fcf
--- /dev/null
+++ b/.changeset/sixty-drinks-search.md
@@ -0,0 +1,5 @@
+---
+'astro': patch
+---
+
+Fix `define:vars` bugs with both `style` and `script`
diff --git a/packages/astro/package.json b/packages/astro/package.json
index ee82786b3..8d5a27e49 100644
--- a/packages/astro/package.json
+++ b/packages/astro/package.json
@@ -82,7 +82,7 @@
     "test:e2e:match": "playwright test -g"
   },
   "dependencies": {
-    "@astrojs/compiler": "^0.20.0",
+    "@astrojs/compiler": "^0.21.0",
     "@astrojs/language-server": "^0.20.0",
     "@astrojs/markdown-remark": "^0.12.0",
     "@astrojs/prism": "0.6.1",
diff --git a/packages/astro/src/runtime/server/index.ts b/packages/astro/src/runtime/server/index.ts
index 59473798d..29ac1ef08 100644
--- a/packages/astro/src/runtime/server/index.ts
+++ b/packages/astro/src/runtime/server/index.ts
@@ -213,20 +213,6 @@ export async function renderComponent(
 	if (Component && (Component as any).isAstroComponentFactory) {
 		async function* renderAstroComponentInline(): AsyncGenerator<string, void, undefined> {
 			let iterable = await renderToIterable(result, Component as any, _props, slots);
-			// If this component added any define:vars styles and the head has already been
-			// sent out, we need to include those inline.
-			if (result.styles.size && alreadyHeadRenderedResults.has(result)) {
-				let styles = Array.from(result.styles);
-				result.styles.clear();
-				for (const style of styles) {
-					if ('define:vars' in style.props) {
-						// We only want to render the property value and not the full stylesheet
-						// which is bundled in the head.
-						style.children = '';
-						yield markHTMLString(renderElement('style', style));
-					}
-				}
-			}
 			yield* iterable;
 		}
 
@@ -583,7 +569,7 @@ Make sure to use the static attribute syntax (\`${key}={value}\`) instead of the
 	}
 
 	// support object styles for better JSX compat
-	if (key === 'style' && typeof value === 'object') {
+	if (key === 'style' && !(value instanceof HTMLString) && typeof value === 'object') {
 		return markHTMLString(` ${key}="${toStyleString(value)}"`);
 	}
 
@@ -633,19 +619,31 @@ export function spreadAttributes(
 }
 
 // Adds CSS variables to an inline style tag
-export function defineStyleVars(selector: string, vars: Record<any, any>) {
-	let output = '\n';
-	for (const [key, value] of Object.entries(vars)) {
-		output += `  --${key}: ${value};\n`;
+export function defineStyleVars(defs: Record<any, any>|Record<any, any>[]) {
+	let output = '';
+	let arr = !Array.isArray(defs) ? [defs] : defs;
+	for (const vars of arr) {
+		for (const [key, value] of Object.entries(vars)) {
+			if (value || value === 0) {
+				output += `--${key}: ${value};`;
+			}
+		}
 	}
-	return markHTMLString(`${selector} {${output}}`);
+	return markHTMLString(output);
 }
 
+// converts (most) arbitrary strings to valid JS identifiers
+const toIdent = (k: string) =>
+  k.trim().replace(/(?:(?<!^)\b\w|\s+|[^\w]+)/g, (match, index) => {
+    if (/[^\w]|\s/.test(match)) return '';
+    return index === 0 ? match : match.toUpperCase();
+  });
+
 // Adds variables to an inline script.
 export function defineScriptVars(vars: Record<any, any>) {
 	let output = '';
 	for (const [key, value] of Object.entries(vars)) {
-		output += `let ${key} = ${JSON.stringify(value)};\n`;
+		output += `let ${toIdent(key)} = ${JSON.stringify(value)};\n`;
 	}
 	return markHTMLString(output);
 }
@@ -945,11 +943,6 @@ function renderElement(
 	const { lang: _, 'data-astro-id': astroId, 'define:vars': defineVars, ...props } = _props;
 	if (defineVars) {
 		if (name === 'style') {
-			if (props['is:global']) {
-				children = defineStyleVars(`:root`, defineVars) + '\n' + children;
-			} else {
-				children = defineStyleVars(`.astro-${astroId}`, defineVars) + '\n' + children;
-			}
 			delete props['is:global'];
 			delete props['is:scoped'];
 		}
diff --git a/packages/astro/src/vite-plugin-astro/index.ts b/packages/astro/src/vite-plugin-astro/index.ts
index 9dfbdddac..333621f6d 100644
--- a/packages/astro/src/vite-plugin-astro/index.ts
+++ b/packages/astro/src/vite-plugin-astro/index.ts
@@ -1,4 +1,4 @@
-import type { PluginContext } from 'rollup';
+import type { PluginContext, SourceDescription } from 'rollup';
 import type * as vite from 'vite';
 import type { AstroConfig } from '../@types/astro';
 import type { LogOptions } from '../core/logger/core.js';
@@ -175,17 +175,29 @@ export default function astro({ config, logging }: AstroPluginOptions): vite.Plu
 						}
 					}
 
-					return {
-						code:
-							hoistedScript.type === 'inline'
-								? hoistedScript.code!
-								: `import "${hoistedScript.src!}";`,
+					let result: SourceDescription & { meta: any } = {
+						code: '',
 						meta: {
 							vite: {
-								lang: 'ts',
-							},
-						},
+								lang: 'ts'
+							}
+						}
 					};
+
+					switch (hoistedScript.type) {
+						case 'inline': {
+							let { code, map } = hoistedScript
+							result.code = appendSourceMap(code, map);
+							break;
+						}
+						case 'external': {
+							const { src } = hoistedScript
+							result.code = `import "${src}"`;
+							break;
+						}
+					}
+					
+					return result
 				}
 				default:
 					return null;
@@ -349,3 +361,8 @@ ${source}
 		},
 	};
 }
+
+function appendSourceMap(content: string, map?: string) {
+	if (!map) return content;
+	return `${content}\n//# sourceMappingURL=data:application/json;charset=utf-8;base64,${Buffer.from(map).toString('base64')}`
+}
diff --git a/packages/astro/test/astro-directives.test.js b/packages/astro/test/astro-directives.test.js
index 8bdccb90b..c2012b169 100644
--- a/packages/astro/test/astro-directives.test.js
+++ b/packages/astro/test/astro-directives.test.js
@@ -14,8 +14,22 @@ describe('Directives', async () => {
 		const html = await fixture.readFile('/define-vars/index.html');
 		const $ = cheerio.load(html);
 
-		expect($('script#inline')).to.have.lengthOf(1);
-		expect($('script#inline').toString()).to.include('let foo = "bar"');
+		expect($('script')).to.have.lengthOf(3);
+
+		let i = 0;
+		for (const script of $('script').toArray()) {
+			// Wrap script in scope ({}) to avoid redeclaration errors
+			expect($(script).text().at(0)).to.equal('{');
+			expect($(script).text().at(-1)).to.equal('}');
+			if (i < 2) {
+				// Inline defined variables
+				expect($(script).toString()).to.include('let foo = "bar"');
+			} else {
+				// Convert invalid keys to valid identifiers
+				expect($(script).toString()).to.include('let dashCase = "bar"');
+			}
+			i++;
+		}
 	});
 
 	it('Passes define:vars to style elements', async () => {
@@ -23,26 +37,13 @@ describe('Directives', async () => {
 		const $ = cheerio.load(html);
 
 		expect($('style')).to.have.lengthOf(2);
-		expect($('style').toString()).to.include('--bg: white;');
-		expect($('style').toString()).to.include('--fg: black;');
-
-		const scopedClass = $('html')
-			.attr('class')
-			.split(' ')
-			.find((name) => /^astro-[A-Za-z0-9-]+/.test(name));
-
-		expect($($('style').get(0)).toString().replace(/\s+/g, '')).to.equal(
-			`<style>.${scopedClass}{--bg:white;--fg:black;}body{background:var(--bg);color:var(--fg)}</style>`
-		);
-
-		const scopedTitleClass = $('.title')
-			.attr('class')
-			.split(' ')
-			.find((name) => /^astro-[A-Za-z0-9-]+/.test(name));
-
-		expect($($('style').get(1)).toString().replace(/\s+/g, '')).to.equal(
-			`<style>.${scopedTitleClass}{--textColor:red;}</style>`
-		);
+
+		// Inject style attribute on top-level element in page
+		expect($('html').attr('style').toString()).to.include('--bg: white;');
+		expect($('html').attr('style').toString()).to.include('--fg: black;');
+
+		// Inject style attribute on top-level elements in component
+		expect($('h1').attr('style').toString()).to.include('--textColor: red;');
 	});
 
 	it('set:html', async () => {
diff --git a/packages/astro/test/fixtures/astro-directives/src/components/Title.astro b/packages/astro/test/fixtures/astro-directives/src/components/Title.astro
index b59303ed6..113d5b1fe 100644
--- a/packages/astro/test/fixtures/astro-directives/src/components/Title.astro
+++ b/packages/astro/test/fixtures/astro-directives/src/components/Title.astro
@@ -1,9 +1,10 @@
 ---
-  const textColor = 'red'
+const textColor = 'red'
 ---
+
 <h1 class="title">hello there</h1>
 
-<style define:vars={{textColor: textColor}}>
+<style define:vars={{ textColor: textColor }}>
   .title {
     color: var(--textColor);
   }
diff --git a/packages/astro/test/fixtures/astro-directives/src/pages/define-vars.astro b/packages/astro/test/fixtures/astro-directives/src/pages/define-vars.astro
index 3df42aea1..86ec2e32d 100644
--- a/packages/astro/test/fixtures/astro-directives/src/pages/define-vars.astro
+++ b/packages/astro/test/fixtures/astro-directives/src/pages/define-vars.astro
@@ -7,9 +7,13 @@ let fg = 'black'
 
 <html>
   <head>
-    <style define:vars={{bg, fg}}>
-      body {
+		<style define:vars={{bg}}>
+      h1 {
         background: var(--bg);
+      }
+    </style>
+    <style define:vars={{fg}}>
+      h1 {
         color: var(--fg);
       }
     </style>
@@ -18,6 +22,12 @@ let fg = 'black'
     <script id="inline" define:vars={{ foo }}>
 			console.log(foo);
 		</script>
+    <script id="inline-2" define:vars={{ foo }}>
+			console.log(foo);
+		</script>
+		<script id="inline-3" define:vars={{ 'dash-case': foo }}>
+			console.log(foo);
+		</script>
 
 		<Title />
   </body>
diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml
index 633ee1b9d..e7816eb5b 100644
--- a/pnpm-lock.yaml
+++ b/pnpm-lock.yaml
@@ -438,7 +438,7 @@ importers:
 
   packages/astro:
     specifiers:
-      '@astrojs/compiler': ^0.20.0
+      '@astrojs/compiler': ^0.21.0
       '@astrojs/language-server': ^0.20.0
       '@astrojs/markdown-remark': ^0.12.0
       '@astrojs/prism': 0.6.1
@@ -523,7 +523,7 @@ importers:
       yargs-parser: ^21.0.1
       zod: ^3.17.3
     dependencies:
-      '@astrojs/compiler': 0.20.0
+      '@astrojs/compiler': 0.21.0
       '@astrojs/language-server': 0.20.0
       '@astrojs/markdown-remark': link:../markdown/remark
       '@astrojs/prism': link:../astro-prism
@@ -2892,8 +2892,8 @@ packages:
       leven: 3.1.0
     dev: true
 
-  /@astrojs/compiler/0.20.0:
-    resolution: {integrity: sha512-6tW9HYJXB8qjdf/YYyurNIGHgBSCbjZe8zN5JNI86ak05eGGKFkAIRwrWilnaRCtp/PNcoyHYFg5l+Hu6p9MXQ==}
+  /@astrojs/compiler/0.21.0:
+    resolution: {integrity: sha512-g+zkKpTsR0UCDiOAhjv0wQW0cPYd+2Hb5/z+ovIEu7K/v8z2jiQZqvhPvIsjI5ni+5rMFgjjoZWhkMCq+e4bOg==}
     dev: false
 
   /@astrojs/language-server/0.20.0: