From 2c33a1e3f236637b11d2f0804914d62739b3dd48 Mon Sep 17 00:00:00 2001
From: Eduardo Asafe <e@asafe.dev>
Date: Mon, 11 Mar 2024 08:47:31 -0300
Subject: [PATCH] Fix Deletion of Trailing Comments
Fix formatter bug that dropped trailing comments.
Co-authored-by: Christian G. Warden <cwarden@xerus.org>
---
formatter/comments_test.go | 70 ++++++++++++++++++++++++++++++++++
formatter/formatter.go | 2 +-
formatter/visitor.go | 78 +++++++++++++++++++++++++++++++++++---
3 files changed, 144 insertions(+), 6 deletions(-)
diff --git a/formatter/comments_test.go b/formatter/comments_test.go
index 37ded9c..16482a9 100644
--- a/formatter/comments_test.go
+++ b/formatter/comments_test.go
@@ -97,3 +97,73 @@ System.debug('I am on a separate line!');`,
}
}
+
+func TestTrailingComments(t *testing.T) {
+ if testing.Verbose() {
+ log.SetLevel(log.DebugLevel)
+
+ }
+ tests :=
+ []struct {
+ input string
+ output string
+ }{
+ {
+ `private class T1Exception extends Exception {} //test`,
+ `private class T1Exception extends Exception {} //test`,
+ },
+ {
+ `public class MyClass { public static void noop() {}
+ // Comment Inside Compilation Unit
+ // Line 2
+}`,
+ `public class MyClass {
+ public static void noop() {}
+ // Comment Inside Compilation Unit
+ // Line 2
+}`},
+ {
+ `public class MyClass { public static void noop() {}}
+// Comment Outside Compilation Unit Moved Inside
+// Line 2`,
+ `public class MyClass {
+ public static void noop() {}
+ // Comment Outside Compilation Unit Moved Inside
+ // Line 2
+}`},
+ {
+ `
+/* comment with whitespace before */
+private class T1Exception {}`,
+ `/* comment with whitespace before */
+private class T1Exception {}`,
+ },
+ {
+ `/* comment with whitespace after */
+
+private class T1Exception {}`,
+ `/* comment with whitespace after */
+
+private class T1Exception {}`,
+ },
+ }
+ for _, tt := range tests {
+ input := antlr.NewInputStream(tt.input)
+ lexer := parser.NewApexLexer(input)
+ stream := antlr.NewCommonTokenStream(lexer, antlr.TokenDefaultChannel)
+
+ p := parser.NewApexParser(stream)
+ p.RemoveErrorListeners()
+ p.AddErrorListener(&testErrorListener{t: t})
+
+ v := NewFormatVisitor(stream)
+ out, ok := v.visitRule(p.CompilationUnit()).(string)
+ if !ok {
+ t.Errorf("Unexpected result parsing apex")
+ }
+ out = removeExtraCommentIndentation(out)
+ if out != tt.output {
+ t.Errorf("unexpected format. expected:\n%q\ngot:\n%q\n", tt.output, out)
+ }
+ }
+}
diff --git a/formatter/formatter.go b/formatter/formatter.go
index 344b994..9afa33c 100644
--- a/formatter/formatter.go
+++ b/formatter/formatter.go
@@ -76,7 +76,7 @@ func (f *Formatter) Format() error {
if f.source == nil {
src, err := readFile(f.filename, f.reader)
if err != nil {
- return fmt.Errorf("Failed to read file %s: %w", f.SourceName(), err)
+ return fmt.Errorf("failed to read file %s: %w", f.SourceName(), err)
}
f.source = src
}
diff --git a/formatter/visitor.go b/formatter/visitor.go
index 88ce56c..81ba464 100644
--- a/formatter/visitor.go
+++ b/formatter/visitor.go
@@ -46,6 +46,7 @@ func (v *FormatVisitor) visitRule(node antlr.RuleNode) interface{} {
panic(fmt.Sprintf("MISSING VISIT FUNCTION FOR %T", node))
}
commentsWithNewlines := commentsWithTrailingNewlines(beforeComments, beforeWhitespace)
+ hasComments := beforeComments != nil && len(beforeComments) > 0
if beforeComments != nil {
comments := []string{}
for _, c := range beforeComments {
@@ -54,8 +55,8 @@ func (v *FormatVisitor) visitRule(node antlr.RuleNode) interface{} {
// Mark the start and end of comments so we can remove indentation
// added to multi-line comments, preserving the whitespace within
// them. See removeIndentationFromComment.
- if _, exists := commentsWithNewlines[index]; exists {
- comments = append(comments, "\uFFFA"+c.GetText()+"\uFFFB\n")
+ if n, exists := commentsWithNewlines[index]; exists {
+ comments = append(comments, "\uFFFA"+c.GetText()+"\uFFFB"+strings.Repeat("\n", n))
} else {
comments = append(comments, "\uFFFA"+c.GetText()+"\uFFFB")
}
@@ -75,7 +76,7 @@ func (v *FormatVisitor) visitRule(node antlr.RuleNode) interface{} {
if beforeWhitespace != nil {
injectNewline := false
for _, c := range beforeWhitespace {
- if len(strings.Split(c.GetText(), "\n")) > 2 {
+ if !hasComments && len(strings.Split(c.GetText(), "\n")) > 2 {
if _, seen := v.newlinesOutput[c.GetTokenIndex()]; !seen {
v.newlinesOutput[c.GetTokenIndex()] = struct{}{}
injectNewline = true
@@ -86,6 +87,46 @@ func (v *FormatVisitor) visitRule(node antlr.RuleNode) interface{} {
result = fmt.Sprintf("\n%s", result)
}
}
+ stop := node.(antlr.ParserRuleContext).GetStop()
+ var afterWhitespace, afterComments []antlr.Token
+ if stop == nil {
+ return result
+ }
+ if len(v.tokens.GetAllTokens()) > 0 {
+ afterWhitespace = v.tokens.GetHiddenTokensToRight(stop.GetTokenIndex(), WHITESPACE_CHANNEL)
+ afterComments = v.tokens.GetHiddenTokensToRight(stop.GetTokenIndex(), COMMENTS_CHANNEL)
+ }
+
+ if afterComments == nil {
+ return result
+ }
+ afterCommentsWithLeadingNewlines := commentsWithLeadingNewlines(afterComments, afterWhitespace)
+ comments := []string{}
+
+ for _, c := range afterComments {
+ index := c.GetTokenIndex()
+ if _, seen := v.commentsOutput[index]; !seen {
+ // Mark the start and end of comments so we can remove indentation
+ // added to multi-line comments, preserving the whitespace within
+ // them. See removeIndentationFromComment.
+ leading := ""
+ if _, exists := afterCommentsWithLeadingNewlines[index]; exists {
+ leading = "\n"
+ }
+ comments = append(comments, leading+"\uFFFA"+c.GetText()+"\uFFFB")
+ v.commentsOutput[index] = struct{}{}
+ }
+ }
+
+ if len(comments) > 0 {
+ allComments := strings.Join(comments, "")
+ containsNewline := strings.Contains(allComments, "\n")
+ if !containsNewline {
+ result = fmt.Sprintf("%s %s", result, strings.TrimSuffix(strings.TrimPrefix(allComments, "\uFFFA"), "\uFFFB"))
+ } else {
+ result = fmt.Sprintf("%s%s", result, allComments)
+ }
+ }
return result
}
@@ -154,8 +195,8 @@ func unwrap(v *FormatVisitor) (*FormatVisitor, bool) {
}
// Find comments that have trailing newlines
-func commentsWithTrailingNewlines(comments []antlr.Token, whitespace []antlr.Token) map[int]struct{} {
- result := make(map[int]struct{})
+func commentsWithTrailingNewlines(comments []antlr.Token, whitespace []antlr.Token) map[int]int {
+ result := make(map[int]int)
whitespaceMap := make(map[int]antlr.Token)
for _, ws := range whitespace {
@@ -170,6 +211,33 @@ func commentsWithTrailingNewlines(comments []antlr.Token, whitespace []antlr.Tok
// Check if the next token is whitespace
if ws, exists := whitespaceMap[nextTokenIndex]; exists {
+ // Check if the whitespace contains a newline
+ if strings.Contains(ws.GetText(), "\n") {
+ result[commentIndex] = len(strings.Split(ws.GetText(), "\n")) - 1
+ }
+ }
+ }
+
+ return result
+}
+
+// Find comments that have leading newlines
+func commentsWithLeadingNewlines(comments []antlr.Token, whitespace []antlr.Token) map[int]struct{} {
+ result := make(map[int]struct{})
+
+ whitespaceMap := make(map[int]antlr.Token)
+ for _, ws := range whitespace {
+ whitespaceMap[ws.GetTokenIndex()] = ws
+ }
+
+ for _, comment := range comments {
+ commentIndex := comment.GetTokenIndex()
+
+ // Find the immediate previous token index
+ prevTokenIndex := commentIndex - 1
+
+ // Check if the next token is whitespace
+ if ws, exists := whitespaceMap[prevTokenIndex]; exists {
// Check if the whitespace contains a newline
if strings.Contains(ws.GetText(), "\n") {
result[commentIndex] = struct{}{}