fmt: fix blank lines after multiline expressions (#5194)

Signed-off-by: Jasper Van der Jeugt <m@jaspervdj.be>

format/format.go

@@ -471,10 +471,19 @@ func (w *writer) insertComments(comments []*ast.Comment, loc *ast.Location) []*a
 
 func (w *writer) writeBody(body ast.Body, comments []*ast.Comment) []*ast.Comment {
 	comments = w.insertComments(comments, body.Loc())
-	offset := 0
 	for i, expr := range body {
-		if i > 0 && expr.Location.Row-body[i-1].Location.Row-offset > 1 {
-			w.blankLine()
+		// Insert a blank line in before the expression if it was not right
+		// after the previous expression.
+		if i > 0 {
+			lastRow := body[i-1].Location.Row
+			for _, c := range body[i-1].Location.Text {
+				if c == '\n' {
+					lastRow++
+				}
+			}
+			if expr.Location.Row > lastRow+1 {
+				w.blankLine()
+			}
 		}
 		w.startLine()
 

format/testfiles/test.rego.formatted

@@ -45,7 +45,6 @@ r = y {
 else = y {
 	y = ["howdy"]
 	x = {"x": {"y": "z"}}
-
 	a = {
 		"a": {"b": "c"},
 		"b": "c", "c": [
@@ -126,7 +125,6 @@ p[x] = y {
 		"d": "e",
 		# Comment before closing object brace.
 	} # Comment on closing object brace.
-
 	a = {"a": "b", "c": "d"}
 	b = [1, 2, 3, 4]
 	c = [
@@ -150,7 +148,6 @@ p[x] = y {
 		split("foo.bar", ".", x)
 		x[_]
 	]
-
 	g = [1 |
 		split("foo.bar", ".", x) # comment in array comprehension
 		x[_]
@@ -163,7 +160,6 @@ p[x] = y {
 		split("foo.bar", ".", x)
 		x[_]
 	}
-
 	k = {1 |
 		split("foo.bar", ".", x) # comment in set comprehension
 		x[_]
@@ -175,12 +171,10 @@ p[x] = y {
 		split("foo.bar", ".", x)
 		y = x[_]
 	}
-
 	n = {y: x |
 		split("foo.bar", ".", x)
 		y = x[_]
 	}
-
 	o = {y: x |
 		split("foo.bar", ".", x) # comment in object comprehension
 		y = x[_]
@@ -197,7 +191,6 @@ nested_infix {
 	f(x, y)
 	y = (x + 1) + 2
 	x = y + z # comment
-
 	x = (a + b) / 2
 	f((a + b) / 2)
 	y = q()

format/testfiles/test_fun_args_with_linebreaks.rego.formatted

@@ -2,12 +2,10 @@ package test
 
 p {
 	x := count([1, 2, 3]) # four
-
 	y := concat(
 		"/",
 		["foo", "bar"],
 	)
-
 	z := concat(
 		"/",
 		[