ast: Extending support for file-level assignments (:=) (#4583)

Updated support for:
* default values
* rules with `else` keyword
* partial rules
* functions

Fixes: #4555

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

ast/errors.go

@@ -121,12 +121,3 @@ func NewError(code string, loc *Location, f string, a ...interface{}) *Error {
 		Message:  fmt.Sprintf(f, a...),
 	}
 }
-
-var (
-	errPartialRuleAssignOperator = fmt.Errorf("partial rules must use = operator (not := operator)")
-	errFunctionAssignOperator    = fmt.Errorf("functions must use = operator (not := operator)")
-)
-
-func errTermAssignOperator(x interface{}) error {
-	return fmt.Errorf("cannot assign to %v", TypeName(x))
-}

ast/parser.go

@@ -569,11 +569,6 @@ func (p *Parser) parseRules() []*Rule {
 
 	if p.s.tok == tokens.Else {
 
-		if rule.Head.Assign {
-			p.error(p.s.Loc(), "else keyword cannot be used on rule declared with := operator")
-			return nil
-		}
-
 		if rule.Head.Key != nil {
 			p.error(p.s.Loc(), "else keyword cannot be used on partial rules")
 			return nil
@@ -639,7 +634,7 @@ func (p *Parser) parseElse(head *Head) *Rule {
 	switch p.s.tok {
 	case tokens.LBrace:
 		rule.Head.Value = BooleanTerm(true)
-	case tokens.Unify:
+	case tokens.Assign, tokens.Unify:
 		p.scan()
 		rule.Head.Value = p.parseTermInfixCall()
 		if rule.Head.Value == nil {
@@ -727,23 +722,22 @@ func (p *Parser) parseHead(defaultRule bool) *Head {
 			p.illegal("expected rule value term (e.g., %s[%s] = <VALUE> { ... })", head.Name, head.Key)
 		}
 	} else if p.s.tok == tokens.Assign {
-
-		if defaultRule {
-			p.error(p.s.Loc(), "default rules must use = operator (not := operator)")
-			return nil
-		} else if head.Key != nil {
-			p.error(p.s.Loc(), "partial rules must use = operator (not := operator)")
-			return nil
-		} else if len(head.Args) > 0 {
-			p.error(p.s.Loc(), "functions must use = operator (not := operator)")
-			return nil
-		}
-
+		s := p.save()
 		p.scan()
 		head.Assign = true
 		head.Value = p.parseTermInfixCall()
 		if head.Value == nil {
-			p.illegal("expected rule value term (e.g., %s := <VALUE> { ... })", head.Name)
+			p.restore(s)
+			switch {
+			case len(head.Args) > 0:
+				p.illegal("expected function value term (e.g., %s(...) := <VALUE> { ... })", head.Name)
+			case head.Key != nil:
+				p.illegal("expected partial rule value term (e.g., %s[...] := <VALUE> { ... })", head.Name)
+			case defaultRule:
+				p.illegal("expected default rule value term (e.g., default %s := <VALUE>)", head.Name)
+			default:
+				p.illegal("expected rule value term (e.g., %s := <VALUE> { ... })", head.Name)
+			}
 		}
 	}
 

ast/parser_ext.go

@@ -173,30 +173,15 @@ func ParseRuleFromExpr(module *Module, expr *Expr) (*Rule, error) {
 		return nil, errors.New("expression cannot be used for rule head")
 	}
 
-	if expr.IsAssignment() {
-
-		lhs, rhs := expr.Operand(0), expr.Operand(1)
-		if lhs == nil || rhs == nil {
-			return nil, errors.New("assignment requires two operands")
-		}
-
-		rule, err := ParseCompleteDocRuleFromAssignmentExpr(module, lhs, rhs)
-
-		if err == nil {
-			rule.Location = expr.Location
-			rule.Head.Location = expr.Location
-			return rule, nil
-		} else if _, ok := lhs.Value.(Call); ok {
-			return nil, errFunctionAssignOperator
-		} else if _, ok := lhs.Value.(Ref); ok {
-			return nil, errPartialRuleAssignOperator
-		}
-
-		return nil, errTermAssignOperator(lhs.Value)
-	}
-
 	if expr.IsEquality() {
 		return parseCompleteRuleFromEq(module, expr)
+	} else if expr.IsAssignment() {
+		rule, err := parseCompleteRuleFromEq(module, expr)
+		if err != nil {
+			return nil, err
+		}
+		rule.Head.Assign = true
+		return rule, nil
 	}
 
 	if _, ok := BuiltinMap[expr.Operator().String()]; ok {

ast/parser_test.go

@@ -1379,6 +1379,16 @@ func TestRule(t *testing.T) {
 		Body:    NewBody(NewExpr(BooleanTerm(true))),
 	})
 
+	assertParseRule(t, "default w/ assignment", `default allow := false`, &Rule{
+		Default: true,
+		Head: &Head{
+			Name:   "allow",
+			Value:  BooleanTerm(false),
+			Assign: true,
+		},
+		Body: NewBody(NewExpr(BooleanTerm(true))),
+	})
+
 	assertParseRule(t, "default w/ comprehension", `default widgets = [x | x = data.fooz[_]]`, &Rule{
 		Default: true,
 		Head:    NewHead(Var("widgets"), nil, MustParseTerm(`[x | x = data.fooz[_]]`)),
@@ -1477,6 +1487,45 @@ func TestRule(t *testing.T) {
 		Body: NewBody(NewExpr(BooleanTerm(true))),
 	})
 
+	assertParseRule(t, "else assignment", `x := 1 { false } else := 2`, &Rule{
+		Head: &Head{
+			Name:   "x",
+			Value:  IntNumberTerm(1),
+			Assign: true,
+		},
+		Body: NewBody(NewExpr(BooleanTerm(false))),
+		Else: &Rule{
+			Head: &Head{
+				Name:   "x",
+				Value:  IntNumberTerm(2),
+				Assign: true,
+			},
+			Body: NewBody(NewExpr(BooleanTerm(true))),
+		},
+	})
+
+	assertParseRule(t, "partial assignment", `p[x] := y { true }`, &Rule{
+		Head: &Head{
+			Name:   "p",
+			Value:  VarTerm("y"),
+			Key:    VarTerm("x"),
+			Assign: true,
+		},
+		Body: NewBody(NewExpr(BooleanTerm(true))),
+	})
+
+	assertParseRule(t, "function assignment", `f(x) := y { true }`, &Rule{
+		Head: &Head{
+			Name:  "f",
+			Value: VarTerm("y"),
+			Args: Args{
+				VarTerm("x"),
+			},
+			Assign: true,
+		},
+		Body: NewBody(NewExpr(BooleanTerm(true))),
+	})
+
 	// TODO: expect expressions instead?
 	assertParseErrorContains(t, "empty body", `f(_) = y {}`, "rego_parse_error: found empty body")
 	assertParseErrorContains(t, "empty rule body", "p {}", "rego_parse_error: found empty body")
@@ -1484,17 +1533,16 @@ func TestRule(t *testing.T) {
 
 	// TODO: how to highlight that assignment is incorrect here?
 	assertParseErrorContains(t, "no output", `f(_) = { "foo" = "bar" }`, "rego_parse_error: unexpected eq token: expected rule value term")
+	assertParseErrorContains(t, "no output", `f(_) := { "foo" = "bar" }`, "rego_parse_error: unexpected assign token: expected function value term")
+	assertParseErrorContains(t, "no output", `f := { "foo" = "bar" }`, "rego_parse_error: unexpected assign token: expected rule value term")
+	assertParseErrorContains(t, "no output", `f[_] := { "foo" = "bar" }`, "rego_parse_error: unexpected assign token: expected partial rule value term")
+	assertParseErrorContains(t, "no output", `default f :=`, "rego_parse_error: unexpected assign token: expected default rule value term")
 
 	// TODO(tsandall): improve error checking here. This is a common mistake
 	// and the current error message is not very good. Need to investigate if the
 	// parser can be improved.
 	assertParseError(t, "dangling semicolon", "p { true; false; }")
 
-	assertParseErrorContains(t, "default assignment", "default p := 1", `default rules must use = operator (not := operator)`)
-	assertParseErrorContains(t, "partial assignment", `p[x] := y { true }`, "partial rules must use = operator (not := operator)")
-	assertParseErrorContains(t, "function assignment", `f(x) := y { true }`, "functions must use = operator (not := operator)")
-	assertParseErrorContains(t, "else assignment", `p := y { true } else = 2 { true } `, "else keyword cannot be used on rule declared with := operator")
-
 	assertParseErrorContains(t, "default invalid rule name", `default 0[0`, "unexpected default keyword")
 	assertParseErrorContains(t, "default invalid rule value", `default a[0`, "illegal default rule (must have a value)")
 	assertParseRule(t, "default missing value", `default a`, &Rule{

docs/content/policy-reference.md

@@ -1236,7 +1236,7 @@ policy          = { rule }
 rule            = [ "default" ] rule-head { rule-body }
 rule-head       = var [ "(" rule-args ")" ] [ "[" term "]" ] [ ( ":=" | "=" ) term ]
 rule-args       = term { "," term }
-rule-body       = [ "else" [ "=" term ] ] "{" query "}"
+rule-body       = [ "else" [ ( ":=" | "=" ) term ] ] "{" query "}"
 query           = literal { ( ";" | ( [CR] LF ) ) literal }
 literal         = ( some-decl | expr | "not" expr ) { with-modifier }
 with-modifier   = "with" term "as" term

format/testfiles/test_assignments.rego

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

format/testfiles/test_assignments.rego.formatted

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

test/cases/testdata/assignments/test-file-level-assignments.yaml

@@ -0,0 +1,48 @@
+cases:
+  - note: assignments/file-level/default_value
+    query: data.test = x
+    modules:
+      - |
+        package test
+
+        default a := 1
+    want_result: [ { "x": { "a": 1 } } ]
+  - note: assignments/file-level/rule
+    query: data.test = x
+    modules:
+      - |
+        package test
+        
+        b := 2
+    want_result: [ { "x": { "b": 2 } } ]
+  - note: assignments/file-level/else_keyword
+    query: data.test = x
+    modules:
+      - |
+        package test
+        
+        c := 3 {
+           false
+        } else := 4 {
+           true
+        }
+    want_result: [ { "x": { "c": 4, } } ]
+  - note: assignments/file-level/partial_rule
+    query: data.test = x
+    modules:
+      - |
+        package test
+        
+        d[msg] := 5 {
+           msg = [1, 2, 3][_]
+        }
+    want_result: [ { "x": { "d": { "1": 5, "2": 5, "3": 5 } } } ]
+  - note: assignments/file-level/function_return_value
+    query: data.test = x
+    modules:
+      - |
+        package test
+        
+        e := f(6)
+        f(x) := x
+    want_result: [ { "x": { "e": 6 } } ]