From e7a95372b274ca4ba3b99915c6fe7a8b187f70b8 Mon Sep 17 00:00:00 2001
From: Stephan Renatus <stephan.renatus@gmail.com>
Date: Fri, 22 Apr 2022 16:12:48 +0200
Subject: [PATCH] format: don't add 'in' keyword import when 'every' is there
 (#4607)

Also ensure that added imports have a location set.

Previously, `opa fmt` on the added test file would have panicked
because the import hadn't had a location.

Fixes #4606.

Signed-off-by: Stephan Renatus <stephan.renatus@gmail.com>
---
 format/format.go                              | 23 ++++++++++++++-----
 ..._in_operator_without_import.rego.formatted |  1 +
 format/testfiles/test_issue_4606.rego         |  6 +++++
 .../testfiles/test_issue_4606.rego.formatted  |  8 +++++++
 4 files changed, 32 insertions(+), 6 deletions(-)
 create mode 100644 format/testfiles/test_issue_4606.rego
 create mode 100644 format/testfiles/test_issue_4606.rego.formatted

diff --git a/format/format.go b/format/format.go
index 18875a7679..7095f79feb 100644
--- a/format/format.go
+++ b/format/format.go
@@ -60,7 +60,7 @@ func Ast(x interface{}) ([]byte, error) {
 	// or internal.member_3, it will sugarize them into usage of the `in`
 	// operator. It has to ensure that the proper future keyword import is
 	// present.
-	extraFutureKeywordImports := map[string]bool{}
+	extraFutureKeywordImports := map[string]struct{}{}
 
 	// Preprocess the AST. Set any required defaults and calculate
 	// values required for printing the formatted output.
@@ -76,9 +76,9 @@ func Ast(x interface{}) ([]byte, error) {
 		case *ast.Expr:
 			switch {
 			case n.IsCall() && ast.Member.Ref().Equal(n.Operator()) || ast.MemberWithKey.Ref().Equal(n.Operator()):
-				extraFutureKeywordImports["in"] = true
+				extraFutureKeywordImports["in"] = struct{}{}
 			case n.IsEvery():
-				extraFutureKeywordImports["every"] = true
+				extraFutureKeywordImports["every"] = struct{}{}
 			}
 		}
 		if x.Loc() == nil {
@@ -1239,11 +1239,22 @@ func (w *writer) down() {
 
 func ensureFutureKeywordImport(imps []*ast.Import, kw string) []*ast.Import {
 	allKeywords := ast.MustParseTerm("future.keywords")
-	kwPath := ast.MustParseTerm("future.keywords." + kw)
+	kwPath := keyword(kw)
+	every := keyword("every")
 	for _, imp := range imps {
-		if allKeywords.Equal(imp.Path) || imp.Path.Equal(kwPath) {
+		if allKeywords.Equal(imp.Path) ||
+			imp.Path.Equal(kwPath) ||
+			(imp.Path.Equal(every) && kw == "in") { // "every" implies "in", so we don't need to add both
 			return imps
 		}
 	}
-	return append(imps, &ast.Import{Path: kwPath})
+	imp := &ast.Import{
+		Path: kwPath,
+	}
+	imp.Location = defaultLocation(imp)
+	return append(imps, imp)
+}
+
+func keyword(kw string) *ast.Term {
+	return ast.MustParseTerm("future.keywords." + kw)
 }
diff --git a/format/testfiles/test_in_operator_without_import.rego.formatted b/format/testfiles/test_in_operator_without_import.rego.formatted
index efd6f755e7..040c28dfdb 100644
--- a/format/testfiles/test_in_operator_without_import.rego.formatted
+++ b/format/testfiles/test_in_operator_without_import.rego.formatted
@@ -1,6 +1,7 @@
 package p
 
 import future.keywords.in
+
 import input.foo
 
 r {
diff --git a/format/testfiles/test_issue_4606.rego b/format/testfiles/test_issue_4606.rego
new file mode 100644
index 0000000000..a71142ed97
--- /dev/null
+++ b/format/testfiles/test_issue_4606.rego
@@ -0,0 +1,6 @@
+package p
+
+#
+import future.keywords.every
+
+r { 1 in [] }
\ No newline at end of file
diff --git a/format/testfiles/test_issue_4606.rego.formatted b/format/testfiles/test_issue_4606.rego.formatted
new file mode 100644
index 0000000000..ce00ead8c9
--- /dev/null
+++ b/format/testfiles/test_issue_4606.rego.formatted
@@ -0,0 +1,8 @@
+package p
+
+#
+import future.keywords.every
+
+r {
+	1 in []
+}