topdown/pe: plug nested every expressions

topdown/eval.go

@@ -3039,9 +3039,12 @@ func (e evalEvery) eval(iter unifyIterator) error {
 }
 
 func (e *evalEvery) save(iter unifyIterator) error {
-	cpy := e.expr.Copy()
-	every := cpy.Terms.(*ast.Every)
+	return e.e.saveExpr(e.plug(e.expr), e.e.bindings, iter)
+}
 
+func (e *evalEvery) plug(expr *ast.Expr) *ast.Expr {
+	cpy := expr.Copy()
+	every := cpy.Terms.(*ast.Every)
 	for i := range every.Body {
 		switch t := every.Body[i].Terms.(type) {
 		case *ast.Term:
@@ -3050,15 +3053,16 @@ func (e *evalEvery) save(iter unifyIterator) error {
 			for j := 1; j < len(t); j++ { // don't plug operator, t[0]
 				t[j] = e.e.bindings.PlugNamespaced(t[j], e.e.caller.bindings)
 			}
+		case *ast.Every:
+			every.Body[i] = e.plug(every.Body[i])
 		}
 	}
 
 	every.Key = e.e.bindings.PlugNamespaced(every.Key, e.e.caller.bindings)
 	every.Value = e.e.bindings.PlugNamespaced(every.Value, e.e.caller.bindings)
 	every.Domain = e.e.bindings.PlugNamespaced(every.Domain, e.e.caller.bindings)
 	cpy.Terms = every
-
-	return e.e.saveExpr(cpy, e.e.bindings, iter)
+	return cpy
 }
 
 func (e *eval) comprehensionIndex(term *ast.Term) *ast.ComprehensionIndex {

topdown/topdown_partial_test.go

@@ -3196,6 +3196,27 @@ func TestTopDownPartialEval(t *testing.T) {
 				}`},
 			wantQueries: []string{`every __local1__2, __local2__2 in [1] { a2 = input; 1 = __local2__2 }`},
 		},
+		{
+			note:  "every: nested and closing over function args",
+			query: "data.test.p",
+			modules: []string{`package test
+				p {
+					f(input)
+				}
+				f(x) {
+					every y in [1] {
+						every z in [2] {
+							a = x
+							z > y
+						}
+					}
+				}`},
+			wantQueries: []string{`every __local1__2, __local2__2 in [1] {
+				__local6__2 = [2]
+				every __local3__2, __local4__2 in __local6__2 {
+					a2 = input; __local4__2 > __local2__2 }
+				}`},
+		},
 	}
 
 	ctx := context.Background()