From 65e2b71b802ec0c8203cbe03bc86f087cab2a977 Mon Sep 17 00:00:00 2001
From: unknown <matthew@skypack.dev>
Date: Mon, 27 Jun 2022 09:07:48 -0400
Subject: [PATCH] Inline define:var styles rendered after the head

---
 packages/astro/src/runtime/server/index.ts    | 26 +++++++++++++++++--
 packages/astro/test/astro-directives.test.js  | 13 ++++++++--
 .../src/components/Title.astro                | 10 +++++++
 .../src/pages/define-vars.astro               |  3 +++
 4 files changed, 48 insertions(+), 4 deletions(-)
 create mode 100644 packages/astro/test/fixtures/astro-directives/src/components/Title.astro

diff --git a/packages/astro/src/runtime/server/index.ts b/packages/astro/src/runtime/server/index.ts
index e1ffd69e5..bfe095f52 100644
--- a/packages/astro/src/runtime/server/index.ts
+++ b/packages/astro/src/runtime/server/index.ts
@@ -183,7 +183,26 @@ export async function renderComponent(
 	}
 
 	if (Component && (Component as any).isAstroComponentFactory) {
-		return renderToIterable(result, Component as any, _props, slots);
+		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;
+		}
+
+		return renderAstroComponentInline();
 	}
 
 	if (!Component && !_props['client:only']) {
@@ -408,6 +427,7 @@ If you're still stuck, please open an issue on GitHub or join us at https://astr
 
 	island.children = `${html ?? ''}${template}`;
 
+	// Scripts to prepend
 	let prescriptType: PrescriptType = needsHydrationScript
 		? 'both'
 		: needsDirectiveScript
@@ -415,7 +435,7 @@ If you're still stuck, please open an issue on GitHub or join us at https://astr
 		: null;
 	let prescripts = getPrescripts(prescriptType, hydration.directive);
 
-	return markHTMLString(prescripts + renderElement('astro-island', island, false));
+	return markHTMLString(prestyles + prescripts + renderElement('astro-island', island, false));
 }
 
 /** Create the Astro.fetchContent() runtime function. */
@@ -719,6 +739,8 @@ export async function renderHead(result: SSRResult): Promise<string> {
 	const styles = Array.from(result.styles)
 		.filter(uniqueElements)
 		.map((style) => renderElement('style', style));
+	// Clear result.styles so that any new styles added will be inlined.
+	result.styles.clear();
 	const scripts = Array.from(result.scripts)
 		.filter(uniqueElements)
 		.map((script, i) => {
diff --git a/packages/astro/test/astro-directives.test.js b/packages/astro/test/astro-directives.test.js
index 161063e97..0a4bab43b 100644
--- a/packages/astro/test/astro-directives.test.js
+++ b/packages/astro/test/astro-directives.test.js
@@ -22,7 +22,7 @@ describe('Directives', async () => {
 		const html = await fixture.readFile('/define-vars/index.html');
 		const $ = cheerio.load(html);
 
-		expect($('style')).to.have.lengthOf(1);
+		expect($('style')).to.have.lengthOf(2);
 		expect($('style').toString()).to.include('--bg: white;');
 		expect($('style').toString()).to.include('--fg: black;');
 
@@ -31,9 +31,18 @@ describe('Directives', async () => {
 			.split(' ')
 			.find((name) => /^astro-[A-Za-z0-9-]+/.test(name));
 
-		expect($('style').toString().replace(/\s+/g, '')).to.equal(
+		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>`
+		);
 	});
 
 	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
new file mode 100644
index 000000000..b59303ed6
--- /dev/null
+++ b/packages/astro/test/fixtures/astro-directives/src/components/Title.astro
@@ -0,0 +1,10 @@
+---
+  const textColor = 'red'
+---
+<h1 class="title">hello there</h1>
+
+<style define:vars={{textColor: textColor}}>
+  .title {
+    color: var(--textColor);
+  }
+</style>
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 db03705ad..3df42aea1 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
@@ -1,4 +1,5 @@
 ---
+import Title from "../components/Title.astro"
 let foo = 'bar'
 let bg = 'white'
 let fg = 'black'
@@ -17,5 +18,7 @@ let fg = 'black'
     <script id="inline" define:vars={{ foo }}>
 			console.log(foo);
 		</script>
+
+		<Title />
   </body>
 </html>