diff --git a/v1/topdown/eval.go b/v1/topdown/eval.go index 01fc104e05..b6700549de 100644 --- a/v1/topdown/eval.go +++ b/v1/topdown/eval.go @@ -4204,36 +4204,64 @@ func isIterableValue(x ast.Value) bool { } func (e *evalEvery) save(iter unifyIterator) error { - return e.e.saveExpr(e.plug(e.expr), e.e.bindings, iter) + plugged, err := e.plug(e.expr) + if err != nil { + return err + } + return e.e.saveExpr(plugged, e.e.bindings, iter) } -func (e *evalEvery) plug(expr *ast.Expr) *ast.Expr { +func (e *evalEvery) plug(expr *ast.Expr) (*ast.Expr, error) { cpy := expr.Copy() every := cpy.Terms.(*ast.Every) - e.plugBody(every.Body) + if err := e.plugBody(every.Body); err != nil { + return nil, err + } 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 cpy + return cpy, nil } -func (e *evalEvery) plugBody(body ast.Body) { +func (e *evalEvery) plugBody(body ast.Body) error { for i := range body { switch t := body[i].Terms.(type) { case *ast.Term: - body[i].Terms = e.e.bindings.PlugNamespaced(t, e.e.caller.bindings) + plugged, err := e.plugTerm(t) + if err != nil { + return err + } + body[i].Terms = plugged case []*ast.Term: 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) + plugged, err := e.plugTerm(t[j]) + if err != nil { + return err + } + t[j] = plugged } case *ast.Every: - body[i] = e.plug(body[i]) + plugged, err := e.plug(body[i]) + if err != nil { + return err + } + body[i] = plugged case *ast.Not: - e.plugBody(t.Body) + if err := e.plugBody(t.Body); err != nil { + return err + } } } + return nil +} + +func (e *evalEvery) plugTerm(t *ast.Term) (*ast.Term, error) { + if ast.IsComprehension(t.Value) { + return e.e.amendComprehension(t, e.e.bindings) + } + return e.e.bindings.PlugNamespaced(t, e.e.caller.bindings), nil } type evalNot struct { diff --git a/v1/topdown/topdown_partial_test.go b/v1/topdown/topdown_partial_test.go index 0a4cf81f87..ce9defb093 100644 --- a/v1/topdown/topdown_partial_test.go +++ b/v1/topdown/topdown_partial_test.go @@ -3051,6 +3051,142 @@ func TestTopDownPartialEval(t *testing.T) { a2 = input; __local4__2 > __local2__2 } }`}, }, + { + note: "every: in-scope var in set comprehension", + query: "data.test.p = true", + modules: []string{`package test + p if { + every a in input.z { + a == 1 + x := input.x + {y | y := input.y; y < x} + } + }`}, + wantQueries: []string{`every __local0__1, __local1__1 in input.z { + __local1__1 = 1 + __local2__1 = input.x + {__local3__1 | __local3__1 = input.y; lt(__local3__1, __local2__1)} + }`}, + }, + { + note: "every: in-scope var in nested set comprehensions", + query: "data.test.p = true", + modules: []string{`package test + p if { + every x in input.x { + {y | + y := input.y + y < x + {z | + z := input.z + z < x + z > y + } + } + } + }`}, + wantQueries: []string{`every __local0__1, __local1__1 in input.x { + {__local2__1 | + __local2__1 = input.y + lt(__local2__1, __local1__1) + {__local3__1 | + __local3__1 = input.z + lt(__local3__1, __local1__1) + gt(__local3__1, __local2__1)} + } + }`}, + }, + { + note: "every: in-scope var in set comprehension, assigned", + query: "data.test.p = true", + modules: []string{`package test + p if { + every a in input.z { + a == 1 + x := input.x + z := {y | y := input.y; y < x} + z != {x} + } + }`}, + wantQueries: []string{`every __local0__1, __local1__1 in input.z { + __local1__1 = 1 + __local2__1 = input.x + __local4__1 = {__local3__1 | __local3__1 = input.y; lt(__local3__1, __local2__1)} + neq(__local4__1, {__local2__1}) + }`}, + }, + { + note: "every: in-scope var in array comprehension", + query: "data.test.p = true", + modules: []string{`package test + p if { + every a in input.z { + a == 1 + x := input.x + [y | y := input.y; y < x] + } + }`}, + wantQueries: []string{`every __local0__1, __local1__1 in input.z { + __local1__1 = 1 + __local2__1 = input.x + [__local3__1 | __local3__1 = input.y; lt(__local3__1, __local2__1)] + }`}, + }, + { + note: "every: in-scope var in array comprehension, assigned", + query: "data.test.p = true", + modules: []string{`package test + p if { + every a in input.z { + a == 1 + x := input.x + z := [y | y := input.y; y < x] + z != [x] + } + }`}, + wantQueries: []string{`every __local0__1, __local1__1 in input.z { + __local1__1 = 1 + __local2__1 = input.x + __local4__1 = [__local3__1 | __local3__1 = input.y; lt(__local3__1, __local2__1)] + neq(__local4__1, [__local2__1]) + }`}, + }, + { + note: "every: in-scope var in object comprehension", + query: "data.test.p = true", + modules: []string{`package test + p if { + every a in input.z { + a == 1 + x := input.x + {k: v | k := input.k; v := x} + } + }`}, + wantQueries: []string{`every __local0__1, __local1__1 in input.z { + __local1__1 = 1 + __local2__1 = input.x + {__local3__1: __local4__1 | __local3__1 = input.k; __local4__1 = __local2__1} + }`}, + }, + { + note: "every: in-scope var in object comprehension, assigned", + query: "data.test.p = true", + modules: []string{`package test + p if { + every a in input.z { + a == 1 + x := input.x + z := {k: v | k := input.k; v := x} + z != {"_": x} + } + }`}, + wantQueries: []string{`every __local0__1, __local1__1 in input.z { + __local1__1 = 1 + __local2__1 = input.x + __local5__1 = {__local3__1: __local4__1 | __local3__1 = input.k; __local4__1 = __local2__1} + neq(__local5__1, {"_": __local2__1}) + }`}, + }, { // https://github.com/open-policy-agent/opa/issues/5367 note: "copypropagation: keep equations that are only found in comprehensions, inlined function call", query: "data.test.p",