releases/v1/format/format_test.go

// Copyright 2017 The OPA Authors.  All rights reserved.
// Use of this source code is governed by an Apache2
// license that can be found in the LICENSE file.

package format

import (
	"bytes"
	"fmt"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"testing"

	"github.com/open-policy-agent/opa/v1/ast"
	"github.com/open-policy-agent/opa/v1/ast/location"
)

func TestFormatNilLocation(t *testing.T) {
	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"
`,
		},
	}

	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

			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))
			}
		})
	}
}

func TestFormatNilLocationEmptyBody(t *testing.T) {
	b := ast.NewBody()
	x, err := Ast(b)
	if len(x) != 0 || err != nil {
		t.Fatalf("Expected empty result but got: %q, err: %v", string(x), err)
	}
}

func TestFormatNilLocationFunctionArgs(t *testing.T) {
	b := ast.NewBody()
	s := ast.StringTerm(" ")
	s.SetLocation(location.NewLocation([]byte("\" \""), "p.rego", 2, 2))
	b.Append(ast.Split.Expr(ast.NewTerm(ast.Var("__local1__")), s, ast.NewTerm(ast.Var("__local2__"))))
	exp := "split(__local1__, \" \", __local2__)\n"
	bs, err := Ast(b)
	if err != nil {
		t.Fatal(err)
	}
	if string(bs) != exp {
		t.Fatalf("Expected %q but got %q", exp, string(bs))
	}
}

func TestFormatSourceError(t *testing.T) {
	rego := "testfiles/v0/test.rego.error"
	contents, err := os.ReadFile(rego)
	if err != nil {
		t.Fatalf("Failed to read rego source: %v", err)
	}

	_, err = Source(rego, contents)
	if err == nil {
		t.Fatal("Expected parsing error, not nil")
	}

	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 TestFormatV0Source(t *testing.T) {
	regoFiles, err := filepath.Glob("testfiles/v0/*.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.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)
			}

			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))
			}

			formattedModule, err := ast.ParseModuleWithOpts(rego+".tmp", string(formatted), popts)
			if err != nil {
				t.Fatalf("Failed to parse formatted bytes: %v", err)
			}

			if expASTPreserved(rego) {
				module := ast.MustParseModuleWithOpts(string(contents), popts)
				if !module.Equal(formattedModule) {
					t.Fatalf("Expected formatting to preserve the AST, but got:\n\n%v\n\nexpected:\n\n%v", formattedModule, module)
				}
			}

			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 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))
			}

			formattedModule, err := ast.ParseModuleWithOpts(rego+".tmp", string(formatted), popts)
			if err != nil {
				t.Fatalf("Failed to parse formatted bytes: %v", err)
			}

			if expASTPreserved(rego) {
				module := ast.MustParseModuleWithOpts(string(contents), popts)
				if !module.Equal(formattedModule) {
					t.Fatalf("Expected formatting to preserve the AST, but got:\n\n%v\n\nexpected:\n\n%v", formattedModule, module)
				}
			}

			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))
			}

		})
	}
}

// astAlteringTestFiles are the test files where formatting is expected to
// change the AST. For all other test files, formatting must preserve it.
var astAlteringTestFiles = map[string]bool{
	// Imports are sorted.
	"testfiles/v0/test.rego": true,
	"testfiles/v0/test_in_operator_with_all_keywords_import.rego": true,
	"testfiles/v0/test_keywords_in_refs.rego":                     true,
	"testfiles/v0/test_keywords_in_refs_keep_brackets.rego":       true,

	// The `future.keywords.in` import is added.
	"testfiles/v0/test_in_operator_without_import.rego": true,

	// `if` is dropped from the rule head, making `q[1] = y` parse as a partial
	// object rule rather than a ref head rule.
	"testfiles/v0/test_ref_heads.rego": true,

	// Rule head `=` is rewritten to `:=`.
	"testfiles/v1/test_contains.rego":                          true,
	"testfiles/v1/test_contains_if.rego":                       true,
	"testfiles/v1/test_every.rego":                             true,
	"testfiles/v1/test_functions.rego":                         true,
	"testfiles/v1/test_future_kw_import.rego":                  true,
	"testfiles/v1/test_if.rego":                                true,
	"testfiles/v1/test_in_operator_with_parenthesis.rego":      true,
	"testfiles/v1/test_issue_2299.rego":                        true,
	"testfiles/v1/test_issue_3836.rego":                        true,
	"testfiles/v1/test_issue_5449.rego":                        true,
	"testfiles/v1/test_issue_5449_with_contains_ref_rule.rego": true,
	"testfiles/v1/test_issue_5449_with_ref_rule.rego":          true,
	"testfiles/v1/test_issue_5798.rego":                        true,
	"testfiles/v1/test_ref_heads.rego":                         true,
	"testfiles/v1/test_unicode.rego":                           true,
	"testfiles/v1/test_with.rego":                              true,

	// Imports are sorted.
	"testfiles/v1/test_keywords_in_refs.rego":               true,
	"testfiles/v1/test_keywords_in_refs_keep_brackets.rego": true,

	// The `rego.v1` import is dropped.
	"testfiles/v1/test_rego_v1.rego": true,

	// Constant template string expressions are folded into string parts.
	"testfiles/v1/test_template_strings.rego": true,

	// Rule head `=` is rewritten to `:=`, imports are sorted, and calls in
	// expression position are rewritten to infix comparisons.
	"testfiles/v1/test.rego": true,
}

func expASTPreserved(rego string) bool {
	return !astAlteringTestFiles[filepath.ToSlash(rego)]
}

func TestFormatV0SourceToRegoV1(t *testing.T) {
	regoFiles, err := filepath.Glob("testfiles/v0_to_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)
			}

			errorExpected := false
			expected, err := os.ReadFile(rego + ".formatted")
			if err != nil {
				if os.IsNotExist(err) {
					errorExpected = true
					expected, err = os.ReadFile(rego + ".error")
					if err != nil {
						t.Fatalf("Failed to read expected error source: %v", err)
					}
				}
				if !errorExpected {
					t.Fatalf("Failed to read expected rego source: %v", err)
				}
			}

			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, sourceOpts)

				if err == nil {
					t.Fatalf("Expected error, got: %s", formatted)
				}
				if err.Error() != string(expected) {
					t.Fatalf("Expected error:\n\n'%s'\n\ngot:\n\n'%s'", expected, err.Error())
				}
			} else {
				formatted, err := SourceWithOpts(rego, contents, sourceOpts)

				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.ParseModule(rego+".tmp", string(formatted)); err != nil {
					t.Fatalf("Failed to parse formatted bytes: %v", err)
				}

				formatted, err = SourceWithOpts(rego, formatted, targetOpts)
				if err != nil {
					t.Fatalf("Failed to double format file: %s", err)
				}

				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))
				}

				// rego-v1 formatted code is still compliant with v0, and should not be changed if formatted as such
				formatted, err = SourceWithOpts(rego, formatted, targetOpts)
				if err != nil {
					t.Fatalf("Failed to double format file as v0: %s", err)
				}

				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 TestFormatAST(t *testing.T) {
	cases := []struct {
		note        string
		regoVersion ast.RegoVersion
		toFmt       any
		expected    string
	}{
		{
			note:     "var",
			toFmt:    ast.Var(`foo`),
			expected: "foo",
		},
		{
			note: "string",
			toFmt: &ast.Term{
				Value:    ast.String("foo"),
				Location: &ast.Location{Text: []byte(`"foo"`)},
			},
			expected: `"foo"`,
		},
		{
			note: "string, raw",
			toFmt: &ast.Term{
				Value:    ast.String(`foo`),
				Location: &ast.Location{Text: []byte("`foo`")},
			},
			expected: "`foo`",
		},
		{
			note:     "var wildcard",
			toFmt:    ast.Var(`$12`),
			expected: "_",
		},
		{
			note: "string with wildcard prefix",
			toFmt: &ast.Term{
				Value:    ast.String("$01"),
				Location: &ast.Location{Text: []byte(`"$01"`)},
			},
			expected: `"$01"`,
		},
		{
			note:     "ref var only",
			toFmt:    ast.MustParseRef(`data.foo`),
			expected: "data.foo",
		},
		{
			note:     "ref multi vars",
			toFmt:    ast.MustParseRef(`data.foo.bar.baz`),
			expected: "data.foo.bar.baz",
		},
		{
			note:     "ref with string",
			toFmt:    ast.MustParseRef(`data["foo"]`),
			expected: `data.foo`,
		},
		{
			note:     "ref multi string",
			toFmt:    ast.MustParseRef(`data["foo"]["bar"]["baz"]`),
			expected: `data.foo.bar.baz`,
		},
		{
			note:     "ref with string needs brackets",
			toFmt:    ast.MustParseRef(`data["foo my-var\nbar"]`),
			expected: `data["foo my-var\nbar"]`,
		},
		{
			note:     "ref multi string needs brackets",
			toFmt:    ast.MustParseRef(`data["foo my-var"]["bar"]["almost.baz"]`),
			expected: `data["foo my-var"].bar["almost.baz"]`,
		},
		{
			note:     "ref var wildcard",
			toFmt:    ast.MustParseRef(`data.foo[_]`),
			expected: "data.foo[_]",
		},
		{
			note:     "ref var wildcard",
			toFmt:    ast.MustParseRef(`foo[_]`),
			expected: "foo[_]",
		},
		{
			note:     "ref string with wildcard prefix",
			toFmt:    ast.MustParseRef(`foo["$01"]`),
			expected: `foo["$01"]`,
		},
		{
			note:     "nested ref var wildcard",
			toFmt:    ast.MustParseRef(`foo[bar[baz[_]]]`),
			expected: "foo[bar[baz[_]]]",
		},
		{
			note:     "ref mixed",
			toFmt:    ast.MustParseRef(`foo["bar"].baz[_]["bar-2"].qux`),
			expected: `foo.bar.baz[_]["bar-2"].qux`,
		},
		{
			note:     "ref empty",
			toFmt:    ast.Ref{},
			expected: ``,
		},
		{
			note:     "ref nil",
			toFmt:    ast.Ref(nil),
			expected: ``,
		},
		{
			note:     "ref operator",
			toFmt:    ast.MustParseRef(`foo[count(foo) - 1]`),
			expected: `foo[count(foo) - 1]`,
		},
		{
			note:     "x in xs",
			toFmt:    ast.Member.Call(ast.VarTerm("x"), ast.VarTerm("xs")),
			expected: `x in xs`,
		},
		{
			note:     "x, y in xs",
			toFmt:    ast.MemberWithKey.Call(ast.VarTerm("x"), ast.VarTerm("y"), ast.VarTerm("xs")),
			expected: `(x, y in xs)`,
		},
		{
			note: "some x in xs",
			toFmt: ast.NewExpr(&ast.SomeDecl{Symbols: []*ast.Term{
				ast.Member.Call(ast.VarTerm("x"), ast.VarTerm("xs")),
			}}),
			expected: `some x in xs`,
		},
		{
			note: "some x, y in xs",
			toFmt: ast.NewExpr(&ast.SomeDecl{Symbols: []*ast.Term{
				ast.MemberWithKey.Call(ast.VarTerm("x"), ast.VarTerm("y"), ast.VarTerm("xs")),
			}}),
			expected: `some x, y in xs`,
		},
		{
			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{
					RegoVersion:    ast.RegoV0,
					FutureKeywords: []string{"every"},
				}),
			expected: `package test

import future.keywords.every

p {
	every k, v in [1, 2] { k != v }
}`,
		},
		{
			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"},
					RegoVersion:    ast.RegoV0,
				}),
			expected: `package test

import future.keywords

p if {
	every k, v in [1, 2] { k != v }
}`,
		},
		{
			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"},
					RegoVersion:    ast.RegoV0,
				}),
			expected: `package test

import future.keywords

p if {
	every k, v in [1, 2] { k != v }
}`,
		},
		{
			note: "body shared wildcard",
			toFmt: ast.Body{
				&ast.Expr{
					Index: 0,
					Terms: []*ast.Term{
						ast.RefTerm(ast.VarTerm("eq")),
						ast.RefTerm(ast.VarTerm("input"), ast.StringTerm("arr"), ast.VarTerm("$01"), ast.StringTerm("some key"), ast.VarTerm("$02")),
						ast.VarTerm("bar"),
					},
				},
				&ast.Expr{
					Index: 1,
					Location: &ast.Location{
						Row: 2,
						Col: 1,
					},
					Terms: []*ast.Term{
						ast.RefTerm(ast.VarTerm("eq")),
						ast.RefTerm(ast.VarTerm("input"), ast.StringTerm("arr"), ast.VarTerm("$01"), ast.StringTerm("bar")),
						ast.VarTerm("qux"),
					},
				},
				&ast.Expr{
					Index: 1,
					Location: &ast.Location{
						Row: 2,
						Col: 1,
					},
					Terms: []*ast.Term{
						ast.RefTerm(ast.VarTerm("eq")),
						ast.RefTerm(ast.VarTerm("foo"), ast.VarTerm("$03"), ast.VarTerm("$01"), ast.StringTerm("bar")),
						ast.RefTerm(ast.VarTerm("bar"), ast.VarTerm("$03"), ast.VarTerm("$04"), ast.VarTerm("$01"), ast.StringTerm("bar")),
					},
				},
			},
			expected: `input.arr[_01]["some key"][_] = bar
input.arr[_01].bar = qux
foo[_03][_01].bar = bar[_03][_][_01].bar
`,
		},
		{
			note: "body shared wildcard - ref head",
			toFmt: ast.Body{
				&ast.Expr{
					Index: 0,
					Terms: ast.VarTerm("$x"),
				},
				&ast.Expr{
					Index: 1,
					Terms: ast.RefTerm(ast.VarTerm("$x"), ast.VarTerm("y")),
				},
			},
			expected: `_x
_x[y]`,
		},
		{
			note: "body shared wildcard - nested ref",
			toFmt: ast.Body{
				&ast.Expr{
					Index: 0,
					Terms: ast.VarTerm("$x"),
				},
				&ast.Expr{
					Index: 1,
					Terms: ast.RefTerm(ast.VarTerm("a"), ast.RefTerm(ast.VarTerm("$x"), ast.VarTerm("y"))),
				},
			},
			expected: `_x
a[_x[y]]`,
		},
		{
			note: "body shared wildcard - nested ref array",
			toFmt: ast.Body{
				&ast.Expr{
					Index: 0,
					Terms: ast.VarTerm("$x"),
				},
				&ast.Expr{
					Index: 1,
					Terms: ast.RefTerm(ast.VarTerm("a"), ast.RefTerm(ast.VarTerm("$x"), ast.VarTerm("y"), ast.ArrayTerm(ast.VarTerm("z"), ast.VarTerm("w")))),
				},
			},
			expected: `_x
a[_x[y][[z, w]]]`,
		},
		{
			note: "expr with wildcard that has a default location",
			toFmt: func() *ast.Expr {
				expr := ast.MustParseExpr(`["foo", _] = split(input.foo, ":")`)
				ast.WalkTerms(expr, func(term *ast.Term) bool {
					v, ok := term.Value.(ast.Var)
					if ok && v.IsWildcard() {
						term.Location = defaultLocation(term)
						return true
					}
					term.Location.File = "foo.rego"
					term.Location.Row = 2
					return false
				})
				return expr
			}(),
			expected: `["foo", _] = split(input.foo, ":")`,
		},
		{
			note: "expr all terms having empty-file locations",
			toFmt: ast.MustParseExpr(`[
					"foo",
					_
					] = split(input.foo, ":")`),
			expected: `
[
	"foo",
	_,
] = split(input.foo, ":")`,
		},
		{
			note: "expr where all terms having empty-file locations, and one is a default location",
			toFmt: func() *ast.Expr {
				expr := ast.MustParseExpr(`
["foo", __local1__] = split(input.foo, ":")`)
				ast.WalkTerms(expr, func(term *ast.Term) bool {
					if ast.VarTerm("__local1__").Equal(term) {
						term.Location = defaultLocation(term)
						return true
					}
					return false
				})
				return expr
			}(),
			expected: `["foo", __local1__] = split(input.foo, ":")`,
		},
		{
			note: "expr where generated var has an AST location not matching its source location",
			toFmt: func() *ast.Expr {
				e := ast.MustParseExpr(`__local0__ = concat(",", [__local1__])`)
				ast.WalkTerms(e, func(t *ast.Term) bool {
					t.Location.File = "t.rego"
					return false
				})
				// mangling that may happen in PE
				return ast.Concat.Expr(
					e.Operand(1).Value.(ast.Call)[1],
					e.Operand(1).Value.(ast.Call)[2],
					e.Operand(0),
				).SetLocation(e.Location)
			}(),
			expected: `concat(",", [__local1__], __local0__)`,
		},
		{
			note: "template-string",
			toFmt: ast.TemplateStringTerm(false,
				ast.StringTerm("foo "),
				&ast.Expr{
					Terms: ast.VarTerm("x"),
				},
				ast.StringTerm(" "),
				&ast.Expr{
					Terms: ast.RefTerm(ast.VarTerm("input"), ast.StringTerm("y")),
				},
				ast.StringTerm(" \n "),
				&ast.Expr{
					Terms: ast.Round.Call(ast.NumberTerm("1.2")),
				},
			),
			expected: `$"foo {x} {input.y} \n {round(1.2)}"`,
		},
		{
			note: "template-string, multi-line",
			toFmt: ast.TemplateStringTerm(true,
				ast.StringTerm(`foo `),
				&ast.Expr{
					Terms: ast.VarTerm("x"),
				},
				ast.StringTerm(` `),
				&ast.Expr{
					Terms: ast.RefTerm(ast.VarTerm("input"), ast.StringTerm("y")),
				},
				ast.StringTerm(` 
 `),
				&ast.Expr{
					Terms: ast.Round.Call(ast.NumberTerm("1.2")),
				},
			),
			expected: "$`foo {x} {input.y} \n {round(1.2)}`",
		},
		{
			note: "template-string, single-line, string template-expressions",
			toFmt: ast.TemplateStringTerm(false,
				ast.StringTerm("foo "),
				&ast.Expr{
					Terms: &ast.Term{
						Value:    ast.String("bar"),
						Location: &ast.Location{Text: []byte(`"bar"`)},
					},
				},
				ast.StringTerm(" "),
				&ast.Expr{
					Terms: &ast.Term{
						Value:    ast.String(`baz`),
						Location: &ast.Location{Text: []byte("`baz`")},
					},
				},
			),
			expected: `$"foo {"bar"} {` + "`baz`" + `}"`,
		},
		{
			note: "template-string, multi-line, string template-expressions",
			toFmt: ast.TemplateStringTerm(true,
				ast.StringTerm("foo "),
				&ast.Expr{
					Terms: &ast.Term{
						Value:    ast.String("bar"),
						Location: &ast.Location{Text: []byte(`"bar"`)},
					},
				},
				ast.StringTerm(" "),
				&ast.Expr{
					Terms: &ast.Term{
						Value:    ast.String(`baz`),
						Location: &ast.Location{Text: []byte("`baz`")},
					},
				},
			),
			expected: "$`foo {\"bar\"} {`baz`}`",
		},
		{
			note: "template-string, left curly in string part",
			toFmt: ast.TemplateStringTerm(false,
				ast.StringTerm(`allow if { `),
				&ast.Expr{
					Terms: ast.VarTerm("x"),
				},
				ast.StringTerm(" }"),
			),
			expected: `$"allow if \{ {x} }"`,
		},
		{
			note: "template-string, multi-line, left curly in string part",
			toFmt: ast.TemplateStringTerm(true,
				ast.StringTerm("allow if {\n\t"),
				&ast.Expr{
					Terms: ast.VarTerm("x"),
				},
				ast.StringTerm("\n}"),
			),
			expected: "$`allow if \\{\n\t{x}\n}`",
		},
		{
			note:        "v0, not-body adds import if missing",
			regoVersion: ast.RegoV0,
			toFmt: ast.MustParseModuleWithOpts(`package test
				p if {
					not 1 + input.x == 2
				}`,
				ast.ParserOptions{
					FutureKeywords: []string{"not"},
				}),
			expected: `package test

import future.keywords.not

p {
	not 1 + input.x == 2
}`,
		},
		{
			note:        "v1, not-body adds import if missing",
			regoVersion: ast.RegoV1,
			toFmt: ast.MustParseModuleWithOpts(`package test
				p if {
					not 1 + input.x == 2
				}`,
				ast.ParserOptions{
					FutureKeywords: []string{"not"},
				}),
			expected: `package test

import future.keywords.not

p if {
	not 1 + input.x == 2
}`,
		},
		{
			note:        "v1, not-body explicit single expression round-trips",
			regoVersion: ast.RegoV1,
			toFmt: ast.MustParseModuleWithOpts(`package test
				import future.keywords.not
				p if {
					not {input.x == 1}
				}`,
				ast.ParserOptions{
					FutureKeywords: []string{"not"},
				}),
			expected: `package test

import future.keywords.not

p if {
	not { input.x == 1 }
}`,
		},
		{
			note:        "v1, not-body multi expression",
			regoVersion: ast.RegoV1,
			toFmt: ast.MustParseModuleWithOpts(`package test
				import future.keywords.not
				p if {
					not {
						input.x == 1
						input.y == 2
					}
				}`,
				ast.ParserOptions{
					FutureKeywords: []string{"not"},
				}),
			expected: `package test

import future.keywords.not

p if {
	not {
		input.x == 1
		input.y == 2
	}
}`,
		},
		{
			note:        "v1, not-body implicit single expression stays inline",
			regoVersion: ast.RegoV1,
			toFmt: ast.MustParseModuleWithOpts(`package test
				import future.keywords.not
				p if {
					not input.x == 1
				}`,
				ast.ParserOptions{
					FutureKeywords: []string{"not"},
				}),
			expected: `package test

import future.keywords.not

p if {
	not input.x == 1
}`,
		},
	}

	for _, tc := range cases {
		t.Run(tc.note, func(t *testing.T) {
			bs, err := AstWithOpts(tc.toFmt, Opts{
				RegoVersion: tc.regoVersion,
				ParserOptions: &ast.ParserOptions{
					RegoVersion: tc.regoVersion,
				},
			})
			if err != nil {
				t.Fatalf("Unexpected error: %s", err)
			}
			expected := strings.TrimSpace(tc.expected)
			actual := strings.TrimSpace(string(bs))
			if actual != expected {
				t.Fatalf("Expected:\n\n%q\n\nGot:\n\n%q\n\n", expected, actual)
			}
		})

		// consistency check: disregarding source locations, it shouldn't panic
		t.Run("no_loc/"+tc.note, func(t *testing.T) {
			if _, err := AstWithOpts(tc.toFmt, Opts{IgnoreLocations: true}); err != nil {
				t.Fatalf("Unexpected error: %s", err)
			}
		})
	}
}

func TestFormatAST_Error(t *testing.T) {
	cases := []struct {
		note        string
		regoVersion ast.RegoVersion
		toFmt       any
		expErr      string
	}{
		{
			note:   "package with only data term",
			toFmt:  &ast.Package{Path: ast.Ref{ast.DefaultRootDocument}},
			expErr: `rego_format_error: invalid package path: data`,
		},
		{
			note: "module with package with only data term",
			toFmt: &ast.Module{
				Package: &ast.Package{Path: ast.Ref{ast.DefaultRootDocument}},
			},
			expErr: `rego_format_error: invalid package path: data`,
		},
	}

	for _, tc := range cases {
		t.Run(tc.note, func(t *testing.T) {
			_, err := AstWithOpts(tc.toFmt, Opts{
				RegoVersion: tc.regoVersion,
				ParserOptions: &ast.ParserOptions{
					RegoVersion: tc.regoVersion,
				},
			})
			if err == nil {
				t.Fatalf("Expected error, got nil")
			}
			if !strings.Contains(err.Error(), tc.expErr) {
				t.Fatalf("Expected error to contain:\n\n%q\n\ngot:\n\n%q", tc.expErr, err.Error())
			}
		})
	}
}

func TestFormatDeepCopy(t *testing.T) {

	original := ast.Body{
		&ast.Expr{
			Index: 0,
			Terms: ast.VarTerm("$x"),
		},
		&ast.Expr{
			Index: 1,
			Terms: ast.RefTerm(ast.VarTerm("$x"), ast.VarTerm("y")),
		},
	}

	cpy := original.Copy()

	_, err := Ast(original)
	if err != nil {
		t.Fatal(err)
	}

	if !cpy.Equal(original) {
		t.Fatal("expected original to be unmodified")
	}

}

func differsAt(a, b []byte) (int, int) {
	if bytes.Equal(a, b) {
		return 0, 0
	}
	minLen := min(len(a), len(b))
	ln := 1
	for i := range minLen {
		if a[i] == '\n' {
			ln++
		}
		if a[i] != b[i] {
			return ln, i
		}
	}
	return ln, minLen - 1
}

func prefixWithLineNumbers(bs []byte) []byte {
	raw := string(bs)
	lines := strings.Split(raw, "\n")
	format := fmt.Sprintf("%%%dd %%s", len(strconv.Itoa(len(lines)+1)))
	for i, line := range lines {
		lines[i] = fmt.Sprintf(format, i+1, line)
	}
	return []byte(strings.Join(lines, "\n"))
}

func TestSource_DefaultRegoVersion(t *testing.T) {
	tests := []struct {
		note         string
		module       string
		expFormatted string
		expErrs      []string
	}{
		{
			note: "v0", // from default rego-version
			module: `package test

p[x]            {
	x = "a"
}`,

			expErrs: []string{
				"test.rego:3: rego_parse_error: `if` keyword is required before rule body",
				"test.rego:3: rego_parse_error: `contains` keyword is required for partial set rules",
			},
		},
		{
			note: "v1",
			module: `package test

p    contains    x    if      {
	x = "a"
}`,
			expFormatted: `package test

p contains x if {
	x = "a"
}
`,
		},
	}

	for _, tc := range tests {
		t.Run(tc.note, func(t *testing.T) {
			formatted, err := Source("test.rego", []byte(tc.module))
			if len(tc.expErrs) > 0 {
				if err == nil {
					t.Fatalf("expected errors but got nil")
				}

				for _, expErr := range tc.expErrs {
					if !strings.Contains(err.Error(), expErr) {
						t.Fatalf("expected error:\n\n%q\n\nbut got:\n\n%q", expErr, err)
					}
				}
			} else {
				if err != nil {
					t.Fatalf("unexpected error: %v", err)
				}

				formattedStr := string(formatted)
				if formattedStr != tc.expFormatted {
					t.Fatalf("expected %q but got %q", tc.expFormatted, formattedStr)
				}
			}
		})
	}
}

func TestSourceWithOpts_DefaultRegoVersion(t *testing.T) {
	tests := []struct {
		note          string
		toRegoVersion ast.RegoVersion
		module        string
		expFormatted  string
		expErrs       []string
	}{
		{
			note:          "v0 -> v0", // from default rego-version
			toRegoVersion: ast.RegoV0,
			module: `package test

p[x]            {
	x = "a"
}`,
			expErrs: []string{
				"test.rego:3: rego_parse_error: `if` keyword is required before rule body",
				"test.rego:3: rego_parse_error: `contains` keyword is required for partial set rules",
			},
		},
		{
			note:          "v0 -> v1", // from default rego-version
			toRegoVersion: ast.RegoV1,
			module: `package test

p[x]            {
	x = "a"
}`,
			expErrs: []string{
				"test.rego:3: rego_parse_error: `if` keyword is required before rule body",
				"test.rego:3: rego_parse_error: `contains` keyword is required for partial set rules",
			},
		},
		{
			note:          "v1 -> v1", // from non-default rego-version
			toRegoVersion: ast.RegoV1,
			module: `package test

p    contains    x    if      {
	x = "a"
}`,
			expFormatted: `package test

p contains x if {
	x = "a"
}
`,
		},
	}

	for _, tc := range tests {
		t.Run(tc.note, func(t *testing.T) {
			formatted, err := SourceWithOpts("test.rego", []byte(tc.module), Opts{RegoVersion: tc.toRegoVersion})
			if len(tc.expErrs) > 0 {
				if err == nil {
					t.Fatalf("expected errors but got nil")
				}

				for _, expErr := range tc.expErrs {
					if !strings.Contains(err.Error(), expErr) {
						t.Fatalf("expected error:\n\n%q\n\nbut got:\n\n%q", expErr, err)
					}
				}
			} else {
				if err != nil {
					t.Fatalf("unexpected error: %v", err)
				}

				formattedStr := string(formatted)
				if formattedStr != tc.expFormatted {
					t.Fatalf("expected %q but got %q", tc.expFormatted, formattedStr)
				}
			}
		})
	}
}

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)
	}
}

func TestFormatKeywordsInRefs(t *testing.T) {
	regoVersions := map[string]ast.RegoVersion{
		"v0": ast.RegoV0,
		"v1": ast.RegoV1,
	}

	for versionName, regoVersion := range regoVersions {
		t.Run(versionName, func(t *testing.T) {
			regoFiles, err := filepath.Glob(fmt.Sprintf("testfiles/%s/*.rego.formatted_no_keywords_in_refs", versionName))
			if err != nil {
				panic(err)
			}

			for _, expected_rego := range regoFiles {
				original_rego := strings.TrimSuffix(expected_rego, ".formatted_no_keywords_in_refs")
				t.Run(original_rego, func(t *testing.T) {
					expected, err := os.ReadFile(expected_rego)
					if err != nil {
						t.Fatalf("Failed to read expected rego source: %v", err)
					}

					original, err := os.ReadFile(original_rego)
					if err != nil {
						t.Fatalf("Failed to read rego source: %v", err)
					}

					caps := ast.CapabilitiesForThisVersion(ast.CapabilitiesRegoVersion(regoVersion))
					feats := []string{}
					for _, f := range caps.Features {
						if f != ast.FeatureKeywordsInRefs {
							feats = append(feats, f)
						}
					}
					caps.Features = feats

					popts := ast.ParserOptions{
						RegoVersion: regoVersion,
					}
					opts := Opts{
						RegoVersion:   regoVersion,
						ParserOptions: &popts,
						Capabilities:  caps,
					}

					formatted, err := SourceWithOpts(original_rego, original, 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(original_rego+".tmp", string(formatted), popts); err != nil {
						t.Fatalf("Failed to parse formatted bytes: %v", err)
					}

					formatted, err = SourceWithOpts(original_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))
					}

				})
			}
		})
	}
}

// Regression test for #8557
func TestSingleItemHonorMultiline(t *testing.T) {
	contents := `package test

x0 := {{"foo", "bar"}}

x1 := {{
	"foo",
	"bar",
}}

x2 := {
	{
		"foo",
		"bar",
	},
}

x3 := {"foo": {
	"bar": "baz",
}}
`
	module := ast.MustParseModuleWithOpts(contents, ast.ParserOptions{RegoVersion: ast.RegoV1})
	formatted := string(MustAstWithOpts(module, Opts{RegoVersion: ast.RegoV1}))

	if contents != formatted {
		t.Fatalf("expected \n%sgot\n%s", contents, formatted)
	}
}

// 3064960 ns/op	 4573131 B/op	   26266 allocs/op // no optimizations
// 1737719 ns/op	 1972193 B/op	   14160 allocs/op // pre-allocate partitionComments
// 1674343 ns/op	 1916700 B/op	   11556 allocs/op // static memberRef & memberWithKeyRef
// 1594546 ns/op	 1882652 B/op	   10644 allocs/op // various minor fixes
// _853508 ns/op	  441730 B/op	    8895 allocs/op // partitionComments early return if unchanged
// _812859 ns/op	  362651 B/op	    8811 allocs/op // partitionComments reuse backing array
// _676179 ns/op	  995204 B/op	    8850 allocs/op // regression in 0fb7526
// _513570 ns/op	  378787 B/op	    8775 allocs/op // addressed regression with sync.Pool
// _481681 ns/op	  352130 B/op	    7954 allocs/op // new util.NewSlicePool using only pointers
// _365116 ns/op	  160528 B/op	    2098 allocs/op // new SkipDefensiveCopying option
func BenchmarkFormatLargePolicy(b *testing.B) {
	contents, err := os.ReadFile("testdata/bench.rego")
	if err != nil {
		b.Fatalf("Failed to read rego source: %v", err)
	}
	module := ast.MustParseModule(string(contents))
	opts := Opts{RegoVersion: ast.RegoV1, SkipDefensiveCopying: true}

	for b.Loop() {
		if _, err := AstWithOpts(module, opts); err != nil {
			b.Fatal(err)
		}
	}
}