Refactor Comment Handling

Update the antlr grammar so that comments are parsed with their
surrounding whitespace, allowing us to more easily preserve the
comment formatting.

Fix lots of comment formatting edge cases.

Fix parsing of class body in which it wasn't visited properly, resulting
in comments outside the body to be moved inside the body.

formatter/comments_test.go

@@ -12,7 +12,6 @@ import (
 func TestComments(t *testing.T) {
 	if testing.Verbose() {
 		log.SetLevel(log.DebugLevel)
-
 	}
 	tests :=
 		[]struct {
@@ -115,27 +114,30 @@ func TestTrailingComments(t *testing.T) {
 			{
 				`public class MyClass { public static void noop() {}
 	// Comment Inside Compilation Unit
-	// Line 2
+	// Line 2 Inside Compilation Unit
 }`,
 				`public class MyClass {
 	public static void noop() {}
 	// Comment Inside Compilation Unit
-	// Line 2
+	// Line 2 Inside Compilation Unit
 }`},
 			{
 				`public class MyClass { public static void noop() {}}
-// Comment Outside Compilation Unit Moved Inside
-// Line 2`,
+// Comment Outside Compilation Unit Not Moved Inside
+// Line 2
+`,
 				`public class MyClass {
 	public static void noop() {}
-	// Comment Outside Compilation Unit Moved Inside
-	// Line 2
-}`},
+}
+// Comment Outside Compilation Unit Not Moved Inside
+// Line 2
+`},
 			{
 				`
 /* comment with whitespace before */
 private class T1Exception {}`,
-				`/* comment with whitespace before */
+				`
+/* comment with whitespace before */
 private class T1Exception {}`,
 			},
 			{
@@ -162,6 +164,18 @@ private class T1Exception {}`,
 		// details about statement2
 		statement2();
 	}
+}`,
+			},
+			{
+				`class TestClass {
+
+	// Blank line before comment
+	private Integer i;
+}`,
+				`class TestClass {
+
+	// Blank line before comment
+	private Integer i;
 }`,
 			},
 		}

formatter/format_test.go

@@ -469,6 +469,98 @@ func TestCompilationUnit(t *testing.T) {
 	}
 
 	void method1(String param1, String param2) {}
+}`,
+			},
+			{
+				`public class MethodRecorder {
+	static {
+		reset();
+	}
+	/**
+	 * Doc string
+	 * line 2
+	 */
+	public static void m1() {
+	}
+}`,
+				`public class MethodRecorder {
+	static {
+		reset();
+	}
+	/**
+	 * Doc string
+	 * line 2
+	 */
+	public static void m1() {}
+}`,
+			},
+			{
+				`public class MethodRecorder {
+
+	/**
+	 * Doc string after newline
+	 * line 2
+	 */
+	public static void m2() {
+	}
+}`,
+				`public class MethodRecorder {
+
+	/**
+	 * Doc string after newline
+	 * line 2
+	 */
+	public static void m2() {}
+}`,
+			},
+			{
+				`public with sharing class Assert {
+class InvalidArgumentException extends Exception {}
+
+
+	private static void checkAssertion(Assertion e) {}
+
+private class Assertion {
+public String message;
+}
+			}`,
+				`public with sharing class Assert {
+	class InvalidArgumentException extends Exception {}
+
+	private static void checkAssertion(Assertion e) {}
+
+	private class Assertion {
+		public String message;
+	}
+}`,
+			},
+			{
+				`
+public class Top {
+	public class Inner {
+		private String doot;
+
+		public void doit() {
+			System.assert(true);
+		}
+	}
+
+	/* comment
+	 */
+	public String it;
+}`,
+				`public class Top {
+	public class Inner {
+		private String doot;
+
+		public void doit() {
+			System.assert(true);
+		}
+	}
+
+	/* comment
+	 */
+	public String it;
 }`,
 			},
 		}
@@ -486,8 +578,9 @@ func TestCompilationUnit(t *testing.T) {
 		if !ok {
 			t.Errorf("Unexpected result parsing apex")
 		}
+		out = removeExtraCommentIndentation(out)
 		if out != tt.output {
-			t.Errorf("unexpected format.  expected:\n%s\ngot:\n%s\n", tt.output, out)
+			t.Errorf("unexpected format.  expected:\n%q\ngot:\n%q\n", tt.output, out)
 		}
 	}
 }

formatter/formatter.go

@@ -130,15 +130,33 @@ func removeIndentationFromComment(comment string) string {
 	modifiedComment := re.ReplaceAllString(commentBody, "\n")
 
 	// Return the modified comment, reattaching any text outside the comment block if necessary
-	return firstLine + modifiedComment
+	leadingTabPattern := regexp.MustCompile("^\t*")
+	unindented := leadingTabPattern.ReplaceAllString(firstLine, "") + modifiedComment
+	// Add leading tab
+	firstCharPattern := regexp.MustCompile(`(?s)^(\s*)(\S)`)
+	unindented = firstCharPattern.ReplaceAllString(unindented, "$1"+strings.Repeat("\t", leadingTabs)+"$2")
+
+	return unindented
 }
 
 // Comments are annotated in FormatVisitor.visitRule.  We preserve whitespace
 // within multi-line comments by removing the indentation added within the
 // comment.
 func removeExtraCommentIndentation(input string) string {
-	commentPattern := regexp.MustCompile(`(?s)\t*` + "\uFFFA" + `.*?` + "\uFFFB")
-	return commentPattern.ReplaceAllStringFunc(input, removeIndentationFromComment)
+	// Remove extra grammar-specific newlines added unaware of newline-preserving comments injected
+	newlinePrefixedMultilineComment := regexp.MustCompile("\n(\t*\uFFFA)")
+	input = newlinePrefixedMultilineComment.ReplaceAllString(input, "$1")
+
+	newlinePrefixedInlineComment := regexp.MustCompile("\n\t*\uFFF9\n")
+	input = newlinePrefixedInlineComment.ReplaceAllString(input, "\uFFF9\n")
+
+	inlineCommentPattern := regexp.MustCompile(`(?s)` + "\uFFF9" + `(.*?)` + "\uFFFB")
+	input = inlineCommentPattern.ReplaceAllString(input, "$1")
+
+	multilineCommentPattern := regexp.MustCompile(`(?s)\t*` + "\uFFFA" + `.*?` + "\uFFFB")
+	unindented := multilineCommentPattern.ReplaceAllStringFunc(input, removeIndentationFromComment)
+
+	return unindented
 }
 
 func readFile(filename string, reader io.Reader) ([]byte, error) {