fmt: restrict = true rewrite to functions with implicit true

Fixes open-policy-agent#6467

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

format/format.go

@@ -535,7 +535,11 @@ func (w *writer) writeHead(head *ast.Head, isDefault, isExpandedConst bool, o fm
 		// * a.b.c.d -> a.b.c.d := true
 		isRegoV1RefConst := o.regoV1 && isExpandedConst && head.Key == nil && len(head.Args) == 0
 
-		if head.Location == head.Value.Location && head.Name != "else" && !isRegoV1RefConst {
+		if len(head.Args) > 0 &&
+			head.Location == head.Value.Location &&
+			head.Name != "else" &&
+			ast.Compare(head.Value, ast.BooleanTerm(true)) == 0 &&
+			!isRegoV1RefConst {
 			// If the value location is the same as the location of the head,
 			// we know that the value is generated, i.e. f(1)
 			// Don't print the value (` = true`) as it is implied.

format/testfiles/test_functions.rego

@@ -25,3 +25,24 @@ f7(x) := 1 {
 } else {
 	false
 }
+
+f(x) = 1 {
+	input.x == x
+} {
+	input.x < 10
+}
+
+f(_) {
+	input.x
+} {
+	input.y
+}
+
+# Non-functions
+foo[x] = y {
+	x = 1
+	y = 2
+} {
+	x = 3
+	y = 4
+}

format/testfiles/test_functions.rego.formatted

@@ -23,3 +23,30 @@ f7(x) := 1 {
 } else {
 	false
 }
+
+f(x) = 1 {
+	input.x == x
+}
+
+f(x) = 1 {
+	input.x < 10
+}
+
+f(_) {
+	input.x
+}
+
+f(_) {
+	input.y
+}
+
+# Non-functions
+foo[x] = y {
+	x = 1
+	y = 2
+}
+
+foo[x] = y {
+	x = 3
+	y = 4
+}