fmt: allow one liner rule grouping (#7453)

While the double newline added by the formatter after each rule makes sense
for most rules, short one-liner rules should be groupable. This PR changes
the behavior of the formatter, so that if the user does:

```rego
x := 1
y := 2
```

That is no longer formatted into:

```rego
x := 1

y := 2
```

If the user **wants** double newlines between one-liner rules, the formatter
respects those when present.

Note that the `default` rules are excepted even when a one-liner, as presenting
these separately helps understanding the policy.

Fixes #6760

Do note that the issued mentioned doing this only for incremental rules, i.e.
to group only rules of the same name. I changed my mind on that though, as grouping
"constants" should be possible too.

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

v1/format/format.go

@@ -437,9 +437,19 @@ func (w *writer) writeComments(comments []*ast.Comment) {
 }
 
 func (w *writer) writeRules(rules []*ast.Rule, comments []*ast.Comment) []*ast.Comment {
-	for _, rule := range rules {
+	for i, rule := range rules {
 		comments = w.insertComments(comments, rule.Location)
 		comments = w.writeRule(rule, false, comments)
+
+		if i < len(rules)-1 && w.groupableOneLiner(rule) {
+			next := rules[i+1]
+			if w.groupableOneLiner(next) && next.Location.Row == rule.Location.Row+1 {
+				// Current rule and the next are both groupable one-liners, and
+				// adjacent in the original policy (i.e. no extra newlines between them).
+				continue
+			}
+		}
+
 		w.blankLine()
 	}
 	return comments
@@ -447,6 +457,18 @@ func (w *writer) writeRules(rules []*ast.Rule, comments []*ast.Comment) []*ast.C
 
 var expandedConst = ast.NewBody(ast.NewExpr(ast.InternedBooleanTerm(true)))
 
+func (w *writer) groupableOneLiner(rule *ast.Rule) bool {
+	// Location required to determine if two rules are adjacent in the policy.
+	// If not, we respect line breaks between rules.
+	if len(rule.Body) > 1 || rule.Default || rule.Location == nil {
+		return false
+	}
+
+	partialSetException := w.fmtOpts.contains || rule.Head.Value != nil
+
+	return (w.fmtOpts.regoV1 || w.fmtOpts.ifs) && partialSetException
+}
+
 func (w *writer) writeRule(rule *ast.Rule, isElse bool, comments []*ast.Comment) []*ast.Comment {
 	if rule == nil {
 		return comments
@@ -468,14 +490,14 @@ func (w *writer) writeRule(rule *ast.Rule, isElse bool, comments []*ast.Comment)
 
 	comments = w.writeHead(rule.Head, rule.Default, isExpandedConst, comments)
 
-	// this excludes partial sets UNLESS `contains` is used
-	partialSetException := w.fmtOpts.contains || rule.Head.Value != nil
-
 	if len(rule.Body) == 0 || isExpandedConst {
 		w.endLine()
 		return comments
 	}
 
+	// this excludes partial sets UNLESS `contains` is used
+	partialSetException := w.fmtOpts.contains || rule.Head.Value != nil
+
 	if (w.fmtOpts.regoV1 || w.fmtOpts.ifs) && partialSetException {
 		w.write(" if")
 		if len(rule.Body) == 1 {

v1/format/format_test.go

@@ -907,6 +907,34 @@ p contains x if {
 	}
 }
 
+func TestGroupableOneLinerRules(t *testing.T) {
+	contents := []byte(`package test
+
+foo := 1 if input.x
+foo := 2 if not input.x
+
+a := 1
+b := 2
+
+c := 3
+
+d := 4
+
+# comment above group
+e := 5
+f := 6
+`)
+
+	formatted, err := Source("test.rego", contents)
+	if err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+
+	if !bytes.Equal(formatted, contents) {
+		t.Fatalf("expected %q but got %q", formatted, contents)
+	}
+}
+
 // 382	   3064960 ns/op	 4573131 B/op	   26266 allocs/op // no optimizations
 // 685	   1737719 ns/op	 1972193 B/op	   14160 allocs/op // pre-allocate partitionComments
 // 708	   1674343 ns/op	 1916700 B/op	   11556 allocs/op // static memberRef & memberWithKeyRef

v1/format/testfiles/v0/test_ref_heads.rego.formatted

@@ -3,19 +3,13 @@ package test
 import future.keywords
 
 a.b.c = "d"
-
 a.b.e = "f"
-
 a.b.g contains x if some x in numbers.range(1, 3)
-
 a.b.h[x] = 1 if x := "one"
 
 q[1] = y
-
 r[x] if x := 10
-
 p.q.r[x] if x := 10
-
 p.q.r[2] = true
 
 g[h].i[j].k = true
@@ -31,7 +25,6 @@ g[3].i[j].k = x if {
 }
 
 g[h].i[j].k[l] = true
-
 g[h].i[j].k[l] contains x if x = "foo"
 
 g[h].i[j].k[l] contains x if {

v1/format/testfiles/v1/test.rego.formatted

@@ -25,7 +25,6 @@ globals := {
 }
 
 partial_obj["x"] := 1
-
 partial_obj["y"] := 2
 
 partial_obj["z"] := 3
@@ -79,7 +78,6 @@ short(x) if {
 }
 
 raw_string := `hi\there`
-
 raw_multiline := `this
 string
   is on
@@ -214,7 +212,6 @@ declare1 := 1
 declare2 := 2 if false
 
 declare3 := {1, 2, 3}
-
 declare4 := {4, 5, 6}
 
 union_object := {"response": (declare3 | declare4)}

v1/format/testfiles/v1/test_grouping.rego

@@ -0,0 +1,49 @@
+package p
+
+x := 1
+y := 2
+
+x1 := 1
+
+x2 := 2
+
+a.b.c := 1
+a.b.d := 2
+
+s contains 1
+s contains 2
+
+s contains 3
+
+s contains 4
+
+rule if foo == bar
+rule if bar == foo
+
+long if {
+    x := 1
+    y := 2
+}
+long if {
+    x := 2
+    y := 3
+}
+
+short if condition
+not_short if {
+	rule
+	body
+}
+
+not_short if {
+	rule
+	body
+}
+short if condition
+
+s contains "foo"
+s contains "bar" if {
+	foo
+	bar 
+}
+

v1/format/testfiles/v1/test_grouping.rego.formatted

@@ -0,0 +1,52 @@
+package p
+
+x := 1
+y := 2
+
+x1 := 1
+
+x2 := 2
+
+a.b.c := 1
+a.b.d := 2
+
+s contains 1
+s contains 2
+
+s contains 3
+
+s contains 4
+
+rule if foo == bar
+rule if bar == foo
+
+long if {
+	x := 1
+	y := 2
+}
+
+long if {
+	x := 2
+	y := 3
+}
+
+short if condition
+
+not_short if {
+	rule
+	body
+}
+
+not_short if {
+	rule
+	body
+}
+
+short if condition
+
+s contains "foo"
+
+s contains "bar" if {
+	foo
+	bar
+}

v1/format/testfiles/v1/test_in_operator_with_parenthesis.rego.formatted

@@ -14,5 +14,4 @@ z if {
 }
 
 f(_, _) := true
-
 g(_) := true

v1/format/testfiles/v1/test_issue_5537_with_ref.rego.formatted

@@ -1,7 +1,6 @@
 package p
 
 first := {"one", "two"}
-
 second := {"two", "three"}
 
 example contains msg if {

v1/format/testfiles/v1/test_ref_heads.rego.formatted

@@ -1,19 +1,13 @@
 package test
 
 a.b.c := "d"
-
 a.b.e := "f"
-
 a.b.g contains x if some x in numbers.range(1, 3)
-
 a.b.h[x] := 1 if x := "one"
 
 q[1] := y
-
 r[x] if x := 10
-
 p.q.r[x] if x := 10
-
 p.q.r[2] := true
 
 g[h].i[j].k := true
@@ -29,7 +23,6 @@ g[3].i[j].k := x if {
 }
 
 g[h].i[j].k[l] := true
-
 g[h].i[j].k[l] contains x if x = "foo"
 
 g[h].i[j].k[l] contains x if {