diff --git a/topdown/copypropagation/copypropagation.go b/topdown/copypropagation/copypropagation.go index c2855cd612..4c6b8c42c7 100644 --- a/topdown/copypropagation/copypropagation.go +++ b/topdown/copypropagation/copypropagation.go @@ -387,34 +387,37 @@ type binding struct { func containedIn(value ast.Value, x interface{}) bool { var stop bool - switch v := value.(type) { - case ast.Ref: - ast.WalkTerms(x, func(t *ast.Term) bool { - switch t := t.Value.(type) { - case *ast.ArrayComprehension, *ast.ObjectComprehension, *ast.SetComprehension: - return true // skip closures - case ast.Ref: - if stop || t.HasPrefix(v) { - stop = true - return stop - } + + var vis *ast.GenericVisitor + vis = ast.NewGenericVisitor(func(x interface{}) bool { + switch x := x.(type) { + case *ast.Every: // skip body + vis.Walk(x.Key) + vis.Walk(x.Value) + vis.Walk(x.Domain) + return true + case *ast.ArrayComprehension, *ast.ObjectComprehension, *ast.SetComprehension: // skip + return true + case ast.Ref: + var match bool + if v, ok := value.(ast.Ref); ok { + match = x.HasPrefix(v) + } else { + match = x.Compare(value) == 0 } - return false - }) - default: - ast.WalkTerms(x, func(t *ast.Term) bool { - switch other := t.Value.(type) { - case *ast.ArrayComprehension, *ast.ObjectComprehension, *ast.SetComprehension: - return true // skip closures - default: - if stop || other.Compare(v) == 0 { - stop = true - return stop - } + if stop || match { + stop = true + return stop } - return false - }) - } + case ast.Value: + if stop || x.Compare(value) == 0 { + stop = true + return stop + } + } + return stop + }) + vis.Walk(x) return stop } diff --git a/topdown/topdown_partial_test.go b/topdown/topdown_partial_test.go index 263800fadd..99c6d62a4f 100644 --- a/topdown/topdown_partial_test.go +++ b/topdown/topdown_partial_test.go @@ -3310,6 +3310,16 @@ func TestTopDownPartialEval(t *testing.T) { }`}, wantQueries: []string{`{true | input.foo} = x_term_1_11; x_term_1_11; x1 = input.foo`}, }, + { + note: "copypropagation: keep equations that are only found in 'every' body", + query: "data.test.p", + modules: []string{`package test + p { + x = input.foo + every y in input.ys { y = input.foo } + }`}, + wantQueries: []string{`every __local0__1, __local1__1 in input.ys { __local1__1 = input.foo }; x1 = input.foo`}, + }, } ctx := context.Background()