Arbitrary vars in rule refs

ast/compile.go

@@ -7,6 +7,7 @@ package ast
 import (
 	"fmt"
 	"io"
+	"os"
 	"sort"
 	"strconv"
 	"strings"
@@ -81,6 +82,28 @@ type Compiler struct {
 	//                     +--- b
 	//                          |
 	//                          +--- c (1 rule)
+	//
+	// Another example with general refs containing vars at arbitrary locations:
+	//
+	//  package ex
+	//  a.b[x].d { x := "c" }            # R1
+	//  a.b.c[x] { x := "d" }            # R2
+	//  a.b[x][y] { x := "c"; y := "d" } # R3
+	//  p := true                        # R4
+	//
+	//  root
+	//    |
+	//    +--- data (no rules)
+	//           |
+	//           +--- ex (no rules)
+	//                |
+	//                +--- a
+	//                |    |
+	//                |    +--- b (R1, R3)
+	//                |         |
+	//                |         +--- c (R2)
+	//                |
+	//                +--- p (R4)
 	RuleTree *TreeNode
 
 	// Graph contains dependencies between rules. An edge (u,v) is added to the
@@ -123,6 +146,7 @@ type Compiler struct {
 	keepModules             bool                          // whether to keep the unprocessed, parse modules (below)
 	parsedModules           map[string]*Module            // parsed, but otherwise unprocessed modules, kept track of when keepModules is true
 	useTypeCheckAnnotations bool                          // whether to provide annotated information (schemas) to the type checker
+	generalRuleRefsEnabled  bool
 }
 
 // CompilerStage defines the interface for stages in the compiler.
@@ -311,6 +335,8 @@ func NewCompiler() *Compiler {
 		{"BuildComprehensionIndices", "compile_stage_rebuild_comprehension_indices", c.buildComprehensionIndices},
 	}
 
+	_, c.generalRuleRefsEnabled = os.LookupEnv("OPA_ENABLE_GENERAL_RULE_REFS")
+
 	return c
 }
 
@@ -790,17 +816,15 @@ func (c *Compiler) buildRuleIndices() {
 		rules := extractRules(node.Values)
 		hasNonGroundKey := false
 		for _, r := range rules {
-			if ref := r.Head.Ref(); len(ref) > 1 {
-				if !ref[len(ref)-1].IsGround() {
-					hasNonGroundKey = true
-				}
-			}
+			hasNonGroundKey = !r.Head.Ref().IsGround()
 		}
 		if hasNonGroundKey {
-			// collect children: as of now, this cannot go deeper than one level,
-			// so we grab those, and abort the DepthFirst processing for this branch
-			for _, n := range node.Children {
-				rules = append(rules, extractRules(n.Values)...)
+			// collect children
+			for _, child := range node.Children {
+				child.DepthFirst(func(c *TreeNode) bool {
+					rules = append(rules, extractRules(c.Values)...)
+					return false
+				})
 			}
 		}
 
@@ -870,10 +894,11 @@ func (c *Compiler) checkRuleConflicts() {
 
 		kinds := make(map[RuleKind]struct{}, len(node.Values))
 		defaultRules := 0
+		completeRules := 0
+		partialRules := 0
 		arities := make(map[int]struct{}, len(node.Values))
 		name := ""
-		var singleValueConflicts []Ref
-		var multiValueConflicts []Ref
+		var conflicts []Ref
 
 		for _, rule := range node.Values {
 			r := rule.(*Rule)
@@ -885,70 +910,99 @@ func (c *Compiler) checkRuleConflicts() {
 				defaultRules++
 			}
 
-			// Single-value rules may not have any other rules in their extent: these pairs are invalid:
+			// Single-value rules may not have any other rules in their extent.
+			// Rules with vars in their ref are allowed to have rules inside their extent.
+			// Only the ground portion (terms before the first var term) of a rule's ref is considered when determining
+			// whether it's inside the extent of another (c.RuleTree is organized this way already).
+			// These pairs are invalid:
 			//
 			//   data.p.q.r { true }          # data.p.q is { "r": true }
 			//   data.p.q.r.s { true }
 			//
-			//   data.p.q[r] { r := input.r } # data.p.q could be { "r": true }
-			//   data.p.q.r.s { true }
+			//   data.p.q.r { true }
+			//   data.p.q.r[s].t { s = input.key }
 			//
-			// data.p[r] := x { r = input.key; x = input.bar }
-			// data.p.q[r] := x { r = input.key; x = input.bar }
-
 			// But this is allowed:
+			//
+			//   data.p.q.r { true }
+			//   data.p.q[r].s.t { r = input.key }
+			//
+			//   data.p[r] := x { r = input.key; x = input.bar }
+			//   data.p.q[r] := x { r = input.key; x = input.bar }
+			//
+			//   data.p.q[r] { r := input.r }
+			//   data.p.q.r.s { true }
+			//
 			//   data.p.q[r] = 1 { r := "r" }
 			//   data.p.q.s = 2
+			//
+			//   data.p[q][r] { q := input.q; r := input.r }
+			//   data.p.q.r { true }
+			//
+			//   data.p.q[r] { r := input.r }
+			//   data.p[q].r { q := input.q }
+			//
+			//   data.p.q[r][s] { r := input.r; s := input.s }
+			//   data.p[q].r.s { q := input.q }
 
-			if r.Head.RuleKind() == SingleValue && len(node.Children) > 0 {
-				if len(ref) > 1 && !ref[len(ref)-1].IsGround() { // p.q[x] and p.q.s.t => check grandchildren
-					for _, c := range node.Children {
-						grandchildrenFound := false
-
-						if len(c.Values) > 0 {
-							childRules := extractRules(c.Values)
-							for _, childRule := range childRules {
-								childRef := childRule.Ref()
-								if childRule.Head.RuleKind() == SingleValue && !childRef[len(childRef)-1].IsGround() {
-									// The child is a partial object rule, so it's effectively "generating" grandchildren.
-									grandchildrenFound = true
-									break
+			if c.generalRuleRefsEnabled {
+				if r.Ref().IsGround() && len(node.Children) > 0 {
+					conflicts = node.flattenChildren()
+				}
+			} else { // TODO: Remove when general rule refs are enabled by default.
+				if r.Head.RuleKind() == SingleValue && len(node.Children) > 0 {
+					if len(ref) > 1 && !ref[len(ref)-1].IsGround() { // p.q[x] and p.q.s.t => check grandchildren
+						for _, c := range node.Children {
+							grandchildrenFound := false
+
+							if len(c.Values) > 0 {
+								childRules := extractRules(c.Values)
+								for _, childRule := range childRules {
+									childRef := childRule.Ref()
+									if childRule.Head.RuleKind() == SingleValue && !childRef[len(childRef)-1].IsGround() {
+										// The child is a partial object rule, so it's effectively "generating" grandchildren.
+										grandchildrenFound = true
+										break
+									}
 								}
 							}
-						}
 
-						if len(c.Children) > 0 {
-							grandchildrenFound = true
-						}
+							if len(c.Children) > 0 {
+								grandchildrenFound = true
+							}
 
-						if grandchildrenFound {
-							singleValueConflicts = node.flattenChildren()
-							break
+							if grandchildrenFound {
+								conflicts = node.flattenChildren()
+								break
+							}
 						}
+					} else { // p.q.s and p.q.s.t => any children are in conflict
+						conflicts = node.flattenChildren()
 					}
-				} else { // p.q.s and p.q.s.t => any children are in conflict
-					singleValueConflicts = node.flattenChildren()
 				}
-			}
 
-			// Multi-value rules may not have any other rules in their extent; e.g.:
-			//
-			// data.p[v] { v := ... }
-			// data.p.q := 42 # In direct conflict with data.p[v], which is constructing a set and cannot have values assigned to a sub-path.
+				// Multi-value rules may not have any other rules in their extent; e.g.:
+				//
+				// data.p[v] { v := ... }
+				// data.p.q := 42 # In direct conflict with data.p[v], which is constructing a set and cannot have values assigned to a sub-path.
 
-			if r.Head.RuleKind() == MultiValue && len(node.Children) > 0 {
-				multiValueConflicts = node.flattenChildren()
+				if r.Head.RuleKind() == MultiValue && len(node.Children) > 0 {
+					conflicts = node.flattenChildren()
+				}
+			}
+
+			if r.Head.RuleKind() == SingleValue && r.Head.Ref().IsGround() {
+				completeRules++
+			} else {
+				partialRules++
 			}
 		}
 
 		switch {
-		case singleValueConflicts != nil:
-			c.err(NewError(TypeErr, node.Values[0].(*Rule).Loc(), "single-value rule %v conflicts with %v", name, singleValueConflicts))
-
-		case multiValueConflicts != nil:
-			c.err(NewError(TypeErr, node.Values[0].(*Rule).Loc(), "multi-value rule %v conflicts with %v", name, multiValueConflicts))
+		case conflicts != nil:
+			c.err(NewError(TypeErr, node.Values[0].(*Rule).Loc(), "rule %v conflicts with %v", name, conflicts))
 
-		case len(kinds) > 1 || len(arities) > 1:
+		case len(kinds) > 1 || len(arities) > 1 || (completeRules >= 1 && partialRules >= 1):
 			c.err(NewError(TypeErr, node.Values[0].(*Rule).Loc(), "conflicting rules %v found", name))
 
 		case defaultRules > 1:
@@ -1697,13 +1751,12 @@ func (c *Compiler) rewriteRuleHeadRefs() {
 			}
 
 			for i := 1; i < len(ref); i++ {
-				// NOTE(sr): In the first iteration, non-string values in the refs are forbidden
+				// NOTE: Unless enabled via the OPA_ENABLE_GENERAL_RULE_REFS env var, non-string values in the refs are forbidden
 				// except for the last position, e.g.
 				//     OK: p.q.r[s]
 				// NOT OK: p[q].r.s
-				// TODO(sr): This is stricter than necessary. We could allow any non-var values there,
-				// but we'll also have to adjust the type tree, for example.
-				if i != len(ref)-1 { // last
+				// TODO: Remove when general rule refs are enabled by default.
+				if !c.generalRuleRefsEnabled && i != len(ref)-1 { // last
 					if _, ok := ref[i].Value.(String); !ok {
 						c.err(NewError(TypeErr, rule.Loc(), "rule head must only contain string terms (except for last): %v", ref[i]))
 						continue