format: Add support for and/or logical expressions

cmd/fmt.go

@@ -10,6 +10,7 @@ import (
 	"io"
 	"os"
 	"path/filepath"
+	"slices"
 
 	"github.com/sergi/go-diff/diffmatchpatch"
 	"github.com/spf13/cobra"
@@ -61,6 +62,49 @@ func (p *fmtCommandParams) regoVersion() ast.RegoVersion {
 	return ast.DefaultRegoVersion
 }
 
+// parserOptions returns the options for parsing the source to format.
+func (p *fmtCommandParams) parserOptions() *ast.ParserOptions {
+	popts := ast.ParserOptions{
+		RegoVersion:  ast.DefaultRegoVersion,
+		Capabilities: p.parserCapabilities(),
+	}
+
+	switch {
+	case p.v0Compatible:
+		popts.RegoVersion = ast.RegoV0
+	case p.v1Compatible:
+		popts.RegoVersion = ast.RegoV1
+	case p.regoV1:
+		// '--rego-v1' formats a v0 module for v1 compatibility, so it is read as v0.
+		popts.RegoVersion = ast.RegoV0
+	}
+
+	return &popts
+}
+
+func (p *fmtCommandParams) parserCapabilities() *ast.Capabilities {
+	caps := *p.capabilities()
+	current := ast.CapabilitiesForThisVersion(ast.CapabilitiesRegoVersion(p.regoVersion()))
+
+	// parsing caps must span current version and target
+	caps.FutureKeywords = union(caps.FutureKeywords, current.FutureKeywords)
+	caps.Features = union(caps.Features, current.Features)
+
+	return &caps
+}
+
+func union(a, b []string) []string {
+	out := slices.Clone(a)
+	for _, s := range b {
+		if !slices.Contains(out, s) {
+			out = append(out, s)
+		}
+	}
+	slices.Sort(out)
+
+	return out
+}
+
 func opaFmt(args []string, fmtParams *fmtCommandParams) int {
 	if len(args) == 0 {
 		if err := formatStdin(fmtParams, os.Stdin, os.Stdout); err != nil {
@@ -118,17 +162,7 @@ func formatFile(params *fmtCommandParams, out io.Writer, filename string, info o
 		RegoVersion:   params.regoVersion(),
 		DropV0Imports: params.dropV0Imports,
 		Capabilities:  params.capabilities(),
-	}
-
-	if params.regoV1 {
-		opts.ParserOptions = &ast.ParserOptions{RegoVersion: ast.RegoV0}
-	}
-
-	if params.v0Compatible {
-		// v0 takes precedence over v1
-		opts.ParserOptions = &ast.ParserOptions{RegoVersion: ast.RegoV0}
-	} else if params.v1Compatible {
-		opts.ParserOptions = &ast.ParserOptions{RegoVersion: ast.RegoV1}
+		ParserOptions: params.parserOptions(),
 	}
 
 	formatted, err := format.SourceWithOpts(filename, contents, opts)
@@ -137,7 +171,10 @@ func formatFile(params *fmtCommandParams, out io.Writer, filename string, info o
 	}
 
 	if params.checkResult {
-		popts := ast.ParserOptions{RegoVersion: params.regoVersion()}
+		popts := ast.ParserOptions{
+			RegoVersion:  params.regoVersion(),
+			Capabilities: params.parserCapabilities(),
+		}
 		_, err := ast.ParseModuleWithOpts("formatted", string(formatted), popts)
 		if err != nil {
 			return newError("%s was successfully formatted, but the result is invalid: %v\n\nTo inspect the formatted Rego, you can turn off this check with --check-result=false.", filename, err)
@@ -203,19 +240,9 @@ func formatStdin(params *fmtCommandParams, r io.Reader, w io.Writer) error {
 	}
 
 	opts := format.Opts{
-		RegoVersion:  params.regoVersion(),
-		Capabilities: params.capabilities(),
-	}
-
-	if params.regoV1 {
-		opts.ParserOptions = &ast.ParserOptions{RegoVersion: ast.RegoV0}
-	}
-
-	if params.v0Compatible {
-		// v0 takes precedence over v1
-		opts.ParserOptions = &ast.ParserOptions{RegoVersion: ast.RegoV0}
-	} else if params.v1Compatible {
-		opts.ParserOptions = &ast.ParserOptions{RegoVersion: ast.RegoV1}
+		RegoVersion:   params.regoVersion(),
+		Capabilities:  params.capabilities(),
+		ParserOptions: params.parserOptions(),
 	}
 
 	formatted, err := format.SourceWithOpts("stdin", contents, opts)

cmd/fmt_test.go

@@ -1346,6 +1346,119 @@ foo["if"]["else"] := true
 	}
 }
 
+// Experimental future keywords are hidden from the default capabilities, so
+// formatting a module that imports one requires --capabilities to list it.
+func TestFmtFormatExperimentalKeywords(t *testing.T) {
+	unformatted := `package test
+
+import future.keywords.and
+import future.keywords.or
+
+p if { input.a   and   input.b or input.c }
+`
+
+	formatted := `package test
+
+import future.keywords.and
+import future.keywords.or
+
+p if input.a and input.b or input.c
+`
+
+	experimental := func() fmtCommandParams {
+		params := newFmtCommandParams()
+		params.capabilitiesFlag.C = ast.CapabilitiesForThisVersion(
+			ast.CapabilitiesRegoVersion(ast.RegoV1),
+			ast.CapabilitiesExperimentalKeywords(true))
+		return *params
+	}
+
+	cases := []struct {
+		note        string
+		params      fmtCommandParams
+		expected    string
+		expectedErr string
+	}{
+		{
+			note:        "default capabilities",
+			params:      *newFmtCommandParams(),
+			expectedErr: "rego_parse_error: unexpected keyword, must be one of [contains every if in not]",
+		},
+		{
+			note:     "capabilities listing the keywords",
+			params:   experimental(),
+			expected: formatted,
+		},
+		{
+			note: "capabilities listing the keywords, --check-result",
+			params: func() fmtCommandParams {
+				params := experimental()
+				params.checkResult = true
+				return params
+			}(),
+			expected: formatted,
+		},
+	}
+
+	for _, tc := range cases {
+		t.Run("file, "+tc.note, func(t *testing.T) {
+			var stdout bytes.Buffer
+
+			files := map[string]string{"policy.rego": unformatted}
+
+			test.WithTempFS(files, func(path string) {
+				policyFile := filepath.Join(path, "policy.rego")
+				info, err := os.Stat(policyFile)
+				if err != nil {
+					t.Fatalf("Unexpected error: %s", err)
+				}
+
+				err = formatFile(&tc.params, &stdout, policyFile, info, err)
+
+				if tc.expectedErr != "" {
+					if err == nil {
+						t.Fatalf("Expected error but got: %s", stdout.String())
+					}
+					if !strings.Contains(err.Error(), tc.expectedErr) {
+						t.Fatalf("Expected error to contain:\n\n%s\n\nGot:\n\n%s", tc.expectedErr, err.Error())
+					}
+					return
+				}
+
+				if err != nil {
+					t.Fatalf("Unexpected error: %s", err)
+				}
+				if actual := stdout.String(); actual != tc.expected {
+					t.Fatalf("Expected:\n%s\n\nGot:\n%s\n\n", tc.expected, actual)
+				}
+			})
+		})
+
+		t.Run("stdin, "+tc.note, func(t *testing.T) {
+			var stdout bytes.Buffer
+
+			err := formatStdin(&tc.params, bytes.NewReader([]byte(unformatted)), &stdout)
+
+			if tc.expectedErr != "" {
+				if err == nil {
+					t.Fatalf("Expected error but got: %s", stdout.String())
+				}
+				if !strings.Contains(err.Error(), tc.expectedErr) {
+					t.Fatalf("Expected error to contain:\n\n%s\n\nGot:\n\n%s", tc.expectedErr, err.Error())
+				}
+				return
+			}
+
+			if err != nil {
+				t.Fatalf("Unexpected error: %s", err)
+			}
+			if actual := stdout.String(); actual != tc.expected {
+				t.Fatalf("Expected:\n%s\n\nGot:\n%s\n\n", tc.expected, actual)
+			}
+		})
+	}
+}
+
 func TestFmtFormatFile_KeywordsInRefs(t *testing.T) {
 	cases := []struct {
 		note        string

v1/ast/parser.go

@@ -1309,7 +1309,9 @@ func (p *Parser) parseLiteral() (expr *Expr) {
 		if nb == nil {
 			return nil
 		}
-		return p.attachWith(nb)
+		// A not-body is a complete operand, so it may lead an and/or chain:
+		// `not { x } and y`.
+		return p.foldLogicalTail(NewBody(nb), false, nb.Location)
 	}
 
 	switch p.s.tok {
@@ -1417,6 +1419,17 @@ func (p *Parser) parseLiteralExpr(negated bool, notLoc *Location) *Expr {
 				expr.SetLoc(startLoc)
 			}
 
+			if notLoc == nil && bytes.HasPrefix(expr.Location.Text, []byte("{")) {
+				// `{}` on its own is an empty body
+				if isEmptyObjectTerm(expr) {
+					p.error(expr.Location, "found empty body")
+					return nil
+				}
+
+				p.errorBraceLedOperand(expr.Location, expr.Location.Text, p.s.tok.String())
+				return nil
+			}
+
 			outer := p.parseLogicalOrChain(NewBody(expr), false, expr.Location)
 			if outer == nil {
 				return nil
@@ -1605,7 +1618,9 @@ func (p *Parser) parseNotBody(notLoc *Location) *Expr {
 		failed := p.save()
 		p.restore(s)
 
-		if term := p.parseTerm(); term != nil {
+		// The operand can extend past the braces (`{1, 2} & input.s == set()`),
+		// and parens group rather than delimit, so it is the whole operand that has to be wrapped.
+		if term := p.parseTermInfixCall(); term != nil {
 			p.errorOperandBraceNeedsBody(braceLoc, p.s.Text(braceOffset, p.s.lastEnd), term, "not ")
 			return nil
 		}
@@ -1662,7 +1677,7 @@ func (p *Parser) parseLogicalOrChain(lhsBody Body, lhsExplicit bool, lhsLoc *Loc
 	for p.s.tok == tokens.LogicalOr {
 		p.scan()
 
-		rhsBody, rhsExplicit, rhsLoc := p.parseLogicalOperand()
+		rhsBody, rhsExplicit, rhsLoc := p.parseLogicalOperand("or")
 		if rhsBody == nil {
 			return nil
 		}
@@ -1709,7 +1724,7 @@ func (p *Parser) parseLogicalAndChain(lhsBody Body, lhsExplicit bool, lhsLoc *Lo
 	for p.s.tok == tokens.LogicalAnd {
 		p.scan()
 
-		rhsBody, rhsExplicit, _ := p.parseLogicalOperand()
+		rhsBody, rhsExplicit, _ := p.parseLogicalOperand("and")
 		if rhsBody == nil {
 			return nil
 		}
@@ -1751,14 +1766,39 @@ func isNegated(p *Parser) bool {
 	return tok != tokens.Dot && tok != tokens.LBrack
 }
 
-// parseLogicalOperand parses a single operand of an `and`/`or` expression.
-func (p *Parser) parseLogicalOperand() (Body, bool, *Location) {
+// parseLogicalOperand parses a single operand of an `and`/`or` expression. op is
+// the operator the operand belongs to, or "" when the caller is speculating and
+// will restore on failure.
+func (p *Parser) parseLogicalOperand(op string) (Body, bool, *Location) {
 	if p.s.tok == tokens.LBrace {
 		braceOffset := p.s.loc.Offset
 		loc := p.s.Loc()
+		s := p.save()
 		p.scan()
+
+		// `{}` is an empty body, which parseBody reports precisely; only non-empty
+		// braces are worth re-reading as a value.
+		empty := p.s.tok == tokens.RBrace
+
 		body := p.parseBody(tokens.RBrace)
 		if body == nil {
+			if empty || op == "" {
+				return nil, false, nil
+			}
+
+			// The braces may hold a value rather than a body.
+			failed := p.save()
+			p.restore(s)
+
+			// The operand can extend past the braces (`{1, 2} & input.s == set()`),
+			// and parens group rather than delimit, so it is the whole operand that has to be wrapped.
+			if term := p.parseTermInfixCall(); term != nil {
+				p.errorBraceLedOperand(loc, p.s.Text(braceOffset, p.s.lastEnd), op)
+				return nil, false, nil
+			}
+
+			p.restore(failed)
+
 			return nil, false, nil
 		}
 		p.scan()
@@ -1861,10 +1901,50 @@ func isAmbiguousUnionBody(b Body) bool {
 }
 
 // errorOperandBraceNeedsBody reports `{...}` in an operand position holding a value instead of expressions.
-func (p *Parser) errorOperandBraceNeedsBody(loc *Location, braces []byte, term *Term, prefix string) {
+func (p *Parser) errorOperandBraceNeedsBody(loc *Location, operand []byte, term *Term, prefix string) {
 	p.hint(fmt.Sprintf("write `%s(%s)` to negate the value, or `%s{%s}` for a body holding it",
-		prefix, braces, prefix, braces))
-	p.errorf(loc, "`{...}` in an operand position must contain expression(s), got: %s", ValueName(term.Value))
+		prefix, operand, prefix, operand))
+	p.errorf(loc, "`{...}` in an operand position must contain expression(s), got: %s", ValueName(braceLedValue(term)))
+}
+
+// braceLedValue returns the value opened by the leading `{` of t. An infix call
+// renders its lhs operand first, so `{1, 2} & s` is brace-led by the set; refs are
+// left alone, as `{"a": 1}["a"]` is reported as the ref it is.
+func braceLedValue(t *Term) Value {
+	if call, ok := t.Value.(Call); ok && len(call) > 0 {
+		if bi, ok := BuiltinMap[call[0].String()]; ok && bi.Infix != "" && len(call) == bi.Decl.Arity()+1 {
+			return braceLedValue(call[1])
+		}
+	}
+
+	return t.Value
+}
+
+// isEmptyObjectTerm reports whether expr is exactly `{}`. In an operand position
+// those braces open a body, so an empty one is an empty body - not the empty
+// object the term parser read.
+func isEmptyObjectTerm(expr *Expr) bool {
+	if len(expr.With) > 0 {
+		return false
+	}
+
+	t, ok := expr.Terms.(*Term)
+	if !ok {
+		return false
+	}
+
+	obj, ok := t.Value.(Object)
+
+	return ok && obj.Len() == 0
+}
+
+// errorBraceLedOperand reports an `and`/`or` operand whose leading `{` opens a
+// value rather than a body. In an operand position the braces are read as an
+// explicit body, so the value form has to be parenthesized, on both sides of the
+// operator.
+func (p *Parser) errorBraceLedOperand(loc *Location, operand []byte, op string) {
+	p.hint(fmt.Sprintf("wrap the operand to keep the value: `(%s) %s ...`", operand, op))
+	p.errorf(loc, "operand of `%s` cannot begin with `{` unless the braces hold a body", op)
 }
 
 // errorParensCannotWrapBody reports `(...)` holding expressions rather than a value.
@@ -1951,7 +2031,7 @@ func (p *Parser) parseLogicalGroup(operandContext bool, prefix string) (Body, bo
 		return nil, false, nil, false
 	}
 
-	lhsBody, lhsExplicit, lhsLoc := p.parseLogicalOperand()
+	lhsBody, lhsExplicit, lhsLoc := p.parseLogicalOperand("")
 	if lhsBody == nil {
 		// Parens are not an operand, so a `{...}` that can't be a body is a value:
 		// restore and let the term parser read it, e.g. `not ({})` is an empty object.