Fix leading style and script hoist (#154)

* fix: properly handle top-level Component nodes with leading style

* chore: add changeset

* chore: fix lint error

.changeset/heavy-masks-jump.md

@@ -0,0 +1,5 @@
+---
+'@astrojs/compiler': patch
+---
+
+Fix handling of top-level component nodes with leading styles

internal/parser.go

@@ -825,9 +825,22 @@ func inHeadIM(p *parser) bool {
 			p.templateStack = append(p.templateStack, inTemplateIM)
 			return true
 		}
+
+		if p.oe.top() != nil && (isComponent(p.oe.top().Data) || isFragment((p.oe.top().Data))) {
+			p.addElement()
+			p.setOriginalIM()
+			p.im = inBodyIM
+			if p.hasSelfClosingToken {
+				p.addLoc()
+				p.oe.pop()
+				p.acknowledgeSelfClosingTag()
+			}
+			return true
+		}
 	case EndTagToken:
 		if isComponent(p.tok.Data) || isFragment(p.tok.Data) {
-			p.addElement()
+			p.addLoc()
+			p.oe.pop()
 			return true
 		}
 

internal/transform/transform.go

@@ -35,6 +35,27 @@ func Transform(doc *tycho.Node, opts TransformOptions) *tycho.Node {
 		script.Parent.RemoveChild(script)
 	}
 
+	// Sometimes files have leading <script hoist> or <style>...
+	// Since we can't detect a "component-only" file until after `parse`, we need to handle
+	// them here. The component will be hoisted to the root of the document, `html` and `head` will be removed.
+	if opts.As != "Fragment" {
+		hasBody := false
+		var onlyComponent *tycho.Node
+		walk(doc, func(n *tycho.Node) {
+			if n.Component && n.Parent != nil && n.Parent.DataAtom == a.Head {
+				onlyComponent = n
+			} else if !hasBody && n.DataAtom == a.Body {
+				hasBody = true
+			}
+		})
+
+		if !hasBody && onlyComponent != nil {
+			onlyComponent.Parent.RemoveChild(onlyComponent)
+			doc.AppendChild(onlyComponent)
+			doc.RemoveChild(onlyComponent.PrevSibling)
+		}
+	}
+
 	// If we've emptied out all the nodes, this was a Fragment that only contained hoisted elements
 	// Add an empty FrontmatterNode to allow the empty component to be printed
 	if doc.FirstChild == nil {

internal/transform/transform_test.go

@@ -109,3 +109,47 @@ func TestTransformScoping(t *testing.T) {
 		})
 	}
 }
+
+func TestFullTransform(t *testing.T) {
+	tests := []struct {
+		name   string
+		source string
+		want   string
+	}{
+		{
+			name: "top-level component with leading style",
+			source: `---
+import Component from "test";
+---
+<style>:root{}</style><Component><h1>Hello world</h1></Component>
+			`,
+			want: `<Component class="astro-XXXXXX"><h1 class="astro-XXXXXX">Hello world</h1></Component>`,
+		},
+		{
+			name: "top-level component with trailing style",
+			source: `---
+import Component from "test";
+---
+<Component><h1>Hello world</h1></Component><style>:root{}</style>
+			`,
+			want: `<Component class="astro-XXXXXX"><h1 class="astro-XXXXXX">Hello world</h1></Component>`,
+		},
+	}
+	var b strings.Builder
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			b.Reset()
+			doc, err := astro.Parse(strings.NewReader(tt.source))
+			if err != nil {
+				t.Error(err)
+			}
+			ExtractStyles(doc)
+			Transform(doc, TransformOptions{Scope: "XXXXXX"})
+			astro.PrintToSource(&b, doc)
+			got := strings.TrimSpace(b.String())
+			if tt.want != got {
+				t.Error(fmt.Sprintf("\nFAIL: %s\n  want: %s\n  got:  %s", tt.name, tt.want, got))
+			}
+		})
+	}
+}