format: Bracketing keyword ref elements in formatter output (#7010)

Also future-proofing format pkg tests to be 1.0 compatible.

Signed-off-by: Johan Fylling <johan.dev@fylling.se>

ast/policy.go

@@ -100,7 +100,7 @@ var Wildcard = &Term{Value: Var("_")}
 var WildcardPrefix = "$"
 
 // Keywords contains strings that map to language keywords.
-var Keywords = KeywordsV0
+var Keywords = KeywordsForRegoVersion(DefaultRegoVersion)
 
 var KeywordsV0 = [...]string{
 	"not",
@@ -134,9 +134,23 @@ var KeywordsV1 = [...]string{
 	"every",
 }
 
+func KeywordsForRegoVersion(v RegoVersion) []string {
+	switch v {
+	case RegoV0:
+		return KeywordsV0[:]
+	case RegoV1, RegoV0CompatV1:
+		return KeywordsV1[:]
+	}
+	return nil
+}
+
 // IsKeyword returns true if s is a language keyword.
 func IsKeyword(s string) bool {
-	for _, x := range Keywords {
+	return IsInKeywords(s, Keywords)
+}
+
+func IsInKeywords(s string, keywords []string) bool {
+	for _, x := range keywords {
 		if x == s {
 			return true
 		}

compile/compile_test.go

@@ -1741,7 +1741,7 @@ contains := 2 {
 import rego.v1
 
 p if {
-	data.foo.contains = input.x
+	data.foo["contains"] = input.x
 }
 `,
 				`package foo

format/format.go

@@ -126,7 +126,16 @@ type fmtOpts struct {
 	// than if they don't.
 	refHeads bool
 
-	regoV1 bool
+	regoV1         bool
+	futureKeywords []string
+}
+
+func (o fmtOpts) keywords() []string {
+	if o.regoV1 {
+		return ast.KeywordsV1[:]
+	}
+	kws := ast.KeywordsV0[:]
+	return append(kws, o.futureKeywords...)
 }
 
 func AstWithOpts(x interface{}, opts Opts) ([]byte, error) {
@@ -171,6 +180,10 @@ func AstWithOpts(x interface{}, opts Opts) ([]byte, error) {
 			}
 
 		case *ast.Import:
+			if kw, ok := future.WhichFutureKeyword(n); ok {
+				o.futureKeywords = append(o.futureKeywords, kw)
+			}
+
 			switch {
 			case isRegoV1Compatible(n):
 				o.contains = true
@@ -200,8 +213,9 @@ func AstWithOpts(x interface{}, opts Opts) ([]byte, error) {
 	})
 
 	w := &writer{
-		indent: "\t",
-		errs:   make([]*ast.Error, 0),
+		indent:  "\t",
+		errs:    make([]*ast.Error, 0),
+		fmtOpts: o,
 	}
 
 	switch x := x.(type) {
@@ -219,18 +233,17 @@ func AstWithOpts(x interface{}, opts Opts) ([]byte, error) {
 				x.Imports = ensureFutureKeywordImport(x.Imports, kw)
 			}
 		}
-		w.writeModule(x, o)
+		w.writeModule(x)
 	case *ast.Package:
 		w.writePackage(x, nil)
 	case *ast.Import:
 		w.writeImports([]*ast.Import{x}, nil)
 	case *ast.Rule:
-		w.writeRule(x, false /* isElse */, o, nil)
+		w.writeRule(x, false /* isElse */, nil)
 	case *ast.Head:
 		w.writeHead(x,
 			false, // isDefault
 			false, // isExpandedConst
-			o,
 			nil)
 	case ast.Body:
 		w.writeBody(x, nil)
@@ -302,9 +315,10 @@ type writer struct {
 	beforeEnd *ast.Comment
 	delay     bool
 	errs      ast.Errors
+	fmtOpts   fmtOpts
 }
 
-func (w *writer) writeModule(module *ast.Module, o fmtOpts) {
+func (w *writer) writeModule(module *ast.Module) {
 	var pkg *ast.Package
 	var others []interface{}
 	var comments []*ast.Comment
@@ -342,7 +356,7 @@ func (w *writer) writeModule(module *ast.Module, o fmtOpts) {
 		imports, others = gatherImports(others)
 		comments = w.writeImports(imports, comments)
 		rules, others = gatherRules(others)
-		comments = w.writeRules(rules, o, comments)
+		comments = w.writeRules(rules, comments)
 	}
 
 	for i, c := range comments {
@@ -365,7 +379,15 @@ func (w *writer) writePackage(pkg *ast.Package, comments []*ast.Comment) []*ast.
 	comments = w.insertComments(comments, pkg.Location)
 
 	w.startLine()
-	w.write(pkg.String())
+
+	// Omit head as all packages have the DefaultRootDocument prepended at parse time.
+	path := make(ast.Ref, len(pkg.Path)-1)
+	path[0] = ast.VarTerm(string(pkg.Path[1].Value.(ast.String)))
+	copy(path[1:], pkg.Path[2:])
+
+	w.write("package ")
+	w.writeRef(path)
+
 	w.blankLine()
 
 	return comments
@@ -380,16 +402,16 @@ func (w *writer) writeComments(comments []*ast.Comment) {
 	}
 }
 
-func (w *writer) writeRules(rules []*ast.Rule, o fmtOpts, comments []*ast.Comment) []*ast.Comment {
+func (w *writer) writeRules(rules []*ast.Rule, comments []*ast.Comment) []*ast.Comment {
 	for _, rule := range rules {
 		comments = w.insertComments(comments, rule.Location)
-		comments = w.writeRule(rule, false, o, comments)
+		comments = w.writeRule(rule, false, comments)
 		w.blankLine()
 	}
 	return comments
 }
 
-func (w *writer) writeRule(rule *ast.Rule, isElse bool, o fmtOpts, comments []*ast.Comment) []*ast.Comment {
+func (w *writer) writeRule(rule *ast.Rule, isElse bool, comments []*ast.Comment) []*ast.Comment {
 	if rule == nil {
 		return comments
 	}
@@ -408,17 +430,17 @@ func (w *writer) writeRule(rule *ast.Rule, isElse bool, o fmtOpts, comments []*a
 	// pretend that the rule has no body in this case.
 	isExpandedConst := rule.Body.Equal(ast.NewBody(ast.NewExpr(ast.BooleanTerm(true)))) && rule.Else == nil
 
-	comments = w.writeHead(rule.Head, rule.Default, isExpandedConst, o, comments)
+	comments = w.writeHead(rule.Head, rule.Default, isExpandedConst, comments)
 
 	// this excludes partial sets UNLESS `contains` is used
-	partialSetException := o.contains || rule.Head.Value != nil
+	partialSetException := w.fmtOpts.contains || rule.Head.Value != nil
 
 	if len(rule.Body) == 0 || isExpandedConst {
 		w.endLine()
 		return comments
 	}
 
-	if (o.regoV1 || o.ifs) && partialSetException {
+	if (w.fmtOpts.regoV1 || w.fmtOpts.ifs) && partialSetException {
 		w.write(" if")
 		if len(rule.Body) == 1 {
 			if rule.Body[0].Location.Row == rule.Head.Location.Row {
@@ -426,7 +448,7 @@ func (w *writer) writeRule(rule *ast.Rule, isElse bool, o fmtOpts, comments []*a
 				comments = w.writeExpr(rule.Body[0], comments)
 				w.endLine()
 				if rule.Else != nil {
-					comments = w.writeElse(rule, o, comments)
+					comments = w.writeElse(rule, comments)
 				}
 				return comments
 			}
@@ -454,12 +476,12 @@ func (w *writer) writeRule(rule *ast.Rule, isElse bool, o fmtOpts, comments []*a
 	w.startLine()
 	w.write("}")
 	if rule.Else != nil {
-		comments = w.writeElse(rule, o, comments)
+		comments = w.writeElse(rule, comments)
 	}
 	return comments
 }
 
-func (w *writer) writeElse(rule *ast.Rule, o fmtOpts, comments []*ast.Comment) []*ast.Comment {
+func (w *writer) writeElse(rule *ast.Rule, comments []*ast.Comment) []*ast.Comment {
 	// If there was nothing else on the line before the "else" starts
 	// then preserve this style of else block, otherwise it will be
 	// started as an "inline" else eg:
@@ -521,16 +543,16 @@ func (w *writer) writeElse(rule *ast.Rule, o fmtOpts, comments []*ast.Comment) [
 		rule.Else.Head.Value.Location = rule.Else.Head.Location
 	}
 
-	return w.writeRule(rule.Else, true, o, comments)
+	return w.writeRule(rule.Else, true, comments)
 }
 
-func (w *writer) writeHead(head *ast.Head, isDefault, isExpandedConst bool, o fmtOpts, comments []*ast.Comment) []*ast.Comment {
+func (w *writer) writeHead(head *ast.Head, isDefault, isExpandedConst bool, comments []*ast.Comment) []*ast.Comment {
 	ref := head.Ref()
 	if head.Key != nil && head.Value == nil && !head.HasDynamicRef() {
 		ref = ref.GroundPrefix()
 	}
-	if o.refHeads || len(ref) == 1 {
-		w.write(ref.String())
+	if w.fmtOpts.refHeads || len(ref) == 1 {
+		w.writeRef(ref)
 	} else {
 		w.write(ref[0].String())
 		w.write("[")
@@ -548,7 +570,7 @@ func (w *writer) writeHead(head *ast.Head, isDefault, isExpandedConst bool, o fm
 		w.write(")")
 	}
 	if head.Key != nil {
-		if o.contains && head.Value == nil {
+		if w.fmtOpts.contains && head.Value == nil {
 			w.write(" contains ")
 			comments = w.writeTerm(head.Key, comments)
 		} else if head.Value == nil { // no `if` for p[x] notation
@@ -566,7 +588,7 @@ func (w *writer) writeHead(head *ast.Head, isDefault, isExpandedConst bool, o fm
 		// * a.b -> a contains "b"
 		// * a.b.c -> a.b.c := true
 		// * a.b.c.d -> a.b.c.d := true
-		isRegoV1RefConst := o.regoV1 && isExpandedConst && head.Key == nil && len(head.Args) == 0
+		isRegoV1RefConst := w.fmtOpts.regoV1 && isExpandedConst && head.Key == nil && len(head.Args) == 0
 
 		if head.Location == head.Value.Location &&
 			head.Name != "else" &&
@@ -578,7 +600,7 @@ func (w *writer) writeHead(head *ast.Head, isDefault, isExpandedConst bool, o fm
 			return comments
 		}
 
-		if head.Assign || o.regoV1 {
+		if head.Assign || w.fmtOpts.regoV1 {
 			// preserve assignment operator, and enforce it if formatting for Rego v1
 			w.write(" := ")
 		} else {
@@ -856,7 +878,7 @@ var varRegexp = regexp.MustCompile("^[[:alpha:]_][[:alpha:][:digit:]_]*$")
 
 func (w *writer) writeRefStringPath(s ast.String) {
 	str := string(s)
-	if varRegexp.MatchString(str) && !ast.IsKeyword(str) {
+	if varRegexp.MatchString(str) && !ast.IsInKeywords(str, w.fmtOpts.keywords()) {
 		w.write("." + str)
 	} else {
 		w.writeBracketed(s.String())
@@ -1067,7 +1089,7 @@ func (w *writer) writeImports(imports []*ast.Import, comments []*ast.Comment) []
 		})
 		for _, i := range group {
 			w.startLine()
-			w.write(i.String())
+			w.writeImport(i)
 			if c, ok := m[i]; ok {
 				w.write(" " + c.String())
 			}
@@ -1079,6 +1101,28 @@ func (w *writer) writeImports(imports []*ast.Import, comments []*ast.Comment) []
 	return comments
 }
 
+func (w *writer) writeImport(imp *ast.Import) {
+	path := imp.Path.Value.(ast.Ref)
+
+	buf := []string{"import"}
+
+	if _, ok := future.WhichFutureKeyword(imp); ok {
+		// We don't want to wrap future.keywords imports in parens, so we create a new writer that doesn't
+		w2 := writer{
+			buf: bytes.Buffer{},
+		}
+		w2.writeRef(path)
+		buf = append(buf, w2.buf.String())
+	} else {
+		buf = append(buf, path.String())
+	}
+
+	if len(imp.Alias) > 0 {
+		buf = append(buf, "as "+imp.Alias.String())
+	}
+	w.write(strings.Join(buf, " "))
+}
+
 type entryWriter func(interface{}, []*ast.Comment) []*ast.Comment
 
 func (w *writer) writeIterable(elements []interface{}, last *ast.Location, close *ast.Location, comments []*ast.Comment, fn entryWriter) []*ast.Comment {
@@ -1505,7 +1549,12 @@ func ensureFutureKeywordImport(imps []*ast.Import, kw string) []*ast.Import {
 		}
 	}
 	imp := &ast.Import{
-		Path: ast.MustParseTerm("future.keywords." + kw),
+		// NOTE: This is a hack to not error on the ref containing a keyword already present in v1.
+		// A cleaner solution would be to instead allow refs to contain keyword terms.
+		// E.g. in v1, `import future.keywords["in"]` is valid, but `import future.keywords.in` is not
+		// as it contains a reserved keyword.
+		Path: ast.MustParseTerm("future.keywords[\"" + kw + "\"]"),
+		//Path: ast.MustParseTerm("future.keywords." + kw),
 	}
 	imp.Location = defaultLocation(imp)
 	return append(imps, imp)

format/format_test.go

@@ -17,21 +17,43 @@ import (
 )
 
 func TestFormatNilLocation(t *testing.T) {
-	rule := ast.MustParseRule(`r = y { y = "foo" }`)
-	rule.Head.Location = nil
-
-	bs, err := Ast(rule)
-	if err != nil {
-		t.Fatal(err)
+	tests := []struct {
+		note        string
+		regoVersion ast.RegoVersion
+		rule        string
+		exp         string
+	}{
+		{
+			note:        "v0",
+			regoVersion: ast.RegoV0,
+			rule:        `r = y { y = "foo" }`,
+			exp: `r = y {
+	y = "foo"
+}`,
+		},
+		{
+			note:        "v1",
+			regoVersion: ast.RegoV1,
+			rule:        `r = y if { y = "foo" }`,
+			exp: `r := y if y = "foo"
+`,
+		},
 	}
 
-	exp := strings.Trim(`
-r = y {
-	y = "foo"
-}`, " \n")
+	for _, tc := range tests {
+		t.Run(tc.note, func(t *testing.T) {
+			rule := ast.MustParseRuleWithOpts(tc.rule, ast.ParserOptions{RegoVersion: tc.regoVersion})
+			rule.Head.Location = nil
 
-	if string(bs) != exp {
-		t.Fatalf("Expected %q but got %q", exp, string(bs))
+			bs, err := AstWithOpts(rule, Opts{RegoVersion: tc.regoVersion})
+			if err != nil {
+				t.Fatal(err)
+			}
+
+			if string(bs) != tc.exp {
+				t.Fatalf("Expected:\n\n%q\n\nbut got:\n\n%q", tc.exp, string(bs))
+			}
+		})
 	}
 }
 
@@ -59,7 +81,7 @@ func TestFormatNilLocationFunctionArgs(t *testing.T) {
 }
 
 func TestFormatSourceError(t *testing.T) {
-	rego := "testfiles/test.rego.error"
+	rego := "testfiles/v0/test.rego.error"
 	contents, err := os.ReadFile(rego)
 	if err != nil {
 		t.Fatalf("Failed to read rego source: %v", err)
@@ -70,15 +92,15 @@ func TestFormatSourceError(t *testing.T) {
 		t.Fatal("Expected parsing error, not nil")
 	}
 
-	exp := "1 error occurred: testfiles/test.rego.error:27: rego_parse_error: unexpected eof token"
+	exp := "1 error occurred: testfiles/v0/test.rego.error:27: rego_parse_error: unexpected eof token"
 
 	if !strings.HasPrefix(err.Error(), exp) {
 		t.Fatalf("Expected error message '%s', got '%s'", exp, err.Error())
 	}
 }
 
-func TestFormatSource(t *testing.T) {
-	regoFiles, err := filepath.Glob("testfiles/*.rego")
+func TestFormatV0Source(t *testing.T) {
+	regoFiles, err := filepath.Glob("testfiles/v0/*.rego")
 	if err != nil {
 		panic(err)
 	}
@@ -95,7 +117,15 @@ func TestFormatSource(t *testing.T) {
 				t.Fatalf("Failed to read expected rego source: %v", err)
 			}
 
-			formatted, err := Source(rego, contents)
+			popts := ast.ParserOptions{
+				RegoVersion: ast.RegoV0,
+			}
+			opts := Opts{
+				RegoVersion:   ast.RegoV0,
+				ParserOptions: &popts,
+			}
+
+			formatted, err := SourceWithOpts(rego, contents, opts)
 			if err != nil {
 				t.Fatalf("Failed to format file: %v", err)
 			}
@@ -104,11 +134,11 @@ func TestFormatSource(t *testing.T) {
 				t.Fatalf("Expected formatted bytes to equal expected bytes but differed near line %d / byte %d (got: %q, expected: %q):\n%s", ln, at, formatted[at], expected[at], prefixWithLineNumbers(formatted))
 			}
 
-			if _, err := ast.ParseModule(rego+".tmp", string(formatted)); err != nil {
+			if _, err := ast.ParseModuleWithOpts(rego+".tmp", string(formatted), popts); err != nil {
 				t.Fatalf("Failed to parse formatted bytes: %v", err)
 			}
 
-			formatted, err = Source(rego, formatted)
+			formatted, err = SourceWithOpts(rego, formatted, opts)
 			if err != nil {
 				t.Fatalf("Failed to double format file")
 			}
@@ -121,8 +151,60 @@ func TestFormatSource(t *testing.T) {
 	}
 }
 
-func TestFormatSourceToRegoV1(t *testing.T) {
-	regoFiles, err := filepath.Glob("testfiles/rego_v1/*.rego")
+func TestFormatV1Source(t *testing.T) {
+	regoFiles, err := filepath.Glob("testfiles/v1/*.rego")
+	if err != nil {
+		panic(err)
+	}
+
+	for _, rego := range regoFiles {
+		t.Run(rego, func(t *testing.T) {
+			contents, err := os.ReadFile(rego)
+			if err != nil {
+				t.Fatalf("Failed to read rego source: %v", err)
+			}
+
+			expected, err := os.ReadFile(rego + ".formatted")
+			if err != nil {
+				t.Fatalf("Failed to read expected rego source: %v", err)
+			}
+
+			popts := ast.ParserOptions{
+				RegoVersion: ast.RegoV1,
+			}
+			opts := Opts{
+				RegoVersion:   ast.RegoV1,
+				ParserOptions: &popts,
+			}
+
+			formatted, err := SourceWithOpts(rego, contents, opts)
+			if err != nil {
+				t.Fatalf("Failed to format file: %v", err)
+			}
+
+			if ln, at := differsAt(formatted, expected); ln != 0 {
+				t.Fatalf("Expected formatted bytes to equal expected bytes but differed near line %d / byte %d (got: %q, expected: %q):\n%s", ln, at, formatted[at], expected[at], prefixWithLineNumbers(formatted))
+			}
+
+			if _, err := ast.ParseModuleWithOpts(rego+".tmp", string(formatted), popts); err != nil {
+				t.Fatalf("Failed to parse formatted bytes: %v", err)
+			}
+
+			formatted, err = SourceWithOpts(rego, formatted, opts)
+			if err != nil {
+				t.Fatalf("Failed to double format file")
+			}
+
+			if ln, at := differsAt(formatted, expected); ln != 0 {
+				t.Fatalf("Expected roundtripped bytes to equal expected bytes but differed near line %d / byte %d:\n%s", ln, at, prefixWithLineNumbers(formatted))
+			}
+
+		})
+	}
+}
+
+func TestFormatV0SourceToRegoV1(t *testing.T) {
+	regoFiles, err := filepath.Glob("testfiles/v0_to_v1/*.rego")
 	if err != nil {
 		panic(err)
 	}
@@ -149,8 +231,19 @@ func TestFormatSourceToRegoV1(t *testing.T) {
 				}
 			}
 
+			sourceOpts := Opts{
+				RegoVersion: ast.RegoV0CompatV1, // Target syntax is v0 compat v1
+				ParserOptions: &ast.ParserOptions{
+					RegoVersion: ast.RegoV0, // Original syntax is v0
+				},
+			}
+			targetOpts := Opts{
+				RegoVersion: ast.RegoV0CompatV1, // Target syntax is v0 compat v1
+			}
+
 			if errorExpected {
-				formatted, err := SourceWithOpts(rego, contents, Opts{RegoVersion: ast.RegoV0CompatV1})
+				formatted, err := SourceWithOpts(rego, contents, sourceOpts)
+
 				if err == nil {
 					t.Fatalf("Expected error, got: %s", formatted)
 				}
@@ -158,7 +251,8 @@ func TestFormatSourceToRegoV1(t *testing.T) {
 					t.Fatalf("Expected error:\n\n'%s'\n\ngot:\n\n'%s'", expected, err.Error())
 				}
 			} else {
-				formatted, err := SourceWithOpts(rego, contents, Opts{RegoVersion: ast.RegoV0CompatV1})
+				formatted, err := SourceWithOpts(rego, contents, sourceOpts)
+
 				if err != nil {
 					t.Fatalf("Failed to format file: %v", err)
 				}
@@ -171,7 +265,7 @@ func TestFormatSourceToRegoV1(t *testing.T) {
 					t.Fatalf("Failed to parse formatted bytes: %v", err)
 				}
 
-				formatted, err = SourceWithOpts(rego, formatted, Opts{RegoVersion: ast.RegoV0CompatV1})
+				formatted, err = SourceWithOpts(rego, formatted, targetOpts)
 				if err != nil {
 					t.Fatalf("Failed to double format file")
 				}
@@ -181,7 +275,7 @@ func TestFormatSourceToRegoV1(t *testing.T) {
 				}
 
 				// rego-v1 formatted code is still compliant with v0, and should not be changed if formatted as such
-				formatted, err = Source(rego, formatted)
+				formatted, err = SourceWithOpts(rego, formatted, targetOpts)
 				if err != nil {
 					t.Fatalf("Failed to double format file as v0")
 				}
@@ -196,9 +290,10 @@ func TestFormatSourceToRegoV1(t *testing.T) {
 
 func TestFormatAST(t *testing.T) {
 	cases := []struct {
-		note     string
-		toFmt    interface{}
-		expected string
+		note        string
+		regoVersion ast.RegoVersion
+		toFmt       interface{}
+		expected    string
 	}{
 		{
 			note:     "var",
@@ -321,12 +416,16 @@ func TestFormatAST(t *testing.T) {
 			expected: `some x, y in xs`,
 		},
 		{
-			note: "every adds import if missing",
+			note:        "v0, every adds import if missing",
+			regoVersion: ast.RegoV0,
 			toFmt: ast.MustParseModuleWithOpts(`package test
 			p {
 				every k, v in [1, 2] { k != v }
 			}`,
-				ast.ParserOptions{FutureKeywords: []string{"every"}}),
+				ast.ParserOptions{
+					RegoVersion:    ast.RegoV0,
+					FutureKeywords: []string{"every"},
+				}),
 			expected: `package test
 
 import future.keywords.every
@@ -336,13 +435,31 @@ p {
 }`,
 		},
 		{
-			note: "every does not add import if all future KWs are there",
+			note:        "v1, every doesn't add import if missing",
+			regoVersion: ast.RegoV1,
+			toFmt: ast.MustParseModuleWithOpts(`package test
+			p if {
+				every k, v in [1, 2] { k != v }
+			}`,
+				ast.ParserOptions{RegoVersion: ast.RegoV1}),
+			expected: `package test
+
+p if {
+	every k, v in [1, 2] { k != v }
+}`,
+		},
+		{
+			note:        "v0: every does not add import if all future KWs are there",
+			regoVersion: ast.RegoV0,
 			toFmt: ast.MustParseModuleWithOpts(`package test
 			import future.keywords
 			p {
 				every k, v in [1, 2] { k != v }
 			}`,
-				ast.ParserOptions{FutureKeywords: []string{"every"}}),
+				ast.ParserOptions{
+					FutureKeywords: []string{"every"},
+					RegoVersion:    ast.RegoV0,
+				}),
 			expected: `package test
 
 import future.keywords
@@ -352,13 +469,17 @@ p if {
 }`,
 		},
 		{
-			note: "every does not add import if already present",
+			note:        "v0: every does not add import if already present",
+			regoVersion: ast.RegoV0,
 			toFmt: ast.MustParseModuleWithOpts(`package test
 			import future.keywords
 			p {
 				every k, v in [1, 2] { k != v }
 			}`,
-				ast.ParserOptions{FutureKeywords: []string{"every"}}),
+				ast.ParserOptions{
+					FutureKeywords: []string{"every"},
+					RegoVersion:    ast.RegoV0,
+				}),
 			expected: `package test
 
 import future.keywords
@@ -520,7 +641,12 @@ a[_x[y][[z, w]]]`,
 
 	for _, tc := range cases {
 		t.Run(tc.note, func(t *testing.T) {
-			bs, err := Ast(tc.toFmt)
+			bs, err := AstWithOpts(tc.toFmt, Opts{
+				RegoVersion: tc.regoVersion,
+				ParserOptions: &ast.ParserOptions{
+					RegoVersion: tc.regoVersion,
+				},
+			})
 			if err != nil {
 				t.Fatalf("Unexpected error: %s", err)
 			}

format/testfiles/rego_v1/deprecated_builtins.rego.error

@@ -1,16 +0,0 @@
-15 errors occurred:
-testfiles/rego_v1/deprecated_builtins.rego:4: rego_type_error: deprecated built-in function calls in expression: any
-testfiles/rego_v1/deprecated_builtins.rego:5: rego_type_error: deprecated built-in function calls in expression: all
-testfiles/rego_v1/deprecated_builtins.rego:6: rego_type_error: deprecated built-in function calls in expression: cast_array
-testfiles/rego_v1/deprecated_builtins.rego:7: rego_type_error: deprecated built-in function calls in expression: cast_boolean
-testfiles/rego_v1/deprecated_builtins.rego:8: rego_type_error: deprecated built-in function calls in expression: cast_null
-testfiles/rego_v1/deprecated_builtins.rego:9: rego_type_error: deprecated built-in function calls in expression: cast_object
-testfiles/rego_v1/deprecated_builtins.rego:10: rego_type_error: deprecated built-in function calls in expression: cast_set
-testfiles/rego_v1/deprecated_builtins.rego:11: rego_type_error: deprecated built-in function calls in expression: cast_string
-testfiles/rego_v1/deprecated_builtins.rego:12: rego_type_error: deprecated built-in function calls in expression: net.cidr_overlap
-testfiles/rego_v1/deprecated_builtins.rego:13: rego_type_error: deprecated built-in function calls in expression: re_match
-testfiles/rego_v1/deprecated_builtins.rego:14: rego_type_error: deprecated built-in function calls in expression: set_diff
-testfiles/rego_v1/deprecated_builtins.rego:17: rego_type_error: deprecated built-in function calls in expression: any
-testfiles/rego_v1/deprecated_builtins.rego:19: rego_type_error: deprecated built-in function calls in expression: any
-testfiles/rego_v1/deprecated_builtins.rego:21: rego_type_error: deprecated built-in function calls in expression: any
-testfiles/rego_v1/deprecated_builtins.rego:21: rego_type_error: deprecated built-in function calls in expression: all

format/testfiles/rego_v1/duplicate_imports.rego.error

@@ -1,3 +0,0 @@
-2 errors occurred:
-testfiles/rego_v1/duplicate_imports.rego:4: rego_compile_error: import must not shadow import data.foo
-testfiles/rego_v1/duplicate_imports.rego:5: rego_compile_error: import must not shadow import data.foo

format/testfiles/rego_v1/keywords.rego.error

@@ -1,5 +0,0 @@
-4 errors occurred:
-testfiles/rego_v1/keywords.rego:3: rego_parse_error: if keyword cannot be used for rule name
-testfiles/rego_v1/keywords.rego:5: rego_parse_error: contains keyword cannot be used for rule name
-testfiles/rego_v1/keywords.rego:7: rego_parse_error: in keyword cannot be used for rule name
-testfiles/rego_v1/keywords.rego:9: rego_parse_error: every keyword cannot be used for rule name

format/testfiles/rego_v1/shadowing.rego.error

@@ -1,13 +0,0 @@
-12 errors occurred:
-testfiles/rego_v1/shadowing.rego:3: rego_compile_error: rules must not shadow input (use a different rule name)
-testfiles/rego_v1/shadowing.rego:5: rego_compile_error: rules must not shadow input (use a different rule name)
-testfiles/rego_v1/shadowing.rego:7: rego_compile_error: rules must not shadow input (use a different rule name)
-testfiles/rego_v1/shadowing.rego:9: rego_compile_error: rules must not shadow data (use a different rule name)
-testfiles/rego_v1/shadowing.rego:11: rego_compile_error: rules must not shadow data (use a different rule name)
-testfiles/rego_v1/shadowing.rego:13: rego_compile_error: rules must not shadow data (use a different rule name)
-testfiles/rego_v1/shadowing.rego:28: rego_compile_error: args must not shadow input (use a different variable name)
-testfiles/rego_v1/shadowing.rego:32: rego_compile_error: args must not shadow data (use a different variable name)
-testfiles/rego_v1/shadowing.rego:16: rego_compile_error: variables must not shadow input (use a different variable name)
-testfiles/rego_v1/shadowing.rego:17: rego_compile_error: variables must not shadow data (use a different variable name)
-testfiles/rego_v1/shadowing.rego:21: rego_compile_error: variables must not shadow input (use a different variable name)
-testfiles/rego_v1/shadowing.rego:25: rego_compile_error: variables must not shadow data (use a different variable name)

format/testfiles/v0/test_contains.rego.formatted

@@ -1,4 +1,4 @@
-package test.contains
+package test["contains"]
 
 import future.keywords.contains
 

format/testfiles/v0/test_if.rego.formatted

@@ -1,4 +1,4 @@
-package test.if
+package test["if"]
 
 import future.keywords.if
 

format/testfiles/v0/test_if_else.rego.formatted

@@ -1,4 +1,4 @@
-package test.if
+package test["if"]
 
 import future.keywords.if
 

format/testfiles/v0/test_in.rego

@@ -0,0 +1,13 @@
+package test.in
+
+import future.keywords.in
+
+a["in"] := "foo"
+
+b.c["in"] := "bar"
+
+c["in"].d := "baz"
+
+p {
+	input["in"]
+}

format/testfiles/v0/test_in.rego.formatted

@@ -0,0 +1,13 @@
+package test["in"]
+
+import future.keywords.in
+
+a["in"] := "foo"
+
+b.c["in"] := "bar"
+
+c["in"].d := "baz"
+
+p {
+	input["in"]
+}

format/testfiles/v0_to_v1/deprecated_builtins.rego.error

@@ -0,0 +1,16 @@
+15 errors occurred:
+testfiles/v0_to_v1/deprecated_builtins.rego:4: rego_type_error: deprecated built-in function calls in expression: any
+testfiles/v0_to_v1/deprecated_builtins.rego:5: rego_type_error: deprecated built-in function calls in expression: all
+testfiles/v0_to_v1/deprecated_builtins.rego:6: rego_type_error: deprecated built-in function calls in expression: cast_array
+testfiles/v0_to_v1/deprecated_builtins.rego:7: rego_type_error: deprecated built-in function calls in expression: cast_boolean
+testfiles/v0_to_v1/deprecated_builtins.rego:8: rego_type_error: deprecated built-in function calls in expression: cast_null
+testfiles/v0_to_v1/deprecated_builtins.rego:9: rego_type_error: deprecated built-in function calls in expression: cast_object
+testfiles/v0_to_v1/deprecated_builtins.rego:10: rego_type_error: deprecated built-in function calls in expression: cast_set
+testfiles/v0_to_v1/deprecated_builtins.rego:11: rego_type_error: deprecated built-in function calls in expression: cast_string
+testfiles/v0_to_v1/deprecated_builtins.rego:12: rego_type_error: deprecated built-in function calls in expression: net.cidr_overlap
+testfiles/v0_to_v1/deprecated_builtins.rego:13: rego_type_error: deprecated built-in function calls in expression: re_match
+testfiles/v0_to_v1/deprecated_builtins.rego:14: rego_type_error: deprecated built-in function calls in expression: set_diff
+testfiles/v0_to_v1/deprecated_builtins.rego:17: rego_type_error: deprecated built-in function calls in expression: any
+testfiles/v0_to_v1/deprecated_builtins.rego:19: rego_type_error: deprecated built-in function calls in expression: any
+testfiles/v0_to_v1/deprecated_builtins.rego:21: rego_type_error: deprecated built-in function calls in expression: any
+testfiles/v0_to_v1/deprecated_builtins.rego:21: rego_type_error: deprecated built-in function calls in expression: all

format/testfiles/v0_to_v1/duplicate_imports.rego.error

@@ -0,0 +1,3 @@
+2 errors occurred:
+testfiles/v0_to_v1/duplicate_imports.rego:4: rego_compile_error: import must not shadow import data.foo
+testfiles/v0_to_v1/duplicate_imports.rego:5: rego_compile_error: import must not shadow import data.foo

format/testfiles/v0_to_v1/keyword_errors.rego

@@ -7,3 +7,7 @@ contains := 2
 in := 3
 
 every := 4
+
+p {
+	data.foo.contains.bar == 42
+}

format/testfiles/v0_to_v1/keyword_errors.rego.error

@@ -0,0 +1,5 @@
+4 errors occurred:
+testfiles/v0_to_v1/keyword_errors.rego:3: rego_parse_error: if keyword cannot be used for rule name
+testfiles/v0_to_v1/keyword_errors.rego:5: rego_parse_error: contains keyword cannot be used for rule name
+testfiles/v0_to_v1/keyword_errors.rego:7: rego_parse_error: in keyword cannot be used for rule name
+testfiles/v0_to_v1/keyword_errors.rego:9: rego_parse_error: every keyword cannot be used for rule name

format/testfiles/v0_to_v1/keywords.rego

@@ -0,0 +1,5 @@
+package test.if.contains.in.every
+
+p {
+	data.if.contains.in.every == 42
+}

format/testfiles/v0_to_v1/keywords.rego.formatted

@@ -0,0 +1,7 @@
+package test["if"]["contains"]["in"]["every"]
+
+import rego.v1
+
+p if {
+	data["if"]["contains"]["in"]["every"] == 42
+}

format/testfiles/v0_to_v1/shadowing.rego.error

@@ -0,0 +1,13 @@
+12 errors occurred:
+testfiles/v0_to_v1/shadowing.rego:3: rego_compile_error: rules must not shadow input (use a different rule name)
+testfiles/v0_to_v1/shadowing.rego:5: rego_compile_error: rules must not shadow input (use a different rule name)
+testfiles/v0_to_v1/shadowing.rego:7: rego_compile_error: rules must not shadow input (use a different rule name)
+testfiles/v0_to_v1/shadowing.rego:9: rego_compile_error: rules must not shadow data (use a different rule name)
+testfiles/v0_to_v1/shadowing.rego:11: rego_compile_error: rules must not shadow data (use a different rule name)
+testfiles/v0_to_v1/shadowing.rego:13: rego_compile_error: rules must not shadow data (use a different rule name)
+testfiles/v0_to_v1/shadowing.rego:28: rego_compile_error: args must not shadow input (use a different variable name)
+testfiles/v0_to_v1/shadowing.rego:32: rego_compile_error: args must not shadow data (use a different variable name)
+testfiles/v0_to_v1/shadowing.rego:16: rego_compile_error: variables must not shadow input (use a different variable name)
+testfiles/v0_to_v1/shadowing.rego:17: rego_compile_error: variables must not shadow data (use a different variable name)
+testfiles/v0_to_v1/shadowing.rego:21: rego_compile_error: variables must not shadow input (use a different variable name)
+testfiles/v0_to_v1/shadowing.rego:25: rego_compile_error: variables must not shadow data (use a different variable name)

format/testfiles/v1/test.rego

@@ -0,0 +1,222 @@
+# The blank lines below me should be gone! (except one)
+
+
+# Comment!
+package a.b
+
+# I also, am a comment.
+import data.x.y.z
+import data.a.b.c              # Another comment!
+# I belong with data.a, there should be a newline before me.
+import data.a
+import data.f.g
+
+default foo = false
+foo contains x if {
+not x = g
+    f(x) = 1
+    g(
+        x, "foo"
+    ) = z
+}
+
+globals = {"foo": "bar",
+"fizz": "buzz"}
+
+partial_obj["x"] = 1
+partial_obj.y = 2
+
+partial_obj["z"] = 3 if {
+    true
+}
+
+partial_set contains "x"
+
+# Latent comment.
+
+r = y if {
+    y = x
+    split("foo.bar", ".", input.x) with input as {"x": x}
+    }
+
+# Comment on else
+        else = y if {
+    y = ["howdy"]
+    x = {"x": {
+    "y": "z",
+    }}
+    a = {"a": {
+    "b": "c",
+    }, "b": "c", "c": [1, 2,
+    3, 4]}
+    }
+
+fn(x) = y if {
+		y = x
+}
+
+fn_else(x) = 1 if {
+    true
+} # foo
+else  =
+# bar
+2
+if {
+    true
+} else = 3
+# baz
+if {
+    false
+}
+
+long(x) = true if {
+    x = "foo %host"
+    }
+
+    short(x) if {
+    x = "bar"
+    }
+
+raw_string = `hi\there`
+raw_multiline = `this
+string
+  is on
+multiple lines`
+
+fn2([x, y,
+z], {"foo": a}) = b if {
+split(x, y, c)
+trim(a, z, d) # function comment 1
+split(c[0], d, b)
+x = sprintf("hello %v",
+["world"])
+#function comment 2
+} # function comment 3
+
+f contains x if {
+    x = "hi"
+} { # Comment on chain
+    x = "bye"
+}
+
+        import data.foo.bar
+    import data.bar.foo # data.bar.foo should be first
+
+p[x] = y if { y = x
+                y = "foo"
+                z = { "a": "b", # Comment inside object 1
+                    "b":    "c"   ,   "c": "d",   # comment on object entry line
+                    # Comment inside object 2
+"d": "e",
+# Comment before closing object brace.
+} # Comment on closing object brace.
+a = {"a": "b", "c": "d"}
+b = [1, 2, 3, 4]
+c = [1, 2,
+# Comment inside array
+3, 4,
+5, 6, 7,
+8,
+# Comment before nested composite.
+[
+["foo"], # Comment inside nested composite.
+["bar"], # Comment after last element in nested composite.
+# Comment before nested composite closing bracket.
+], # Comment on nested composite closing bracket.
+# Comment before closing array bracket.
+] # Comment on closing array bracket.
+
+d = [1 | b[_]]
+e = [1 | split("foo.bar", ".", x); x[_]]
+f = [1 | split("foo.bar", ".", x)
+x[_]]
+g = [1 |
+split("foo.bar", ".", x) # comment in array comprehension
+x[_]
+# inner comment
+]
+
+h = {1 | b[_]}
+i = {1 | split("foo.bar", ".", x); x[_]}
+j = {1 | split("foo.bar", ".", x)
+x[_]}
+k = {1 |
+split("foo.bar", ".", x) # comment in set comprehension
+x[_]
+
+# inner comment
+}
+
+l = {"foo":1 | b[_]}
+m = {y:
+x | split("foo.bar", ".", x); y = x[_]}
+n = {y:  x | split("foo.bar", ".", x)
+y = x[_]}
+o = {y: x |
+split("foo.bar", ".", x) # comment in object comprehension
+y = x[_]
+
+# inner comment
+}
+} # Comment on rule closing brace
+
+nested_infix if {
+    x + 1
+    x = y + 2
+    plus(x, 1, 2)
+    plus(x, 1)
+    y = f(x)
+    f(x, y)
+    y = x + 1 + 2
+    x = y + # comment
+        z
+    x = (a + b) / 2
+    f((a+b)/2)
+    y = q()
+}
+
+expanded_const = true
+
+partial_obj["why"] = true if { false }
+
+empty_sets if {
+    set()
+    set() # comment at end of set
+}
+
+vardecls if {
+    some v1,  v2,# c1
+        v3,v4, # c2
+        v5 # c3
+}
+
+declare1 := 1
+
+declare2 := 2 if { false }
+
+declare3 := {1,2,3}
+declare4 := {4,5,6}
+
+union_object := {"response": (declare3|declare4)}
+
+union_set := {(declare3|declare4)}
+
+union_list := [(declare3|declare4)]
+
+union_set_2 := {(((declare3|declare4)))}
+
+union_object_multi_line := {"response": (
+  declare3 | declare4
+)}
+
+union_object_key := {(declare3|declare4): "foo"}
+
+union_object_key_multi_line := {(declare3|
+declare4):
+"foo"
+}
+
+# more comments!
+# more comments!
+# more comments!
+# more comments!

format/testfiles/v1/test.rego.error

@@ -0,0 +1,26 @@
+# The blank lines below me should be gone! (except one)
+
+
+# Comment!
+package a.b
+
+# I also, am a comment.
+import data.x.y.z
+import data.a.b.c              # Another comment!
+# I belong with data.a, there should be a newline before me.
+import data.a
+import data.f.g
+
+default foo = false
+foo[x] {
+not x = g
+}
+
+globals = {"foo": "bar",
+"fizz": "buzz"}
+
+# Latent comment.
+
+r = y {
+    y = x
+    split("foo.bar", ".", input.x) with input as {"x": x}

format/testfiles/v1/test.rego.formatted

@@ -0,0 +1,237 @@
+# The blank lines below me should be gone! (except one)
+
+# Comment!
+package a.b
+
+# I also, am a comment.
+import data.a.b.c # Another comment!
+import data.x.y.z
+
+# I belong with data.a, there should be a newline before me.
+import data.a
+import data.f.g
+
+default foo := false
+
+foo contains x if {
+	not x = g
+	f(x) = 1
+	g(x, "foo") = z
+}
+
+globals := {
+	"foo": "bar",
+	"fizz": "buzz",
+}
+
+partial_obj["x"] := 1
+
+partial_obj["y"] := 2
+
+partial_obj["z"] := 3
+
+partial_set contains "x"
+
+# Latent comment.
+
+r := y if {
+	y = x
+	split("foo.bar", ".", input.x) with input as {"x": x}
+}
+
+# Comment on else
+else := y if {
+	y = ["howdy"]
+	x = {"x": {"y": "z"}}
+	a = {
+		"a": {"b": "c"},
+		"b": "c", "c": [
+			1, 2,
+			3, 4,
+		],
+	}
+}
+
+fn(x) := y if {
+	y = x
+}
+
+fn_else(x) := 1 if {
+	true
+} # foo
+
+else := 2 if {
+	# bar
+
+	true
+} else := 3 if {
+	# baz
+
+	false
+}
+
+long(x) if {
+	x = "foo %host"
+}
+
+short(x) if {
+	x = "bar"
+}
+
+raw_string := `hi\there`
+
+raw_multiline := `this
+string
+  is on
+multiple lines`
+
+fn2(
+	[
+		x, y,
+		z,
+	],
+	{"foo": a},
+) := b if {
+	split(x, y, c)
+	trim(a, z, d) # function comment 1
+	split(c[0], d, b)
+	x = sprintf(
+		"hello %v",
+		["world"],
+	)
+	#function comment 2
+} # function comment 3
+
+f contains x if {
+	x = "hi"
+}
+
+f contains x if { # Comment on chain
+	x = "bye"
+}
+
+import data.bar.foo # data.bar.foo should be first
+import data.foo.bar
+
+p[x] := y if {
+	y = x
+	y = "foo"
+	z = {
+		"a": "b", # Comment inside object 1
+		"b": "c", "c": "d", # comment on object entry line
+		# Comment inside object 2
+		"d": "e",
+		# Comment before closing object brace.
+	} # Comment on closing object brace.
+	a = {"a": "b", "c": "d"}
+	b = [1, 2, 3, 4]
+	c = [
+		1, 2,
+		# Comment inside array
+		3, 4,
+		5, 6, 7,
+		8,
+		# Comment before nested composite.
+		[
+			["foo"], # Comment inside nested composite.
+			["bar"], # Comment after last element in nested composite.
+			# Comment before nested composite closing bracket.
+		], # Comment on nested composite closing bracket.
+		# Comment before closing array bracket.
+	] # Comment on closing array bracket.
+
+	d = [1 | b[_]]
+	e = [1 | split("foo.bar", ".", x); x[_]]
+	f = [1 |
+		split("foo.bar", ".", x)
+		x[_]
+	]
+	g = [1 |
+		split("foo.bar", ".", x) # comment in array comprehension
+		x[_]
+		# inner comment
+	]
+
+	h = {1 | b[_]}
+	i = {1 | split("foo.bar", ".", x); x[_]}
+	j = {1 |
+		split("foo.bar", ".", x)
+		x[_]
+	}
+	k = {1 |
+		split("foo.bar", ".", x) # comment in set comprehension
+		x[_]
+		# inner comment
+	}
+
+	l = {"foo": 1 | b[_]}
+	m = {y: x |
+		split("foo.bar", ".", x)
+		y = x[_]
+	}
+	n = {y: x |
+		split("foo.bar", ".", x)
+		y = x[_]
+	}
+	o = {y: x |
+		split("foo.bar", ".", x) # comment in object comprehension
+		y = x[_]
+		# inner comment
+	}
+} # Comment on rule closing brace
+
+nested_infix if {
+	x + 1
+	x = y + 2
+	2 = x + 1
+	x + 1
+	y = f(x)
+	f(x, y)
+	y = (x + 1) + 2
+	x = y + z # comment
+	x = (a + b) / 2
+	f((a + b) / 2)
+	y = q()
+}
+
+expanded_const := true
+
+partial_obj["why"] if false
+
+empty_sets if {
+	set()
+	set() # comment at end of set
+}
+
+vardecls if {
+	some v1, v2, # c1
+		v3, v4, # c2
+		v5 # c3
+}
+
+declare1 := 1
+
+declare2 := 2 if false
+
+declare3 := {1, 2, 3}
+
+declare4 := {4, 5, 6}
+
+union_object := {"response": (declare3 | declare4)}
+
+union_set := {(declare3 | declare4)}
+
+union_list := [(declare3 | declare4)]
+
+union_set_2 := {(declare3 | declare4)}
+
+union_object_multi_line := {"response": (declare3 | declare4)}
+
+union_object_key := {(declare3 | declare4): "foo"}
+
+union_object_key_multi_line := {(declare3 | declare4): "foo"}
+
+# more comments!
+# more comments!
+# more comments!
+# more comments!

format/testfiles/v1/test_assignments.rego

@@ -0,0 +1,24 @@
+package assignments
+
+# default value assignment
+default a := 1
+
+# rule
+b := 2
+
+# else keyword
+c := 3 if {
+	false
+} else := 4 if {
+	true
+}
+
+# partial rule
+d[msg] := 5 if {
+	msg = [1, 2, 3][_]
+}
+
+# function return value
+e := f(6)
+
+f(x) := x

format/testfiles/v1/test_assignments.rego.formatted

@@ -0,0 +1,22 @@
+package assignments
+
+# default value assignment
+default a := 1
+
+# rule
+b := 2
+
+# else keyword
+c := 3 if {
+	false
+} else := 4
+
+# partial rule
+d[msg] := 5 if {
+	msg = [1, 2, 3][_]
+}
+
+# function return value
+e := f(6)
+
+f(x) := x