ast+topdown: exit-early for (constant) complete virtual docs and func…

…tions (open-policy-agent#3898)

If the following conditions hold for a set of rules returned by the indexer,
it will set EarlyExit: true, and change how the complete virtual doc or
function is evaluated:

- all rule head values are ground
- all rule head values match

This implies that some cases where early exit would be possible will not be
covered:

    p = x {
      x := true
      input.foo == "bar"
    }

    p = x {
      x := true
      input.baz == "quz"
    }

To indicate that "early exit" is possible, the indexer result message is
amended. Also, the "Exit" trace event will have a message of "early" when
"early exit" actually happens in eval:

    $ echo '{"x":"x", "y":"y"}' | opa eval -I -fpretty --explain=full -d r.rego data.r.r
    query:1       Enter data.r.r = _
    query:1       | Eval data.r.r = _
    query:1       | Index data.r.r (matched 2 rules, early exit)
    r.rego:11     | Enter data.r.r
    r.rego:12     | | Eval input.y = "y"
    r.rego:11     | | Exit data.r.r
    query:1       | Exit data.r.r = _
    query:1       Redo data.r.r = _
    query:1       | Redo data.r.r = _
    r.rego:11     | Redo data.r.r
    r.rego:12     | | Redo input.y = "y"
    r.rego:11     | Exit data.r.r early

With `r.rego` as

    package r
    r {
      input.x = "x"
    }
    r = 2 {
      input.z = "z"
    }
    r {
      input.y = "y"
    }

This is done in in a way such that early-exit will abort array/set/object
iterations on data:

    r {
      data.i[_] = "one"
      data.j[_] = "four"
    }

    f(x, y) {
      data.i[_] = x
      data.j[_] = y
    }

Complete rules (r) and functions (f) that iterate over sets, arrays, and
objects from either data (evalTree) or a term that's returned by some
other rule etc (evalTerm).

The CLI and golang packages expose ways to disable 'early-exit':

This is in line with how indexing can be disabled. It's supposed to be
used as a debugging measure, so it's only exposed as a CLI flag to
`opa eval`.

Co-authored-by: Torin Sandall <[email protected]>
Signed-off-by: Stephan Renatus <[email protected]>

ast/index.go

@@ -32,10 +32,11 @@ type RuleIndex interface {
 
 // IndexResult contains the result of an index lookup.
 type IndexResult struct {
-	Kind    DocKind
-	Rules   []*Rule
-	Else    map[*Rule][]*Rule
-	Default *Rule
+	Kind      DocKind
+	Rules     []*Rule
+	Else      map[*Rule][]*Rule
+	Default   *Rule
+	EarlyExit bool
 }
 
 // NewIndexResult returns a new IndexResult object.
@@ -114,13 +115,11 @@ func (i *baseDocEqIndex) Build(rules []*Rule) bool {
 			// Insert rule into trie with (insertion order, priority order)
 			// tuple. Retaining the insertion order allows us to return rules
 			// in the order they were passed to this function.
-			node.rules = append(node.rules, &ruleNode{[...]int{idx, prio}, rule})
+			node.append([...]int{idx, prio}, rule)
 			prio++
 			return false
 		})
-
 	}
-
 	return true
 }
 
@@ -143,6 +142,7 @@ func (i *baseDocEqIndex) Lookup(resolver ValueResolver) (*IndexResult, error) {
 		})
 		nodes := tr.unordered[pos]
 		root := nodes[0].rule
+
 		result.Rules = append(result.Rules, root)
 		if len(nodes) > 1 {
 			result.Else[root] = make([]*Rule, len(nodes)-1)
@@ -152,6 +152,8 @@ func (i *baseDocEqIndex) Lookup(resolver ValueResolver) (*IndexResult, error) {
 		}
 	}
 
+	result.EarlyExit = tr.values.Len() == 1 && tr.values.Slice()[0].IsGround()
+
 	return result, nil
 }
 
@@ -181,6 +183,8 @@ func (i *baseDocEqIndex) AllRules(resolver ValueResolver) (*IndexResult, error)
 		}
 	}
 
+	result.EarlyExit = tr.values.Len() == 1 && tr.values.Slice()[0].IsGround()
+
 	return result, nil
 }
 
@@ -190,9 +194,7 @@ type ruleWalker struct {
 
 func (r *ruleWalker) Do(x interface{}) trieWalker {
 	tn := x.(*trieNode)
-	for _, rn := range tn.rules {
-		r.result.Add(rn)
-	}
+	r.result.Add(tn)
 	return r
 }
 
@@ -397,25 +399,33 @@ type trieWalker interface {
 type trieTraversalResult struct {
 	unordered map[int][]*ruleNode
 	ordering  []int
+	values    Set
 }
 
 func newTrieTraversalResult() *trieTraversalResult {
 	return &trieTraversalResult{
 		unordered: map[int][]*ruleNode{},
+		values:    NewSet(),
 	}
 }
 
-func (tr *trieTraversalResult) Add(node *ruleNode) {
-	root := node.prio[0]
-	nodes, ok := tr.unordered[root]
-	if !ok {
-		tr.ordering = append(tr.ordering, root)
+func (tr *trieTraversalResult) Add(t *trieNode) {
+	for _, node := range t.rules {
+		root := node.prio[0]
+		nodes, ok := tr.unordered[root]
+		if !ok {
+			tr.ordering = append(tr.ordering, root)
+		}
+		tr.unordered[root] = append(nodes, node)
+	}
+	if t.values != nil {
+		t.values.Foreach(func(v *Term) { tr.values.Add(v) })
 	}
-	tr.unordered[root] = append(nodes, node)
 }
 
 type trieNode struct {
 	ref       Ref
+	values    Set
 	mappers   []*valueMapper
 	next      *trieNode
 	any       *trieNode
@@ -457,9 +467,25 @@ func (node *trieNode) String() string {
 	if len(node.mappers) > 0 {
 		flags = append(flags, fmt.Sprintf("%d mapper(s)", len(node.mappers)))
 	}
+	if l := node.values.Len(); l > 0 {
+		flags = append(flags, fmt.Sprintf("%d value(s)", l))
+	}
 	return strings.Join(flags, " ")
 }
 
+func (node *trieNode) append(prio [2]int, rule *Rule) {
+	node.rules = append(node.rules, &ruleNode{prio, rule})
+
+	if node.values != nil {
+		node.values.Add(rule.Head.Value)
+		return
+	}
+
+	if node.values == nil && rule.Head.DocKind() == CompleteDoc {
+		node.values = NewSet(rule.Head.Value)
+	}
+}
+
 type ruleNode struct {
 	prio [2]int
 	rule *Rule
@@ -513,9 +539,7 @@ func (node *trieNode) Traverse(resolver ValueResolver, tr *trieTraversalResult)
 		return nil
 	}
 
-	for i := range node.rules {
-		tr.Add(node.rules[i])
-	}
+	tr.Add(node)
 
 	return node.next.traverse(resolver, tr)
 }