From bfd984ddf93ef2c4963a08d4fdadae0bcf1a3717 Mon Sep 17 00:00:00 2001
From: Stephan Renatus <stephan.renatus@gmail.com>
Date: Mon, 4 Oct 2021 08:11:41 +0200
Subject: [PATCH] format: make groupIterable sort by row

Before, it depended on the elements being passed in ordered by their rows.
Before https://github.com/open-policy-agent/opa/pull/3823, the iteration
order was the same as the row order; but with sorting the keys slice (which
determines iteration order) on creation, that was changed.

Now, we'll sort the elements within `groupIterable`.

Fixes #3849.

Signed-off-by: Stephan Renatus <stephan.renatus@gmail.com>
---
 format/format.go                              | 13 ++++++--
 format/testfiles/test.rego.formatted          |  5 ++-
 format/testfiles/test_issue_3849.rego         | 33 +++++++++++++++++++
 .../testfiles/test_issue_3849.rego.formatted  | 33 +++++++++++++++++++
 4 files changed, 80 insertions(+), 4 deletions(-)
 create mode 100644 format/testfiles/test_issue_3849.rego
 create mode 100644 format/testfiles/test_issue_3849.rego.formatted

diff --git a/format/format.go b/format/format.go
index be1ff2739b..8af159510d 100644
--- a/format/format.go
+++ b/format/format.go
@@ -810,10 +810,17 @@ func (w *writer) listWriter() entryWriter {
 	}
 }
 
-func groupIterable(elements []interface{}, last *ast.Location) (lines [][]interface{}) {
+// groupIterable will group the `elements` slice into slices according to their
+// location: anything on the same line will be put into a slice.
+func groupIterable(elements []interface{}, last *ast.Location) [][]interface{} {
+	sort.Slice(elements, func(i, j int) bool {
+		return locLess(elements[i], elements[j])
+	})
+	var lines [][]interface{}
 	var cur []interface{}
 	for i, t := range elements {
-		loc := getLoc(t)
+		elem := t
+		loc := getLoc(elem)
 		lineDiff := loc.Row - last.Row
 		if lineDiff > 0 && i > 0 {
 			lines = append(lines, cur)
@@ -821,7 +828,7 @@ func groupIterable(elements []interface{}, last *ast.Location) (lines [][]interf
 		}
 
 		last = loc
-		cur = append(cur, t)
+		cur = append(cur, elem)
 	}
 	return append(lines, cur)
 }
diff --git a/format/testfiles/test.rego.formatted b/format/testfiles/test.rego.formatted
index 7dd3acc41c..b7bff96823 100644
--- a/format/testfiles/test.rego.formatted
+++ b/format/testfiles/test.rego.formatted
@@ -19,7 +19,10 @@ foo[x] {
 	g(x, "foo") = z
 }
 
-globals = {"fizz": "buzz", "foo": "bar"}
+globals = {
+	"foo": "bar",
+	"fizz": "buzz",
+}
 
 partial_obj["x"] = 1
 
diff --git a/format/testfiles/test_issue_3849.rego b/format/testfiles/test_issue_3849.rego
new file mode 100644
index 0000000000..4ebfeb60ce
--- /dev/null
+++ b/format/testfiles/test_issue_3849.rego
@@ -0,0 +1,33 @@
+package test_issue_3849
+
+test_require_context {
+	require_context("monkey", "eat", "banana") with input as {
+		"principal": {"id": 101, "type": "monkey"},
+		"action": "eat",
+		"entity": {"id": 102, "type": "banana"},
+	}
+}
+
+test_contrived {
+	allow with input as {
+		"a": 101,
+		"b": 101,
+		"z": 101,
+		"y": 101,
+		"x": 101,
+		"w": 101,
+		"v": 101,
+		"u": 101,
+		"t": 101,
+		"s": 101,
+		"r": 101,
+		"q": 101,
+		"p": 101,
+		"o": 101,
+		"n": 101,
+		"j": 101,
+		"k": 101,
+		"l": 101,
+		"m": 101,
+	}
+}
\ No newline at end of file
diff --git a/format/testfiles/test_issue_3849.rego.formatted b/format/testfiles/test_issue_3849.rego.formatted
new file mode 100644
index 0000000000..d3194bbe6a
--- /dev/null
+++ b/format/testfiles/test_issue_3849.rego.formatted
@@ -0,0 +1,33 @@
+package test_issue_3849
+
+test_require_context {
+	require_context("monkey", "eat", "banana") with input as {
+		"principal": {"id": 101, "type": "monkey"},
+		"action": "eat",
+		"entity": {"id": 102, "type": "banana"},
+	}
+}
+
+test_contrived {
+	allow with input as {
+		"a": 101,
+		"b": 101,
+		"z": 101,
+		"y": 101,
+		"x": 101,
+		"w": 101,
+		"v": 101,
+		"u": 101,
+		"t": 101,
+		"s": 101,
+		"r": 101,
+		"q": 101,
+		"p": 101,
+		"o": 101,
+		"n": 101,
+		"j": 101,
+		"k": 101,
+		"l": 101,
+		"m": 101,
+	}
+}