From 4580762b04bb1f74af48f4824f9c3451430f3f85 Mon Sep 17 00:00:00 2001
From: Edmond Winston <132081506+edmondwinston@users.noreply.github.com>
Date: Tue, 26 Dec 2023 22:13:09 +0700
Subject: [PATCH] fix(fmt): stop adding extra lines between Go comments and
 templates (#351)

Co-authored-by: Adrian Hesketh <adrianhesketh@hushmail.com>
---
 .version                                   |  2 +-
 generator/generator.go                     |  9 ++++++
 generator/generator_test.go                | 22 ++++++++++++---
 generator/test-css-usage/template.templ    |  6 ----
 generator/test-css-usage/template_templ.go |  7 -----
 parser/v2/types.go                         | 32 ++++++++++++++--------
 parser/v2/types_test.go                    | 19 +++++++++++++
 7 files changed, 68 insertions(+), 29 deletions(-)

diff --git a/.version b/.version
index 4bfdbd046..994d91615 100644
--- a/.version
+++ b/.version
@@ -1 +1 @@
-0.2.503
\ No newline at end of file
+0.2.507
\ No newline at end of file
diff --git a/generator/generator.go b/generator/generator.go
index a4194771a..cb6c35e35 100644
--- a/generator/generator.go
+++ b/generator/generator.go
@@ -292,6 +292,15 @@ func (g *generator) writeGoExpression(n parser.TemplateFileGoExpression) (err er
 		return err
 	}
 	g.sourceMap.Add(n.Expression, r)
+	v := n.Expression.Value
+	lineSlice := strings.Split(v, "\n")
+	lastLine := lineSlice[len(lineSlice)-1]
+	if strings.HasPrefix(lastLine, "//") {
+		if _, err = g.w.WriteIndent(0, "\n"); err != nil {
+			return err
+		}
+		return err
+	}
 	if _, err = g.w.WriteIndent(0, "\n\n"); err != nil {
 		return err
 	}
diff --git a/generator/generator_test.go b/generator/generator_test.go
index 562a2e827..fcf139e19 100644
--- a/generator/generator_test.go
+++ b/generator/generator_test.go
@@ -14,17 +14,16 @@ func TestGeneratorSourceMap(t *testing.T) {
 		w:         NewRangeWriter(w),
 		sourceMap: parser.NewSourceMap(),
 	}
-	exp := parser.TemplateFileGoExpression{
+	invalidExp := parser.TemplateFileGoExpression{
 		Expression: parser.Expression{
 			Value: "line1\nline2",
 		},
 	}
-	err := g.writeGoExpression(exp)
-	if err != nil {
+	if err := g.writeGoExpression(invalidExp); err != nil {
 		t.Fatalf("failed to write Go expression: %v", err)
 	}
-	expected := parser.NewPosition(0, 0, 0)
 
+	expected := parser.NewPosition(0, 0, 0)
 	actual, ok := g.sourceMap.TargetPositionFromSource(0, 0)
 	if !ok {
 		t.Errorf("failed to get matching target")
@@ -32,4 +31,19 @@ func TestGeneratorSourceMap(t *testing.T) {
 	if diff := cmp.Diff(expected, actual); diff != "" {
 		t.Errorf("unexpected target:\n%v", diff)
 	}
+
+	withCommentExp := parser.TemplateFileGoExpression{
+		Expression: parser.Expression{
+			Value: `package main
+
+// A comment.
+templ h1() {
+	<h1></h1>
+}
+			`,
+		},
+	}
+	if err := g.writeGoExpression(withCommentExp); err != nil {
+		t.Fatalf("failed to write Go expression: %v", err)
+	}
 }
diff --git a/generator/test-css-usage/template.templ b/generator/test-css-usage/template.templ
index a71d0bfd0..262f27473 100644
--- a/generator/test-css-usage/template.templ
+++ b/generator/test-css-usage/template.templ
@@ -1,7 +1,6 @@
 package testcssusage
 
 // Constant class.
-
 templ StyleTagsAreSupported() {
 	<style>
 	.test {
@@ -25,7 +24,6 @@ templ CSSComponentsAreSupported() {
 
 // Both CSS components and constants are supported.
 // Only string names are really required. There is no need to use templ.Class or templ.SafeClass.
-
 templ CSSComponentsAndConstantsAreSupported() {
 	<div class={ cssComponentGreen(), "classA", templ.Class("&&&classB"), templ.SafeClass("classC"), "d e" } type="button">Both CSS components and constants are supported</div>
 	// The following is also valid, but not required - you can put the class names in directly.
@@ -33,7 +31,6 @@ templ CSSComponentsAndConstantsAreSupported() {
 }
 
 // Maps can be used to determine if a class should be added or not.
-
 templ MapsCanBeUsedToConditionallySetClasses() {
 	<div class={ map[string]bool{"a": true, "b": false, "c": true} }>Maps can be used to determine if a class should be added or not.</div>
 }
@@ -53,19 +50,16 @@ templ KVCanBeUsedToConditionallySetClasses() {
 }
 
 // Pseudo attributes can be used without any special syntax.
-
 templ PsuedoAttributesAndComplexClassNamesAreSupported() {
 	<div class={ "bg-violet-500", "hover:bg-red-600", "hover:bg-sky-700", "text-[#50d71e]", "w-[calc(100%-4rem)" }>Psuedo attributes and complex class names are supported.</div>
 }
 
 // Class names are HTML escaped.
-
 templ ClassNamesAreHTMLEscaped() {
 	<div class={ "a\" onClick=\"alert('hello')\"" }>Class names are HTML escaped.</div>
 }
 
 // Combine all tests.
-
 templ TestComponent() {
 	@StyleTagsAreSupported()
 	@CSSComponentsAreSupported()
diff --git a/generator/test-css-usage/template_templ.go b/generator/test-css-usage/template_templ.go
index ded80499d..f0ad75bd3 100644
--- a/generator/test-css-usage/template_templ.go
+++ b/generator/test-css-usage/template_templ.go
@@ -11,7 +11,6 @@ import "bytes"
 import "strings"
 
 // Constant class.
-
 func StyleTagsAreSupported() templ.Component {
 	return templ.ComponentFunc(func(ctx context.Context, templ_7745c5c3_W io.Writer) (templ_7745c5c3_Err error) {
 		templ_7745c5c3_Buffer, templ_7745c5c3_IsBuffer := templ_7745c5c3_W.(*bytes.Buffer)
@@ -120,7 +119,6 @@ func CSSComponentsAreSupported() templ.Component {
 
 // Both CSS components and constants are supported.
 // Only string names are really required. There is no need to use templ.Class or templ.SafeClass.
-
 func CSSComponentsAndConstantsAreSupported() templ.Component {
 	return templ.ComponentFunc(func(ctx context.Context, templ_7745c5c3_W io.Writer) (templ_7745c5c3_Err error) {
 		templ_7745c5c3_Buffer, templ_7745c5c3_IsBuffer := templ_7745c5c3_W.(*bytes.Buffer)
@@ -194,7 +192,6 @@ func CSSComponentsAndConstantsAreSupported() templ.Component {
 }
 
 // Maps can be used to determine if a class should be added or not.
-
 func MapsCanBeUsedToConditionallySetClasses() templ.Component {
 	return templ.ComponentFunc(func(ctx context.Context, templ_7745c5c3_W io.Writer) (templ_7745c5c3_Err error) {
 		templ_7745c5c3_Buffer, templ_7745c5c3_IsBuffer := templ_7745c5c3_W.(*bytes.Buffer)
@@ -242,7 +239,6 @@ func MapsCanBeUsedToConditionallySetClasses() templ.Component {
 }
 
 // The templ.KV function can be used to add a class if a condition is true.
-
 func d() templ.CSSClass {
 	var templ_7745c5c3_CSSBuilder strings.Builder
 	templ_7745c5c3_CSSBuilder.WriteString(`font-size:12pt;`)
@@ -310,7 +306,6 @@ func KVCanBeUsedToConditionallySetClasses() templ.Component {
 }
 
 // Pseudo attributes can be used without any special syntax.
-
 func PsuedoAttributesAndComplexClassNamesAreSupported() templ.Component {
 	return templ.ComponentFunc(func(ctx context.Context, templ_7745c5c3_W io.Writer) (templ_7745c5c3_Err error) {
 		templ_7745c5c3_Buffer, templ_7745c5c3_IsBuffer := templ_7745c5c3_W.(*bytes.Buffer)
@@ -358,7 +353,6 @@ func PsuedoAttributesAndComplexClassNamesAreSupported() templ.Component {
 }
 
 // Class names are HTML escaped.
-
 func ClassNamesAreHTMLEscaped() templ.Component {
 	return templ.ComponentFunc(func(ctx context.Context, templ_7745c5c3_W io.Writer) (templ_7745c5c3_Err error) {
 		templ_7745c5c3_Buffer, templ_7745c5c3_IsBuffer := templ_7745c5c3_W.(*bytes.Buffer)
@@ -406,7 +400,6 @@ func ClassNamesAreHTMLEscaped() templ.Component {
 }
 
 // Combine all tests.
-
 func TestComponent() templ.Component {
 	return templ.ComponentFunc(func(ctx context.Context, templ_7745c5c3_W io.Writer) (templ_7745c5c3_Err error) {
 		templ_7745c5c3_Buffer, templ_7745c5c3_IsBuffer := templ_7745c5c3_W.(*bytes.Buffer)
diff --git a/parser/v2/types.go b/parser/v2/types.go
index 2f4dd23d2..fc3dcec7e 100644
--- a/parser/v2/types.go
+++ b/parser/v2/types.go
@@ -125,27 +125,37 @@ func (tf TemplateFile) Write(w io.Writer) error {
 	if err := tf.Package.Write(w, indent); err != nil {
 		return err
 	}
-	if _, err := w.Write([]byte("\n\n")); err != nil {
+	if _, err := io.WriteString(w, "\n\n"); err != nil {
 		return err
 	}
-	count := len(tf.Nodes)
-	for i := 0; i < count; i++ {
-		var err error
-		if err = tf.Nodes[i].Write(w, indent); err != nil {
+	for i := 0; i < len(tf.Nodes); i++ {
+		if err := tf.Nodes[i].Write(w, indent); err != nil {
 			return err
 		}
-		if i == count-1 {
-			_, err = w.Write([]byte("\n"))
-		} else {
-			_, err = w.Write([]byte("\n\n"))
-		}
-		if err != nil {
+		if _, err := io.WriteString(w, getNodeWhitespace(tf.Nodes, i)); err != nil {
 			return err
 		}
 	}
 	return nil
 }
 
+func getNodeWhitespace(nodes []TemplateFileNode, i int) string {
+	if i == len(nodes)-1 {
+		return "\n"
+	}
+	if _, nextIsTemplate := nodes[i+1].(HTMLTemplate); nextIsTemplate {
+		if e, isGo := nodes[i].(TemplateFileGoExpression); isGo && endsWithComment(e.Expression.Value) {
+			return "\n"
+		}
+	}
+	return "\n\n"
+}
+
+func endsWithComment(s string) bool {
+	lineSlice := strings.Split(s, "\n")
+	return strings.HasPrefix(lineSlice[len(lineSlice)-1], "//")
+}
+
 // TemplateFileNode can be a Template, CSS, Script or Go.
 type TemplateFileNode interface {
 	IsTemplateFileNode() bool
diff --git a/parser/v2/types_test.go b/parser/v2/types_test.go
index 1d5fb90b0..e0ef4926c 100644
--- a/parser/v2/types_test.go
+++ b/parser/v2/types_test.go
@@ -776,6 +776,25 @@ templ test() {
 		Leave this alone.
 	*/
 }
+`,
+		},
+		{
+			name: "godoc comments are preserved",
+			input: ` // first line removed to make indentation clear
+package main
+
+// test the comment handling.
+templ test() {
+	Test
+}
+`,
+			expected: ` // first line removed to make indentation clear
+package main
+
+// test the comment handling.
+templ test() {
+	Test
+}
 `,
 		},
 	}