Fix formatting of trailing comments in composites

Previously trailing comments inside arrays, objects, and sets were not
being formatted correctly. For example:

[
    1,
    2,
    # foo
]

Would result in:

[
    1,
    2,
 # foo ]

The problem was that when the sequence was ended, the comments were not
being emitted. As a result when the comments were finally emitted, the
indenting was wrong and the state of the formatter was not consistent
(and so the closing bracket appeared on the same line the comment.)

These changes modify the formatter to emit the comments when ending the
sequence, as that's the point where the indenting state is known.

Also, as part of these changes, the fix for extra newlines (#1032) has
been modified. Instead of changing the startLine and endLine behaviours
(which are a bit sensitive) we just squash trailing newlines at the end
of the formatting process.

Fixes #1060

Signed-off-by: Torin Sandall <torinsandall@gmail.com>

format/format.go

@@ -94,7 +94,15 @@ func Ast(x interface{}) (formatted []byte, err error) {
 	default:
 		return nil, fmt.Errorf("not an ast element: %v", x)
 	}
-	return w.buf.Bytes(), nil
+
+	return squashTrailingNewlines(w.buf.Bytes()), nil
+}
+
+func squashTrailingNewlines(bs []byte) []byte {
+	if bytes.HasSuffix(bs, []byte("\n")) {
+		return append(bytes.TrimRight(bs, "\n"), '\n')
+	}
+	return bs
 }
 
 func defaultLocation(x ast.Node) *ast.Location {
@@ -249,7 +257,7 @@ func (w *writer) writeHead(head *ast.Head, isExpandedConst bool, comments []*ast
 		for _, arg := range head.Args {
 			args = append(args, arg)
 		}
-		comments = w.writeIterable(args, head.Location, comments, w.listWriter())
+		comments = w.writeIterable(args, head.Location, closingLoc(0, 0, '(', ')', head.Location), comments, w.listWriter())
 		w.write(")")
 	}
 	if head.Key != nil {
@@ -447,8 +455,7 @@ func (w *writer) writeObject(obj ast.Object, loc *ast.Location, comments []*ast.
 	obj.Foreach(func(k, v *ast.Term) {
 		s = append(s, ast.Item(k, v))
 	})
-	comments = w.writeIterable(s, loc, comments, w.objectWriter())
-	return w.insertComments(comments, closingLoc(0, 0, '{', '}', loc))
+	return w.writeIterable(s, loc, closingLoc(0, 0, '{', '}', loc), comments, w.objectWriter())
 }
 
 func (w *writer) writeArray(arr ast.Array, loc *ast.Location, comments []*ast.Comment) []*ast.Comment {
@@ -459,8 +466,7 @@ func (w *writer) writeArray(arr ast.Array, loc *ast.Location, comments []*ast.Co
 	for _, t := range arr {
 		s = append(s, t)
 	}
-	comments = w.writeIterable(s, loc, comments, w.listWriter())
-	return w.insertComments(comments, closingLoc(0, 0, '[', ']', loc))
+	return w.writeIterable(s, loc, closingLoc(0, 0, '[', ']', loc), comments, w.listWriter())
 }
 
 func (w *writer) writeSet(set ast.Set, loc *ast.Location, comments []*ast.Comment) []*ast.Comment {
@@ -471,8 +477,7 @@ func (w *writer) writeSet(set ast.Set, loc *ast.Location, comments []*ast.Commen
 	set.Foreach(func(t *ast.Term) {
 		s = append(s, t)
 	})
-	comments = w.writeIterable(s, loc, comments, w.listWriter())
-	return w.insertComments(comments, closingLoc(0, 0, '{', '}', loc))
+	return w.writeIterable(s, loc, closingLoc(0, 0, '{', '}', loc), comments, w.listWriter())
 }
 
 func (w *writer) writeArrayComprehension(arr *ast.ArrayComprehension, loc *ast.Location, comments []*ast.Comment) []*ast.Comment {
@@ -572,12 +577,11 @@ func (w *writer) writeImports(imports []*ast.Import, comments []*ast.Comment) []
 
 type entryWriter func(interface{}, []*ast.Comment) []*ast.Comment
 
-func (w *writer) writeIterable(elements []interface{}, last *ast.Location, comments []*ast.Comment, fn entryWriter) []*ast.Comment {
+func (w *writer) writeIterable(elements []interface{}, last *ast.Location, close *ast.Location, comments []*ast.Comment, fn entryWriter) []*ast.Comment {
 	lines := groupIterable(elements, last)
 	if len(lines) > 1 {
 		w.delayBeforeEnd()
 		w.startMultilineSeq()
-		defer w.endMultilineSeq()
 	}
 
 	i := 0
@@ -588,7 +592,17 @@ func (w *writer) writeIterable(elements []interface{}, last *ast.Location, comme
 		w.endLine()
 		w.startLine()
 	}
+
 	comments = w.writeIterableLine(lines[i], comments, fn)
+
+	if len(lines) > 1 {
+		w.write(",")
+		w.endLine()
+		comments = w.insertComments(comments, close)
+		w.down()
+		w.startLine()
+	}
+
 	return comments
 }
 
@@ -840,9 +854,6 @@ func (w *writer) startLine() {
 		panic("currently in a line")
 	}
 	w.inline = true
-	if w.buf.Len() > 0 {
-		w.write("\n")
-	}
 	for i := 0; i < w.level; i++ {
 		w.write(w.indent)
 	}
@@ -859,6 +870,7 @@ func (w *writer) endLine() {
 		w.beforeEnd = nil
 	}
 	w.delay = false
+	w.write("\n")
 }
 
 // beforeLineEnd registers a comment to be printed at the end of the current line.

format/format_test.go

@@ -122,7 +122,7 @@ func differsAt(a, b []byte) (int, int) {
 			return ln, i
 		}
 	}
-	return ln, minLen
+	return ln, minLen - 1
 }
 
 func prefixWithLineNumbers(bs []byte) []byte {

format/testfiles/test.rego

@@ -109,6 +109,7 @@ p[x] = y { y = x
                     "b":    "c"   ,   "c": "d",   # comment on object entry line
                     # Comment inside object 2
 "d": "e",
+# Comment before closing object brace.
 } # Comment on closing object brace.
 a = {"a": "b", "c": "d"}
 b = [1, 2, 3, 4]
@@ -117,6 +118,13 @@ c = [1, 2,
 3, 4,
 5, 6, 7,
 8,
+# Comment before nested composite.
+[
+["foo"], # Comment inside nested composite.
+["bar"], # Comment after last element in nested composite.
+# Comment before nested composite closing bracket.
+], # Comment on nested composite closing bracket.
+# Comment before closing array bracket.
 ] # Comment on closing array bracket.
 
 d = [1 | b[_]]

format/testfiles/test.rego.formatted

@@ -122,6 +122,7 @@ p[x] = y {
 		"b": "c", "c": "d", # comment on object entry line
 		# Comment inside object 2
 		"d": "e",
+		# Comment before closing object brace.
 	} # Comment on closing object brace.
 
 	a = {"a": "b", "c": "d"}
@@ -132,6 +133,13 @@ p[x] = y {
 		3, 4,
 		5, 6, 7,
 		8,
+		# Comment before nested composite.
+		[
+			["foo"], # Comment inside nested composite.
+			["bar"], # Comment after last element in nested composite.
+			# Comment before nested composite closing bracket.
+		], # Comment on nested composite closing bracket.
+		# Comment before closing array bracket.
 	] # Comment on closing array bracket.
 
 	d = [1 | b[_]]