ast: Disqualify indexing if candidate found in nested comprehension

ast/compile.go

@@ -1601,27 +1601,38 @@ func getComprehensionIndex(arity func(Ref) int, candidates VarSet, expr *Expr) *
 
 	outputs := outputVarsForBody(body, arity, ReservedVars)
 	unsafe := body.Vars(SafetyCheckVisitorParams).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 {
+	regressionVis := newComprehensionIndexRegressionCheckVisitor(candidates)
+	regressionVis.Walk(body)
+	if regressionVis.worse {
 		return nil
 	}
 
-	indexVars := candidates.Intersect(outputs)
-	if len(indexVars) == 0 {
+	// Check if any nested comprehensions close over candidates. If any intersection is found
+	// the comprehension cannot be cached because it would require closing over the candidates
+	// which the evaluator does not support today.
+	nestedVis := newComprehensionIndexNestedCandidateVisitor(candidates)
+	nestedVis.Walk(body)
+	if nestedVis.found {
 		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.
+	// if we can decide that one ordering is better than another. If the set is
+	// empty, there is no indexing to do.
+	indexVars := candidates.Intersect(outputs)
+	if len(indexVars) == 0 {
+		return nil
+	}
+
 	result := make([]*Term, 0, len(indexVars))
 
 	for v := range indexVars {
@@ -1690,6 +1701,38 @@ func (vis *comprehensionIndexRegressionCheckVisitor) assertEmptyIntersection(vs
 	}
 }
 
+type comprehensionIndexNestedCandidateVisitor struct {
+	candidates VarSet
+	nested     bool
+	found      bool
+}
+
+func newComprehensionIndexNestedCandidateVisitor(candidates VarSet) *comprehensionIndexNestedCandidateVisitor {
+	return &comprehensionIndexNestedCandidateVisitor{
+		candidates: candidates,
+	}
+}
+
+func (vis *comprehensionIndexNestedCandidateVisitor) Walk(x interface{}) {
+	NewGenericVisitor(vis.visit).Walk(x)
+}
+
+func (vis *comprehensionIndexNestedCandidateVisitor) visit(x interface{}) bool {
+
+	if vis.found {
+		return true
+	}
+
+	if v, ok := x.(Value); ok && IsComprehension(v) {
+		varVis := NewVarVisitor().WithParams(VarVisitorParams{SkipRefHead: true})
+		varVis.Walk(v)
+		vis.found = len(varVis.Vars().Intersect(vis.candidates)) > 0
+		return true
+	}
+
+	return false
+}
+
 // ModuleTreeNode represents a node in the module tree. The module
 // tree is keyed by the package path.
 type ModuleTreeNode struct {

ast/compile_test.go

@@ -3526,6 +3526,21 @@ func TestCompilerBuildComprehensionIndexKeySet(t *testing.T) {
 					ys = [y | y = input[j]]
 				}`,
 		},
+		{
+			note: "skip: due to nested comprehension containing candidate",
+			module: `
+				package test
+
+				p {
+					x = input[i]                # 'x' is a candidate
+					y = 2                       # 'y' is a candidate
+					z = [1 |
+						x = data.foo[j]             # 'x' is an index key
+						t = [1 | data.bar[k] = y]   # 'y' disqualifies indexing because it is nested inside a comprehension
+					]
+				}
+			`,
+		},
 		{
 			note: "skip: avoid increasing runtime (func arg)",
 			module: `

docs/content/policy-performance.md

@@ -228,7 +228,7 @@ In order to be indexed, comprehensions must meet the following conditions:
 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.
+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 or inside nested comprehensions.
 
 The following examples show cases that are NOT indexed:
 
@@ -262,6 +262,14 @@ not_indexed_because_reference_operand_closure {
   x := input[y].x
   ys := [y | x == input[y].z[_]]
 }
+
+not_indexed_because_nested_closure {
+  x = 1
+  y = 2
+  _ = [i |
+    x == input.foo[i]
+    _ = [j | y == input.bar[j]]]
+}
 ```
 
 > The 4th and 5th restrictions may be relaxed in the future.