format: keep whitespaces for multiple indented same-line withs

Fixes #4634.

Signed-off-by: Stephan Renatus <stephan.renatus@gmail.com>

format/format.go

@@ -455,7 +455,9 @@ func (w *writer) writeExpr(expr *ast.Expr, comments []*ast.Comment) []*ast.Comme
 
 	var indented bool
 	for i, with := range expr.With {
-		if i > 0 && with.Location.Row-expr.With[i-1].Location.Row > 0 {
+		if i == 0 || with.Location.Row == expr.With[i-1].Location.Row { // we're on the same line
+			comments = w.writeWith(with, comments, false)
+		} else { // we're on a new line
 			if !indented {
 				indented = true
 
@@ -464,8 +466,8 @@ func (w *writer) writeExpr(expr *ast.Expr, comments []*ast.Comment) []*ast.Comme
 			}
 			w.endLine()
 			w.startLine()
+			comments = w.writeWith(with, comments, true)
 		}
-		comments = w.writeWith(with, comments, indented)
 	}
 
 	return comments

format/testfiles/test_with.rego

@@ -16,6 +16,12 @@ multi_line_with {
                     3]
 }
 
+mixed_new_lines_with {
+    true with input.a as "a"
+      with input.b as "b" with input.c as "c"
+      with input.d as "d"
+}
+
 mock_f(_) = 123
 
 func_replacements {

format/testfiles/test_with.rego.formatted

@@ -15,6 +15,12 @@ multi_line_with {
 		]
 }
 
+mixed_new_lines_with {
+	true with input.a as "a"
+		with input.b as "b" with input.c as "c"
+		with input.d as "d"
+}
+
 mock_f(_) = 123
 
 func_replacements {