ast, topdown: Index comprehensions to avoid unnecessary work

This commit adds a new kind of indexing to the compiler and topdown to
help avoid recomputing comprehensions. This helps with queries that
perform "group by" operations.

This optimization allows policies to perform group-by/aggregation in
O(n) instead of O(n^2). The optimization works by computing a set of
index keys for the comprehension at compile-time and then computing
the collection once at evaluation-time and indexing the result based
on the keys.

The index keys are variables in the outer query that limit the values
produced by the comprehension. In the simple group-by case these are
the object values themselves. During evaluation, topdown checks if
indexing is possible and builds the index by computing the
comprehension without creating a closure over the outer query. This
computes ALL values in the collection defined by the
comprehension. The results are keyed by the assignments to the
variables indicated in the comprehension index. This way the
comprehension does not have to be recomputed for each set of
assignments in the outer query.

The index is exposed on both the compiler and the query compiler so
that ad-hoc queries can benefit from the indexing as well. This is
important for things like the playground where users may select a rule
body and run it. If that exhibited n^2 behaviour it would be quite
confusing.

In order to be indexed, the comprehension must meet a few
conditions. Importantly, the indexing should not worsen overall
performance. To ensure this, comprehensions containing refs or walk()
calls that include output vars that close over the outer query are not
indexed. This means that if the caller were pushing down assignments
to those vars, OPA will not compute the entire collection.

In the future we can improve the index to cover more kinds of
comprehensions. One improvement that would be particularly nice would
be to allow the comprehension index to close over specific local
variables in the parent scope. This would let us build the index in
more cases--however, the analysis would need to be careful to take
into account the count of closure variables. Variables with multiple
assignments would be poor candidates.

Benchmark results (before, O(n^2) runtime):

BenchmarkComprehensionIndexing/10-16 	   13831	     85821 ns/op
BenchmarkComprehensionIndexing/100-16         	     208	   5662625 ns/op
BenchmarkComprehensionIndexing/1000-16        	       2	 549295038 ns/op

Benchmark results (after, O(n) runtime):

BenchmarkComprehensionIndexing/10-16 	   35809	     33369 ns/op
BenchmarkComprehensionIndexing/100-16         	    3756	    274546 ns/op
BenchmarkComprehensionIndexing/1000-16        	     438	   2725152 ns/op

Fixes #2276

Signed-off-by: Torin Sandall <torinsandall@gmail.com>

ast/compile.go

@@ -90,13 +90,14 @@ type Compiler struct {
 		metricName string
 		f          func()
 	}
-	maxErrs           int
-	sorted            []string // list of sorted module names
-	pathExists        func([]string) (bool, error)
-	after             map[string][]CompilerStageDefinition
-	metrics           metrics.Metrics
-	builtins          map[string]*Builtin
-	unsafeBuiltinsMap map[string]struct{}
+	maxErrs              int
+	sorted               []string // list of sorted module names
+	pathExists           func([]string) (bool, error)
+	after                map[string][]CompilerStageDefinition
+	metrics              metrics.Metrics
+	builtins             map[string]*Builtin
+	unsafeBuiltinsMap    map[string]struct{}
+	comprehensionIndices map[*Term]*ComprehensionIndex
 }
 
 // CompilerStage defines the interface for stages in the compiler.
@@ -187,6 +188,10 @@ type QueryCompiler interface {
 	// parsed query. For example, given the query "input := 1" the rewritten
 	// query would be "__local0__ = 1". The mapping would then be {__local0__: input}.
 	RewrittenVars() map[Var]Var
+
+	// ComprehensionIndex returns an index data structure for the given comprehension
+	// term. If no index is found, returns nil.
+	ComprehensionIndex(term *Term) *ComprehensionIndex
 }
 
 // QueryCompilerStage defines the interface for stages in the query compiler.
@@ -214,9 +219,10 @@ func NewCompiler() *Compiler {
 		}, func(x util.T) int {
 			return x.(Ref).Hash()
 		}),
-		maxErrs:           CompileErrorLimitDefault,
-		after:             map[string][]CompilerStageDefinition{},
-		unsafeBuiltinsMap: map[string]struct{}{},
+		maxErrs:              CompileErrorLimitDefault,
+		after:                map[string][]CompilerStageDefinition{},
+		unsafeBuiltinsMap:    map[string]struct{}{},
+		comprehensionIndices: map[*Term]*ComprehensionIndex{},
 	}
 
 	c.ModuleTree = NewModuleTree(nil)
@@ -261,6 +267,7 @@ func NewCompiler() *Compiler {
 		{"CheckTypes", "compile_stage_check_types", c.checkTypes},
 		{"CheckUnsafeBuiltins", "compile_state_check_unsafe_builtins", c.checkUnsafeBuiltins},
 		{"BuildRuleIndices", "compile_stage_rebuild_indices", c.buildRuleIndices},
+		{"BuildComprehensionIndices", "compile_stage_rebuild_comprehension_indices", c.buildComprehensionIndices},
 	}
 
 	return c
@@ -350,6 +357,13 @@ func (c *Compiler) Failed() bool {
 	return len(c.Errors) > 0
 }
 
+// ComprehensionIndex returns a data structure specifying how to index comprehension
+// results so that callers do not have to recompute the comprehension more than once.
+// If no index is found, returns nil.
+func (c *Compiler) ComprehensionIndex(term *Term) *ComprehensionIndex {
+	return c.comprehensionIndices[term]
+}
+
 // GetArity returns the number of args a function referred to by ref takes. If
 // ref refers to built-in function, the built-in declaration is consulted,
 // otherwise, the ref is used to perform a ruleset lookup.
@@ -610,7 +624,13 @@ func (c *Compiler) WithModuleLoader(f ModuleLoader) *Compiler {
 	return c
 }
 
-// buildRuleIndices constructs indices for rules.
+func (c *Compiler) counterAdd(name string, n uint64) {
+	if c.metrics == nil {
+		return
+	}
+	c.metrics.Counter(name).Add(n)
+}
+
 func (c *Compiler) buildRuleIndices() {
 
 	c.RuleTree.DepthFirst(func(node *TreeNode) bool {
@@ -628,6 +648,18 @@ func (c *Compiler) buildRuleIndices() {
 
 }
 
+func (c *Compiler) buildComprehensionIndices() {
+	for _, name := range c.sorted {
+		WalkRules(c.Modules[name], func(r *Rule) bool {
+			candidates := r.Head.Args.Vars()
+			candidates.Update(ReservedVars)
+			n := buildComprehensionIndices(c.GetArity, candidates, r.Body, c.comprehensionIndices)
+			c.counterAdd(compileStageComprehensionIndexBuild, n)
+			return false
+		})
+	}
+}
+
 // checkRecursion ensures that there are no recursive definitions, i.e., there are
 // no cycles in the Graph.
 func (c *Compiler) checkRecursion() {
@@ -1250,19 +1282,21 @@ func (c *Compiler) setGraph() {
 }
 
 type queryCompiler struct {
-	compiler       *Compiler
-	qctx           *QueryContext
-	typeEnv        *TypeEnv
-	rewritten      map[Var]Var
-	after          map[string][]QueryCompilerStageDefinition
-	unsafeBuiltins map[string]struct{}
+	compiler             *Compiler
+	qctx                 *QueryContext
+	typeEnv              *TypeEnv
+	rewritten            map[Var]Var
+	after                map[string][]QueryCompilerStageDefinition
+	unsafeBuiltins       map[string]struct{}
+	comprehensionIndices map[*Term]*ComprehensionIndex
 }
 
 func newQueryCompiler(compiler *Compiler) QueryCompiler {
 	qc := &queryCompiler{
-		compiler: compiler,
-		qctx:     nil,
-		after:    map[string][]QueryCompilerStageDefinition{},
+		compiler:             compiler,
+		qctx:                 nil,
+		after:                map[string][]QueryCompilerStageDefinition{},
+		comprehensionIndices: map[*Term]*ComprehensionIndex{},
 	}
 	return qc
 }
@@ -1286,6 +1320,15 @@ func (qc *queryCompiler) RewrittenVars() map[Var]Var {
 	return qc.rewritten
 }
 
+func (qc *queryCompiler) ComprehensionIndex(term *Term) *ComprehensionIndex {
+	if result, ok := qc.comprehensionIndices[term]; ok {
+		return result
+	} else if result, ok := qc.compiler.comprehensionIndices[term]; ok {
+		return result
+	}
+	return nil
+}
+
 func (qc *queryCompiler) runStage(metricName string, qctx *QueryContext, query Body, s func(*QueryContext, Body) (Body, error)) (Body, error) {
 	if qc.compiler.metrics != nil {
 		qc.compiler.metrics.Timer(metricName).Start()
@@ -1321,6 +1364,7 @@ func (qc *queryCompiler) Compile(query Body) (Body, error) {
 		{"RewriteDynamicTerms", "query_compile_stage_rewrite_dynamic_terms", qc.rewriteDynamicTerms},
 		{"CheckTypes", "query_compile_stage_check_types", qc.checkTypes},
 		{"CheckUnsafeBuiltins", "query_compile_stage_check_unsafe_builtins", qc.checkUnsafeBuiltins},
+		{"BuildComprehensionIndex", "query_compile_stage_build_comprehension_index", qc.buildComprehensionIndices},
 	}
 
 	qctx := qc.qctx.Copy()
@@ -1462,6 +1506,189 @@ func (qc *queryCompiler) rewriteWithModifiers(qctx *QueryContext, body Body) (Bo
 	return body, nil
 }
 
+func (qc *queryCompiler) buildComprehensionIndices(qctx *QueryContext, body Body) (Body, error) {
+	// NOTE(tsandall): The query compiler does not have a metrics object so we
+	// cannot record index metrics currently.
+	_ = buildComprehensionIndices(qc.compiler.GetArity, ReservedVars, body, qc.comprehensionIndices)
+	return body, nil
+}
+
+// ComprehensionIndex specifies how the comprehension term can be indexed. The keys
+// tell the evaluator what variables to use for indexing. In the future, the index
+// could be expanded with more information that would allow the evaluator to index
+// a larger fragment of comprehensions (e.g., by closing over variables in the outer
+// query.)
+type ComprehensionIndex struct {
+	Term *Term
+	Keys []*Term
+}
+
+func (ci *ComprehensionIndex) String() string {
+	if ci == nil {
+		return ""
+	}
+	return fmt.Sprintf("<keys: %v>", Array(ci.Keys))
+}
+
+func buildComprehensionIndices(arity func(Ref) int, candidates VarSet, node interface{}, result map[*Term]*ComprehensionIndex) (n uint64) {
+	WalkBodies(node, func(b Body) bool {
+		cpy := candidates.Copy()
+		for _, expr := range b {
+			if index := getComprehensionIndex(arity, cpy, expr); index != nil {
+				result[index.Term] = index
+				n++
+			}
+			// Any variables appearing in the expressions leading up to the comprehension
+			// are fair-game to be used as index keys.
+			cpy.Update(expr.Vars(VarVisitorParams{SkipClosures: true, SkipRefCallHead: true}))
+		}
+		return false
+	})
+	return n
+}
+
+func getComprehensionIndex(arity func(Ref) int, candidates VarSet, expr *Expr) *ComprehensionIndex {
+
+	// Ignore everything except <var> = <comprehension> expressions. Extract
+	// the comprehension term from the expression.
+	if !expr.IsEquality() || expr.Negated || len(expr.With) > 0 {
+		return nil
+	}
+
+	var term *Term
+
+	lhs, rhs := expr.Operand(0), expr.Operand(1)
+
+	if _, ok := lhs.Value.(Var); ok && IsComprehension(rhs.Value) {
+		term = rhs
+	} else if _, ok := rhs.Value.(Var); ok && IsComprehension(lhs.Value) {
+		term = lhs
+	}
+
+	if term == nil {
+		return nil
+	}
+
+	// Ignore comprehensions that contain expressions that close over variables
+	// in the outer body if those variables are not also output variables in the
+	// comprehension body. In other words, ignore comprehensions that we cannot
+	// safely evaluate without bindings from the outer body. For example:
+	//
+	// 	x = [1]
+	//	[true | data.y[z] = x]     # safe to evaluate w/o outer body
+	//	[true | data.y[z] = x[0]]  # NOT safe to evaluate because 'x' would be unsafe.
+	//
+	// By identifying output variables in the body we also know what to index on by
+	// intersecting with candidate variables from the outer query.
+	//
+	// For example:
+	//
+	//	x = data.foo[_]
+	//	_ = [y | data.bar[y] = x]      # index on 'x'
+	//
+	// This query goes from O(data.foo*data.bar) to O(data.foo+data.bar).
+	var body Body
+
+	switch x := term.Value.(type) {
+	case *ArrayComprehension:
+		body = x.Body
+	case *SetComprehension:
+		body = x.Body
+	case *ObjectComprehension:
+		body = x.Body
+	}
+
+	outputs := outputVarsForBody(body, arity, ReservedVars)
+	unsafe := body.Vars(safetyCheckVarVisitorParams).Diff(outputs).Diff(ReservedVars)
+	if len(unsafe) > 0 {
+		return nil
+	}
+
+	// Similarly, ignore comprehensions that contain references with output variables
+	// that intersect with the candidates. Indexing these comprehensions could worsen
+	// performance.
+	vis := newComprehensionIndexRegressionCheckVisitor(candidates)
+	vis.Walk(body)
+	if vis.worse {
+		return nil
+	}
+
+	indexVars := candidates.Intersect(outputs)
+	if len(indexVars) == 0 {
+		return nil
+	}
+
+	// Make a sorted set of variable names that will serve as the index key set.
+	// Sort to ensure deterministic indexing. In future this could be relaxed
+	// if we can decide that one ordering is better than another.
+	result := make([]*Term, 0, len(indexVars))
+
+	for v := range indexVars {
+		result = append(result, NewTerm(v))
+	}
+
+	sort.Slice(result, func(i, j int) bool {
+		return result[i].Value.Compare(result[j].Value) < 0
+	})
+
+	return &ComprehensionIndex{Term: term, Keys: result}
+}
+
+type comprehensionIndexRegressionCheckVisitor struct {
+	candidates VarSet
+	seen       VarSet
+	worse      bool
+}
+
+// TOOD(tsandall): Improve this so that users can either supply this list explicitly
+// or the information is maintained on the built-in function declaration. What we really
+// need to know is whether the built-in function allows callers to push down output
+// values or not. It's unlikely that anything outside of OPA does this today so this
+// solution is fine for now.
+var comprehensionIndexBlacklist = map[string]int{
+	WalkBuiltin.Name: len(WalkBuiltin.Decl.Args()),
+}
+
+func newComprehensionIndexRegressionCheckVisitor(candidates VarSet) *comprehensionIndexRegressionCheckVisitor {
+	return &comprehensionIndexRegressionCheckVisitor{
+		candidates: candidates,
+		seen:       NewVarSet(),
+	}
+}
+
+func (vis *comprehensionIndexRegressionCheckVisitor) Walk(x interface{}) {
+	NewGenericVisitor(vis.visit).Walk(x)
+}
+
+func (vis *comprehensionIndexRegressionCheckVisitor) visit(x interface{}) bool {
+	if !vis.worse {
+		switch x := x.(type) {
+		case *Expr:
+			operands := x.Operands()
+			if pos := comprehensionIndexBlacklist[x.Operator().String()]; pos > 0 && pos < len(operands) {
+				vis.assertEmptyIntersection(operands[pos].Vars())
+			}
+		case Ref:
+			vis.assertEmptyIntersection(x.OutputVars())
+		case Var:
+			vis.seen.Add(x)
+		// Always skip comprehensions. We do not have to visit their bodies here.
+		case *ArrayComprehension, *SetComprehension, *ObjectComprehension:
+			return true
+		}
+	}
+	return vis.worse
+}
+
+func (vis *comprehensionIndexRegressionCheckVisitor) assertEmptyIntersection(vs VarSet) {
+	for v := range vs {
+		if vis.candidates.Contains(v) && !vis.seen.Contains(v) {
+			vis.worse = true
+			return
+		}
+	}
+}
+
 // ModuleTreeNode represents a node in the module tree. The module
 // tree is keyed by the package path.
 type ModuleTreeNode struct {

ast/compile_test.go

@@ -3377,6 +3377,225 @@ func TestCompilerWithStageAfterWithMetrics(t *testing.T) {
 	}
 }
 
+func TestCompilerBuildComprehensionIndexKeySet(t *testing.T) {
+
+	tests := []struct {
+		note     string
+		module   string
+		atRow    int
+		wantTerm string
+		wantKeys string
+	}{
+		{
+			note: "example: invert object",
+			module: `
+				package test
+
+				p {
+					value = input[i]
+					keys = [j | value = input[j]]
+				}
+			`,
+			atRow:    6,
+			wantTerm: `[j | value = input[j]]`,
+			wantKeys: `[value]`,
+		},
+		{
+			note: "example: multiple keys from body",
+			module: `
+				package test
+
+				p {
+					v1 = input[i].v1
+					v2 = input[i].v2
+					keys = [j | v1 = input[j].v1; v2 = input[j].v2]
+				}
+			`,
+			atRow:    7,
+			wantTerm: `[j | v1 = input[j].v1; v2 = input[j].v2]`,
+			wantKeys: `[v1, v2]`,
+		},
+		{
+			note: "example: nested comprehensions are supported",
+			module: `
+				package test
+
+				p = {x: ys |
+					x = input[i]
+					ys = {y | x = input[y]}
+				}
+			`,
+			atRow:    6,
+			wantTerm: `{y | x = input[y]}`,
+			wantKeys: `[x]`,
+		},
+		{
+			note: "skip: lone comprehensions",
+			module: `
+				package test
+
+				p {
+					[v | input[i] = v]  # skip because no assignment
+				}`,
+		},
+		{
+			note: "skip: due to with modifier",
+			module: `
+				package test
+
+				p {
+					v = input[i]
+					ks = [j | input[j] = v] with data.x as 1  # skip because of with modifier
+				}`,
+		},
+		{
+			note: "skip: due to negation",
+			module: `
+				package test
+
+				p {
+					v = input[i]
+					a = []
+					not a = [j | input[j] = v] # skip due to negation
+				}`,
+		},
+		{
+			note: "skip: due to lack of comprehension",
+			module: `
+				package test
+
+				p {
+					v = input[i]
+				}`,
+		},
+		{
+			note: "skip: due to unsafe comprehension body",
+			module: `
+				package test
+
+				f(x) {
+					v = input[i]
+					ys = [y | y = x[j]]  # x is not safe
+				}`,
+		},
+		{
+			note: "skip: due to no candidates",
+			module: `
+				package test
+
+				p {
+					ys = [y | y = input[j]]
+				}`,
+		},
+		{
+			note: "skip: avoid increasing runtime (func arg)",
+			module: `
+				package test
+
+				f(x) {
+					y = input[x]
+					ys = [y | y = input[x]]
+				}`,
+		},
+		{
+			note: "skip: avoid increasing runtime (head key)",
+			module: `
+				package test
+
+				p[x] {
+					y = input[x]
+					ys = [y | y = input[x]]
+				}`,
+		},
+		{
+			note: "skip: avoid increasing runtime (walk)",
+			module: `
+				package test
+
+				p[x] {
+					y = input.bar[x]
+					ys = [y | a = input.foo; walk(a, [x, y])]
+				}`,
+		},
+		{
+			note: "bypass: use intermediate var to skip regression check",
+			module: `
+				package test
+
+				p[x] {
+					y = input[x]
+					ys = [y | y = input[z]; z = x]
+				}`,
+			atRow:    6,
+			wantTerm: ` [y | y = input[z]; z = x]`,
+			wantKeys: `[x, y]`,
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.note, func(t *testing.T) {
+
+			m := metrics.New()
+			compiler := NewCompiler().WithMetrics(m)
+			compiler.Compile(map[string]*Module{"test.rego": MustParseModule(tc.module)})
+			if compiler.Failed() {
+				t.Fatal(compiler.Errors)
+			}
+
+			n := m.Counter(compileStageComprehensionIndexBuild).Value().(uint64)
+
+			if tc.atRow == 0 {
+				if n > 0 || len(compiler.comprehensionIndices) > 0 {
+					t.Fatal("expected no indices to be built. got:", compiler.comprehensionIndices)
+				}
+				return
+			}
+
+			if n != 1 {
+				t.Fatal("expected counter to be incremented")
+			}
+
+			var comprehension *Term
+
+			WalkTerms(compiler.Modules["test.rego"], func(x *Term) bool {
+				if !IsComprehension(x.Value) {
+					return true
+				}
+				if x.Location.Row != tc.atRow {
+					return false
+				} else if comprehension != nil {
+					t.Fatal("expected at most one comprehension per line in test module")
+				}
+				comprehension = x
+				return false
+			})
+
+			if comprehension == nil {
+				t.Fatal("expected comprehension at line:", tc.atRow)
+			}
+
+			result := compiler.ComprehensionIndex(comprehension)
+
+			if result == nil {
+				t.Fatal("expected result")
+			}
+
+			expTerm := MustParseTerm(tc.wantTerm)
+
+			if !result.Term.Equal(expTerm) {
+				t.Fatalf("expected term to be %v but got: %v", expTerm, result.Term)
+			}
+
+			expKeys := MustParseTerm(tc.wantKeys).Value.(Array)
+
+			if Array(result.Keys).Compare(expKeys) != 0 {
+				t.Fatalf("expected keys to be %v but got: %v", expKeys, result.Keys)
+			}
+
+		})
+	}
+}
+
 func TestQueryCompiler(t *testing.T) {
 	tests := []struct {
 		note     string

ast/compilemetrics.go

@@ -0,0 +1,9 @@
+// Copyright 2020 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 ast
+
+const (
+	compileStageComprehensionIndexBuild = "compile_stage_comprehension_index_build"
+)

docs/content/policy-performance.md

@@ -3,6 +3,7 @@ title: Policy Performance
 kind: documentation
 weight: 3
 ---
+
 ## High Performance Policy Decisions
 
 For low-latency/high-performance use-cases, e.g. microservice API authorization, policy evaluation has a budget on the order of 1 millisecond.  Not all use cases require that kind of performance, and OPA is powerful enough that you can write expressive policies that take longer than 1 millisecond to evaluate.  But for high-performance use cases, there is a fragment of the policy language that has been engineered to evaluate quickly.  Even as the size of the policies grow, the performance for this fragment can be nearly constant-time.
@@ -131,6 +132,140 @@ For `glob.match(pattern, delimiter, match)` statements to be indexed the pattern
 | `glob.match("foo:**:bar", [":"], input.x)` | no | pattern contains `**` |
 | `glob.match("foo:*:bar", [":"], input.x[i])` | no | match contains variable(s) |
 
+### Comprehension Indexing
+
+Rego does not support mutation. As a result, certain operations like "group by" require
+use of comprehensions to aggregate values. To avoid O(n^2) runtime complexity in
+queries/rules that perform group-by, OPA may compute and memoize the entire collection
+produced by comprehensions at once. This ensures that runtime complexity is O(n) where
+n is the size of the collection that group-by/aggregation is being performed on.
+
+For example, suppose the policy must check if the number of ports exposed on an interface
+exceeds some threshold (e.g., any interface may expose up to 100 ports.) The policy is given
+the port->interface mapping as a JSON array under `input`:
+
+```json
+{
+  "exposed": [
+    {
+      "interface": "eth0",
+      "port": 8080,
+    },
+    {
+      "interface": "eth0",
+      "port": 8081,
+    },
+    {
+      "interface": "eth1",
+      "port": 443,
+    },
+    {
+      "interface": "lo1",
+      "port": 5000,
+    }
+  ]
+}
+```
+
+In this case, the policy must count the number of ports exposed on each interface. To do this,
+the policy must first aggregate/group the ports by the interface name. Conceptually,
+the policy should generate a document like this:
+
+```json
+{
+  "exposed_ports_by_interface": {
+    "eth0": [8080, 8081],
+    "eth1": [443],
+    "lo1": [5000]
+  }
+}
+```
+
+Since multiple ports could be exposed on a single interface, the policy must use a comprehension to
+aggregate the port values by the interface names. To implement this logic in Rego, we would write:
+
+```rego
+some i
+intf := input.exposed[i].interface
+ports := [port | some j; input.exposed[j].interface == intf; port := input.exposed[j].port]
+```
+
+Without comprehension indexing, this query would be O(n^2) where n is the size of `input.exposed`.
+However, with comprehension indexing, the query remains O(n) because OPA only computes the comprehension
+_once_. In this case, the comprehension is evaluated and all possible values of `ports` are computed
+at once. These values are indexed by the assignments of `intf`.
+
+To implement the policy above we could write:
+
+```rego
+deny[msg] {
+  some i
+  count(exposed_ports_by_interface[i]) > 100
+  msg := sprintf("interface '%v' exposes too many ports", [i])
+}
+
+exposed_ports_by_interface := {intf: ports |
+  some i
+  intf := input.exposed[i].interface
+  ports := [port |
+    some j
+    input.exposed[j].interface == intf
+    port := input.exposed[j].port
+  ]
+}
+```
+
+Indices can be built for comprehensions (nested or not) that generate collections (i.e., arrays, sets, or objects)
+based on variables in an outer query. In the example above:
+
+* `intf` is the variable in the outer query.
+* `[port | some j; input.exposed[j].interface == intf; port := input.exposed[j].port]` is the comprehension.
+* `ports` is the variable the collection is assigned to.
+
+In order to be indexed, comprehensions must meet the following conditions:
+
+1. The comprehension appears in an assignment or unification statement.
+1. The expression containing the comprehension does not include a `with` statement.
+1. The expression containing the comprehension is not negated.
+1. The comprehension body is safe when considered independent from the outer query.
+1. The comprehension body closes over at least one variable in the outer query and none of these variables appear as outputs in references or `walk()` calls.
+
+The following examples show cases that are NOT indexed:
+
+```rego
+not_indexed_because_missing_assignment {
+  x := input[_]
+  [y | some y; x == input[y]]
+}
+
+not_indexed_because_includes_with {
+  x := input[_]
+  ys := [y | some y; x := input[y]] with input as {}
+}
+
+not_indexed_because_negated {
+  x := input[_]
+  not data.arr = [y | some y; x := input[y]]
+}
+
+not_indexed_because_safety {
+  obj := input.foo.bar
+  x := obj[_]
+  ys := [y | some y; x == obj[y]]
+}
+
+not_indexed_because_no_closure {
+  ys := [y | x := input[y]]
+}
+
+not_indexed_because_reference_operand_closure {
+  x := input[y].x
+  ys := [y | x == input[y].z[_]]
+}
+```
+
+> The 4th and 5th restrictions may be relaxed in the future.
+
 ### Profiling
 
 You can also _profile_ your policies using `opa eval`. The profiler is useful if you need to understand

topdown/cache.go

@@ -164,3 +164,74 @@ func (s *refStack) Prefixed(ref ast.Ref) bool {
 	}
 	return false
 }
+
+type comprehensionCache struct {
+	stack []map[*ast.Term]*comprehensionCacheElem
+}
+
+type comprehensionCacheElem struct {
+	value    *ast.Term
+	children *util.HashMap
+}
+
+func newComprehensionCache() *comprehensionCache {
+	cache := &comprehensionCache{}
+	cache.Push()
+	return cache
+}
+
+func (c *comprehensionCache) Push() {
+	c.stack = append(c.stack, map[*ast.Term]*comprehensionCacheElem{})
+}
+
+func (c *comprehensionCache) Pop() {
+	c.stack = c.stack[:len(c.stack)-1]
+}
+
+func (c *comprehensionCache) Elem(t *ast.Term) (*comprehensionCacheElem, bool) {
+	elem, ok := c.stack[len(c.stack)-1][t]
+	return elem, ok
+}
+
+func (c *comprehensionCache) Set(t *ast.Term, elem *comprehensionCacheElem) {
+	c.stack[len(c.stack)-1][t] = elem
+}
+
+func newComprehensionCacheElem() *comprehensionCacheElem {
+	return &comprehensionCacheElem{children: newComprehensionCacheHashMap()}
+}
+
+func (c *comprehensionCacheElem) Get(key []*ast.Term) *ast.Term {
+	node := c
+	for i := 0; i < len(key); i++ {
+		x, ok := node.children.Get(key[i])
+		if !ok {
+			return nil
+		}
+		node = x.(*comprehensionCacheElem)
+	}
+	return node.value
+}
+
+func (c *comprehensionCacheElem) Put(key []*ast.Term, value *ast.Term) {
+	node := c
+	for i := 0; i < len(key); i++ {
+		x, ok := node.children.Get(key[i])
+		if ok {
+			node = x.(*comprehensionCacheElem)
+		} else {
+			next := newComprehensionCacheElem()
+			node.children.Put(key[i], next)
+			node = next
+		}
+	}
+	node.value = value
+}
+
+func newComprehensionCacheHashMap() *util.HashMap {
+	return util.NewHashMap(func(a, b util.T) bool {
+		return a.(*ast.Term).Equal(b.(*ast.Term))
+	}, func(x util.T) int {
+		return x.(*ast.Term).Hash()
+	})
+}

topdown/errors.go

@@ -117,3 +117,11 @@ func mergeConflictErr(loc *ast.Location) error {
 		Message:  "real and replacement data could not be merged",
 	}
 }
+
+func internalErr(loc *ast.Location, msg string) error {
+	return &Error{
+		Code:     InternalErr,
+		Location: loc,
+		Message:  msg,
+	}
+}

topdown/eval.go

@@ -30,37 +30,38 @@ func (f *queryIDFactory) Next() uint64 {
 }
 
 type eval struct {
-	ctx             context.Context
-	seed            io.Reader
-	queryID         uint64
-	queryIDFact     *queryIDFactory
-	parent          *eval
-	caller          *eval
-	cancel          Cancel
-	query           ast.Body
-	queryCompiler   ast.QueryCompiler
-	index           int
-	indexing        bool
-	bindings        *bindings
-	store           storage.Store
-	baseCache       *baseCache
-	txn             storage.Transaction
-	compiler        *ast.Compiler
-	input           *ast.Term
-	data            *ast.Term
-	targetStack     *refStack
-	tracers         []Tracer
-	instr           *Instrumentation
-	builtins        map[string]*Builtin
-	builtinCache    builtins.Cache
-	virtualCache    *virtualCache
-	saveSet         *saveSet
-	saveStack       *saveStack
-	saveSupport     *saveSupport
-	saveNamespace   *ast.Term
-	disableInlining [][]ast.Ref
-	genvarprefix    string
-	runtime         *ast.Term
+	ctx                context.Context
+	seed               io.Reader
+	queryID            uint64
+	queryIDFact        *queryIDFactory
+	parent             *eval
+	caller             *eval
+	cancel             Cancel
+	query              ast.Body
+	queryCompiler      ast.QueryCompiler
+	index              int
+	indexing           bool
+	bindings           *bindings
+	store              storage.Store
+	baseCache          *baseCache
+	txn                storage.Transaction
+	compiler           *ast.Compiler
+	input              *ast.Term
+	data               *ast.Term
+	targetStack        *refStack
+	tracers            []Tracer
+	instr              *Instrumentation
+	builtins           map[string]*Builtin
+	builtinCache       builtins.Cache
+	virtualCache       *virtualCache
+	comprehensionCache *comprehensionCache
+	saveSet            *saveSet
+	saveStack          *saveStack
+	saveSupport        *saveSupport
+	saveNamespace      *ast.Term
+	disableInlining    [][]ast.Ref
+	genvarprefix       string
+	runtime            *ast.Term
 }
 
 func (e *eval) Run(iter evalIterator) error {
@@ -384,7 +385,6 @@ func (e *eval) evalWith(iter evalIterator) error {
 	}
 
 	input, err := mergeTermWithValues(e.input, pairsInput)
-
 	if err != nil {
 		return &Error{
 			Code:     ConflictErr,
@@ -432,6 +432,7 @@ func (e *eval) evalWithPush(input *ast.Term, data *ast.Term, targets []ast.Ref,
 		e.data = data
 	}
 
+	e.comprehensionCache.Push()
 	e.virtualCache.Push()
 	e.targetStack.Push(targets)
 	e.disableInlining = append(e.disableInlining, disable)
@@ -443,6 +444,7 @@ func (e *eval) evalWithPop(input *ast.Term, data *ast.Term) {
 	e.disableInlining = e.disableInlining[:len(e.disableInlining)-1]
 	e.targetStack.Pop()
 	e.virtualCache.Pop()
+	e.comprehensionCache.Pop()
 	e.data = data
 	e.input = input
 }
@@ -850,6 +852,14 @@ func (e *eval) biunifyComprehension(a, b *ast.Term, b1, b2 *bindings, swap bool,
 		return e.biunifyComprehensionPartial(a, b, b1, b2, swap, iter)
 	}
 
+	value, err := e.buildComprehensionCache(a)
+
+	if err != nil {
+		return err
+	} else if value != nil {
+		return e.biunify(value, b, b1, b2, iter)
+	}
+
 	switch a := a.Value.(type) {
 	case *ast.ArrayComprehension:
 		return e.biunifyComprehensionArray(a, b, b1, b2, iter)
@@ -859,7 +869,106 @@ func (e *eval) biunifyComprehension(a, b *ast.Term, b1, b2 *bindings, swap bool,
 		return e.biunifyComprehensionObject(a, b, b1, b2, iter)
 	}
 
-	return fmt.Errorf("illegal comprehension %T", a)
+	return internalErr(e.query[e.index].Location, "illegal comprehension type")
+}
+
+func (e *eval) buildComprehensionCache(a *ast.Term) (*ast.Term, error) {
+
+	index := e.comprehensionIndex(a)
+	if index == nil {
+		e.instr.counterIncr(evalOpComprehensionCacheSkip)
+		return nil, nil
+	}
+
+	cache, ok := e.comprehensionCache.Elem(a)
+	if !ok {
+		var err error
+		switch x := a.Value.(type) {
+		case *ast.ArrayComprehension:
+			cache, err = e.buildComprehensionCacheArray(x, index.Keys)
+		case *ast.SetComprehension:
+			cache, err = e.buildComprehensionCacheSet(x, index.Keys)
+		case *ast.ObjectComprehension:
+			cache, err = e.buildComprehensionCacheObject(x, index.Keys)
+		default:
+			err = internalErr(e.query[e.index].Location, "illegal comprehension type")
+		}
+		if err != nil {
+			return nil, err
+		}
+		e.comprehensionCache.Set(a, cache)
+		e.instr.counterIncr(evalOpComprehensionCacheBuild)
+	} else {
+		e.instr.counterIncr(evalOpComprehensionCacheHit)
+	}
+
+	values := make([]*ast.Term, len(index.Keys))
+
+	for i := range index.Keys {
+		values[i] = e.bindings.Plug(index.Keys[i])
+	}
+
+	return cache.Get(values), nil
+}
+
+func (e *eval) buildComprehensionCacheArray(x *ast.ArrayComprehension, keys []*ast.Term) (*comprehensionCacheElem, error) {
+	child := e.child(x.Body)
+	node := newComprehensionCacheElem()
+	return node, child.Run(func(child *eval) error {
+		values := make([]*ast.Term, len(keys))
+		for i := range keys {
+			values[i] = child.bindings.Plug(keys[i])
+		}
+		head := child.bindings.Plug(x.Term)
+		cached := node.Get(values)
+		if cached != nil {
+			cached.Value = append(cached.Value.(ast.Array), head)
+		} else {
+			node.Put(values, ast.ArrayTerm(head))
+		}
+		return nil
+	})
+}
+
+func (e *eval) buildComprehensionCacheSet(x *ast.SetComprehension, keys []*ast.Term) (*comprehensionCacheElem, error) {
+	child := e.closure(x.Body)
+	node := newComprehensionCacheElem()
+	return node, child.Run(func(child *eval) error {
+		values := make([]*ast.Term, len(keys))
+		for i := range keys {
+			values[i] = child.bindings.Plug(keys[i])
+		}
+		head := child.bindings.Plug(x.Term)
+		cached := node.Get(values)
+		if cached != nil {
+			set := cached.Value.(ast.Set)
+			set.Add(head)
+		} else {
+			node.Put(values, ast.SetTerm(head))
+		}
+		return nil
+	})
+}
+
+func (e *eval) buildComprehensionCacheObject(x *ast.ObjectComprehension, keys []*ast.Term) (*comprehensionCacheElem, error) {
+	child := e.closure(x.Body)
+	node := newComprehensionCacheElem()
+	return node, child.Run(func(child *eval) error {
+		values := make([]*ast.Term, len(keys))
+		for i := range keys {
+			values[i] = child.bindings.Plug(keys[i])
+		}
+		headKey := child.bindings.Plug(x.Key)
+		headValue := child.bindings.Plug(x.Value)
+		cached := node.Get(values)
+		if cached != nil {
+			obj := cached.Value.(ast.Object)
+			obj.Insert(headKey, headValue)
+		} else {
+			node.Put(values, ast.ObjectTerm(ast.Item(headKey, headValue)))
+		}
+		return nil
+	})
 }
 
 func (e *eval) biunifyComprehensionPartial(a, b *ast.Term, b1, b2 *bindings, swap bool, iter unifyIterator) error {
@@ -2227,6 +2336,13 @@ func (e evalTerm) save(iter unifyIterator) error {
 	return e.e.biunify(ast.NewTerm(ref), e.rterm, e.termbindings, e.rbindings, iter)
 }
 
+func (e *eval) comprehensionIndex(term *ast.Term) *ast.ComprehensionIndex {
+	if e.queryCompiler != nil {
+		return e.queryCompiler.ComprehensionIndex(term)
+	}
+	return e.compiler.ComprehensionIndex(term)
+}
+
 func applyCopyPropagation(p *copypropagation.CopyPropagator, instr *Instrumentation, body ast.Body) ast.Body {
 	instr.startTimer(partialOpCopyPropagation)
 	result := p.Apply(body)

topdown/instrumentation.go

@@ -7,18 +7,21 @@ package topdown
 import "github.com/open-policy-agent/opa/metrics"
 
 const (
-	evalOpPlug                  = "eval_op_plug"
-	evalOpResolve               = "eval_op_resolve"
-	evalOpRuleIndex             = "eval_op_rule_index"
-	evalOpBuiltinCall           = "eval_op_builtin_call"
-	evalOpVirtualCacheHit       = "eval_op_virtual_cache_hit"
-	evalOpVirtualCacheMiss      = "eval_op_virtual_cache_miss"
-	evalOpBaseCacheHit          = "eval_op_base_cache_hit"
-	evalOpBaseCacheMiss         = "eval_op_base_cache_miss"
-	partialOpSaveUnify          = "partial_op_save_unify"
-	partialOpSaveSetContains    = "partial_op_save_set_contains"
-	partialOpSaveSetContainsRec = "partial_op_save_set_contains_rec"
-	partialOpCopyPropagation    = "partial_op_copy_propagation"
+	evalOpPlug                    = "eval_op_plug"
+	evalOpResolve                 = "eval_op_resolve"
+	evalOpRuleIndex               = "eval_op_rule_index"
+	evalOpBuiltinCall             = "eval_op_builtin_call"
+	evalOpVirtualCacheHit         = "eval_op_virtual_cache_hit"
+	evalOpVirtualCacheMiss        = "eval_op_virtual_cache_miss"
+	evalOpBaseCacheHit            = "eval_op_base_cache_hit"
+	evalOpBaseCacheMiss           = "eval_op_base_cache_miss"
+	evalOpComprehensionCacheSkip  = "eval_op_comprehension_cache_skip"
+	evalOpComprehensionCacheBuild = "eval_op_comprehension_cache_build"
+	evalOpComprehensionCacheHit   = "eval_op_comprehension_cache_hit"
+	partialOpSaveUnify            = "partial_op_save_unify"
+	partialOpSaveSetContains      = "partial_op_save_set_contains"
+	partialOpSaveSetContainsRec   = "partial_op_save_set_contains_rec"
+	partialOpCopyPropagation      = "partial_op_copy_propagation"
 )
 
 // Instrumentation implements helper functions to instrument query evaluation

topdown/query.go

@@ -185,32 +185,33 @@ func (q *Query) PartialRun(ctx context.Context) (partials []ast.Body, support []
 	f := &queryIDFactory{}
 	b := newBindings(0, q.instr)
 	e := &eval{
-		ctx:           ctx,
-		seed:          q.seed,
-		cancel:        q.cancel,
-		query:         q.query,
-		queryCompiler: q.queryCompiler,
-		queryIDFact:   f,
-		queryID:       f.Next(),
-		bindings:      b,
-		compiler:      q.compiler,
-		store:         q.store,
-		baseCache:     newBaseCache(),
-		targetStack:   newRefStack(),
-		txn:           q.txn,
-		input:         q.input,
-		tracers:       q.tracers,
-		instr:         q.instr,
-		builtins:      q.builtins,
-		builtinCache:  builtins.Cache{},
-		virtualCache:  newVirtualCache(),
-		saveSet:       newSaveSet(q.unknowns, b, q.instr),
-		saveStack:     newSaveStack(),
-		saveSupport:   newSaveSupport(),
-		saveNamespace: ast.StringTerm(q.partialNamespace),
-		genvarprefix:  q.genvarprefix,
-		runtime:       q.runtime,
-		indexing:      q.indexing,
+		ctx:                ctx,
+		seed:               q.seed,
+		cancel:             q.cancel,
+		query:              q.query,
+		queryCompiler:      q.queryCompiler,
+		queryIDFact:        f,
+		queryID:            f.Next(),
+		bindings:           b,
+		compiler:           q.compiler,
+		store:              q.store,
+		baseCache:          newBaseCache(),
+		targetStack:        newRefStack(),
+		txn:                q.txn,
+		input:              q.input,
+		tracers:            q.tracers,
+		instr:              q.instr,
+		builtins:           q.builtins,
+		builtinCache:       builtins.Cache{},
+		virtualCache:       newVirtualCache(),
+		comprehensionCache: newComprehensionCache(),
+		saveSet:            newSaveSet(q.unknowns, b, q.instr),
+		saveStack:          newSaveStack(),
+		saveSupport:        newSaveSupport(),
+		saveNamespace:      ast.StringTerm(q.partialNamespace),
+		genvarprefix:       q.genvarprefix,
+		runtime:            q.runtime,
+		indexing:           q.indexing,
 	}
 
 	if len(q.disableInlining) > 0 {
@@ -285,28 +286,29 @@ func (q *Query) Iter(ctx context.Context, iter func(QueryResult) error) error {
 	}
 	f := &queryIDFactory{}
 	e := &eval{
-		ctx:           ctx,
-		seed:          q.seed,
-		cancel:        q.cancel,
-		query:         q.query,
-		queryCompiler: q.queryCompiler,
-		queryIDFact:   f,
-		queryID:       f.Next(),
-		bindings:      newBindings(0, q.instr),
-		compiler:      q.compiler,
-		store:         q.store,
-		baseCache:     newBaseCache(),
-		targetStack:   newRefStack(),
-		txn:           q.txn,
-		input:         q.input,
-		tracers:       q.tracers,
-		instr:         q.instr,
-		builtins:      q.builtins,
-		builtinCache:  builtins.Cache{},
-		virtualCache:  newVirtualCache(),
-		genvarprefix:  q.genvarprefix,
-		runtime:       q.runtime,
-		indexing:      q.indexing,
+		ctx:                ctx,
+		seed:               q.seed,
+		cancel:             q.cancel,
+		query:              q.query,
+		queryCompiler:      q.queryCompiler,
+		queryIDFact:        f,
+		queryID:            f.Next(),
+		bindings:           newBindings(0, q.instr),
+		compiler:           q.compiler,
+		store:              q.store,
+		baseCache:          newBaseCache(),
+		targetStack:        newRefStack(),
+		txn:                q.txn,
+		input:              q.input,
+		tracers:            q.tracers,
+		instr:              q.instr,
+		builtins:           q.builtins,
+		builtinCache:       builtins.Cache{},
+		virtualCache:       newVirtualCache(),
+		comprehensionCache: newComprehensionCache(),
+		genvarprefix:       q.genvarprefix,
+		runtime:            q.runtime,
+		indexing:           q.indexing,
 	}
 	e.caller = e
 	q.startTimer(metrics.RegoQueryEval)

topdown/topdown_test.go

@@ -2385,6 +2385,25 @@ func TestTopDownWithKeyword(t *testing.T) {
 			setl[x] { data.foo[x] }`},
 			rules: []string{`p = true { data.ex.setl[1] with data.foo as {1} }`},
 		},
+		{
+			// NOTE(tsandall): This case assumes that partial sets are not memoized.
+			// If we change that, it'll be harder to test that the comprehension
+			// cache is invalidated.
+			note: "invalidate comprehension cache",
+			exp:  `[[{"b": ["a", "c"]}], [{"b": ["a"]}]]`,
+			modules: []string{`package ex
+				s[x] {
+					x = {v: ks |
+						v = input[i]
+						ks = {k | v = input[k]}
+					}
+				}
+			`},
+			rules: []string{`p = [x, y] {
+				x = data.ex.s with input as {"a": "b", "c": "b"}
+				y = data.ex.s with input as {"a": "b"}
+			}`},
+		},
 	}
 
 	for _, tc := range tests {