Fix Line Comments Following Statement

octoberswimmer · Jan 8, 2025 · 0f3dd8a · 0f3dd8a
1 parent a1abbc5
commit 0f3dd8a
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 3 deletions.
diff --git a/formatter/comments_test.go b/formatter/comments_test.go
@@ -84,6 +84,30 @@ System.debug('I am on a separate line!');`,
 				`// Test trailing whitespace
 go();`,
 			},
+			{
+				`if (true) {
+	if (true) {
+		go();
+	} // line comment
+}`,
+				`if (true) {
+	if (true) {
+		go();
+	} // line comment
+}`,
+			},
+			{
+				`if (true) {
+	if (true) {
+		go(); 			// line comment
+	}
+}`,
+				`if (true) {
+	if (true) {
+		go(); // line comment
+	}
+}`,
+			},
 		}
 	for i, tt := range tests {
 		t.Run(fmt.Sprintf("%d", i), func(t *testing.T) {

diff --git a/formatter/visitor.go b/formatter/visitor.go
@@ -378,18 +378,20 @@ func appendHiddenTokens(v *FormatVisitor, result interface{}, tokens []antlr.Tok
 				// Strip leading whitespace so the comment can be indented to the right location
 				containsNewline := strings.Contains(text, "\n")
 				text = strings.TrimSpace(text)
-				toEndOfLine := strings.HasPrefix(text, "//")
+				lineComment := strings.HasPrefix(text, "//")
 
 				if n := countNewlines(trailingWhitespace); n > 0 {
 					trailing = strings.Repeat("\n", n)
-				} else if len(trailingWhitespace) > 0 && !toEndOfLine {
+				} else if len(trailingWhitespace) > 0 && !lineComment {
 					trailing = " "
 				}
 
 				text = fmt.Sprintf("%s%s%s", leading, text, trailing)
 				log.Trace(fmt.Sprintf("NORMALIZED COMMENT: %q\n", text))
-				if containsNewline || toEndOfLine {
+				if containsNewline {
 					text = "\uFFFA" + text + "\uFFFB" + "\n"
+				} else if lineComment {
+					text = "\uFFF9" + text + "\uFFFB" + "\n"
 				} else {
 					text = "\uFFF9" + text + "\uFFFB"
 				}