Fix opa fmt location for non-key rules (#4695)

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

@@ -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)

format/testfiles/test_end_of_rule_comment.rego

@@ -0,0 +1,7 @@
+package foo
+
+bar {
+	# before
+	input.bar
+	# after
+}

format/testfiles/test_end_of_rule_comment.rego.formatted

@@ -0,0 +1,7 @@
+package foo
+
+bar {
+	# before
+	input.bar
+	# after
+}

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 {