topdown/copyprop: skip 'every' body when re-adding removed eqs (open-…

…policy-agent#5375)

Follow-up to 232b0c7.

For the following policy, the second expression does make a difference:

    every x in input.xs { x = input.foo }
    _ = input.foo

Although it's part of the `every` body, when `input.xs` is `[]`, the body is
not evaluated and the overall `every` expression is true. So, we cannot drop
that second equation because its ref already is in the body of `every`.

Signed-off-by: Stephan Renatus <[email protected]>

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
 }
 

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()