From 3cc3e3017f51d35121baf3d024a64c9b909b9968 Mon Sep 17 00:00:00 2001
From: Jasper Van der Jeugt <m@jaspervdj.be>
Date: Mon, 16 May 2022 16:46:14 +0200
Subject: [PATCH] Fix opa fmt location for non-key rules

Running the following through `opa fmt`:

    package foo

    bar {
    	# before
    	input.bar
    	# after
    }

Causes the `after` comment to be moved outside of the rule:

    package foo

    bar {
    	# before
    	input.bar
    }

    # after

This was caused by `skipPast` in `closingLoc` being called even when there is no
`[key]` part in the rule head.  Adding a third clause fixed this; it seems
like `closingLoc` is designed to take `0` in this case because of the
`skipOpen > 0`.

This did affect one other test case, where I had to add an extra newline
to separate the comment from the rule head.  Without that, `insertComments`
(correctly, I guess) inserts:

    } # some special case

Signed-off-by: Jasper Van der Jeugt <m@jaspervdj.be>
---
 format/format.go                                         | 4 +++-
 format/testfiles/test_end_of_rule_comment.rego           | 7 +++++++
 format/testfiles/test_end_of_rule_comment.rego.formatted | 7 +++++++
 format/testfiles/test_issue_2299.rego                    | 3 ++-
 4 files changed, 19 insertions(+), 2 deletions(-)
 create mode 100644 format/testfiles/test_end_of_rule_comment.rego
 create mode 100644 format/testfiles/test_end_of_rule_comment.rego.formatted

diff --git a/format/format.go b/format/format.go
index c55d7998d1..2376c94db3 100644
--- a/format/format.go
+++ b/format/format.go
@@ -300,8 +300,10 @@ func (w *writer) writeRule(rule *ast.Rule, isElse bool, comments []*ast.Comment)
 
 	if len(rule.Head.Args) > 0 {
 		closeLoc = closingLoc('(', ')', '{', '}', rule.Location)
-	} else {
+	} else if rule.Head.Key != nil {
 		closeLoc = closingLoc('[', ']', '{', '}', rule.Location)
+	} else {
+		closeLoc = closingLoc(0, 0, '{', '}', rule.Location)
 	}
 
 	comments = w.insertComments(comments, closeLoc)
diff --git a/format/testfiles/test_end_of_rule_comment.rego b/format/testfiles/test_end_of_rule_comment.rego
new file mode 100644
index 0000000000..0c3bd75bee
--- /dev/null
+++ b/format/testfiles/test_end_of_rule_comment.rego
@@ -0,0 +1,7 @@
+package foo
+
+bar {
+	# before
+	input.bar
+	# after
+}
diff --git a/format/testfiles/test_end_of_rule_comment.rego.formatted b/format/testfiles/test_end_of_rule_comment.rego.formatted
new file mode 100644
index 0000000000..0c3bd75bee
--- /dev/null
+++ b/format/testfiles/test_end_of_rule_comment.rego.formatted
@@ -0,0 +1,7 @@
+package foo
+
+bar {
+	# before
+	input.bar
+	# after
+}
diff --git a/format/testfiles/test_issue_2299.rego b/format/testfiles/test_issue_2299.rego
index 5f1c5b74f0..5af69145d9 100644
--- a/format/testfiles/test_issue_2299.rego
+++ b/format/testfiles/test_issue_2299.rego
@@ -25,7 +25,8 @@ else = z {
 # Mixed compact and newline separated
 p = x {
     foo == "bar"
-} # some special case
+}
+# some special case
 # with lots of comments
 # describing it
 else = y {