Fixes open-policy-agent#4376

Do note though that this does not change how multiple
with statements are grouped. Although I agree with that,
it's IMHO a separate feature request, while the spacing
issue is a bug.

Signed-off-by: Anders Eknert <anders@eknert.com>

format/format.go

@@ -110,7 +110,7 @@ func Ast(x interface{}) ([]byte, error) {
 	case *ast.Expr:
 		w.writeExpr(x, nil)
 	case *ast.With:
-		w.writeWith(x, nil)
+		w.writeWith(x, nil, false)
 	case *ast.Term:
 		w.writeTerm(x, nil)
 	case ast.Value:
@@ -452,7 +452,7 @@ func (w *writer) writeExpr(expr *ast.Expr, comments []*ast.Comment) []*ast.Comme
 			w.endLine()
 			w.startLine()
 		}
-		comments = w.writeWith(with, comments)
+		comments = w.writeWith(with, comments, indented)
 	}
 
 	return comments
@@ -556,9 +556,12 @@ func (w *writer) writeFunctionCallPlain(terms []*ast.Term, comments []*ast.Comme
 	return w.writeIterable(args, loc, closingLoc(0, 0, '(', ')', loc), comments, w.listWriter())
 }
 
-func (w *writer) writeWith(with *ast.With, comments []*ast.Comment) []*ast.Comment {
+func (w *writer) writeWith(with *ast.With, comments []*ast.Comment, indented bool) []*ast.Comment {
 	comments = w.insertComments(comments, with.Location)
-	w.write(" with ")
+	if !indented {
+		w.write(" ")
+	}
+	w.write("with ")
 	comments = w.writeTerm(with.Target, comments)
 	w.write(" as ")
 	return w.writeTerm(with.Value, comments)

format/testfiles/test.rego

@@ -195,18 +195,6 @@ declare1 := 1
 
 declare2 := 2 { false }
 
-multi_line_with {
-    fn(1) with input.a as "a"
-                with input.b as "b"
-            with input.c as {
-                "foo": "bar",
-            }
-                with input.d as [
-                    1,
-                    2,
-                    3]
-}
-
 # more comments!
 # more comments!
 # more comments!

format/testfiles/test.rego.formatted

@@ -226,17 +226,6 @@ declare2 := 2 {
 	false
 }
 
-multi_line_with {
-	fn(1) with input.a as "a"
-		 with input.b as "b"
-		 with input.c as {"foo": "bar"}
-		 with input.d as [
-			1,
-			2,
-			3,
-		]
-}
-
 # more comments!
 # more comments!
 # more comments!

format/testfiles/test_with.rego

@@ -0,0 +1,17 @@
+package p
+
+single_line_with {
+  fn(1)   with input.a as "a"
+}
+
+multi_line_with {
+    fn(1) with input.a as "a"
+                with input.b as "b"
+            with input.c as {
+                "foo": "bar",
+            }
+                with input.d as [
+                    1,
+                    2,
+                    3]
+}

format/testfiles/test_with.rego.formatted

@@ -0,0 +1,16 @@
+package p
+
+single_line_with {
+	fn(1) with input.a as "a"
+}
+
+multi_line_with {
+	fn(1) with input.a as "a"
+		with input.b as "b"
+		with input.c as {"foo": "bar"}
+		with input.d as [
+			1,
+			2,
+			3,
+		]
+}