From 13dc7ec66f1c0da0958c5d2423acb0a8e284684a Mon Sep 17 00:00:00 2001
From: Anders Eknert <anders@eknert.com>
Date: Sat, 26 Mar 2022 00:05:58 +0100
Subject: [PATCH] Fixes #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                          | 11 +++++++----
 format/testfiles/test.rego                | 12 ------------
 format/testfiles/test.rego.formatted      | 11 -----------
 format/testfiles/test_with.rego           | 17 +++++++++++++++++
 format/testfiles/test_with.rego.formatted | 16 ++++++++++++++++
 5 files changed, 40 insertions(+), 27 deletions(-)
 create mode 100644 format/testfiles/test_with.rego
 create mode 100644 format/testfiles/test_with.rego.formatted

diff --git a/format/format.go b/format/format.go
index 0d54a85225..18875a7679 100644
--- a/format/format.go
+++ b/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)
diff --git a/format/testfiles/test.rego b/format/testfiles/test.rego
index d73f41ac00..d24ddff211 100644
--- a/format/testfiles/test.rego
+++ b/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!
diff --git a/format/testfiles/test.rego.formatted b/format/testfiles/test.rego.formatted
index cbb8e93132..18d270a592 100644
--- a/format/testfiles/test.rego.formatted
+++ b/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!
diff --git a/format/testfiles/test_with.rego b/format/testfiles/test_with.rego
new file mode 100644
index 0000000000..0d727bccaa
--- /dev/null
+++ b/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]
+}
diff --git a/format/testfiles/test_with.rego.formatted b/format/testfiles/test_with.rego.formatted
new file mode 100644
index 0000000000..546523c312
--- /dev/null
+++ b/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,
+		]
+}