Fix partial evaluation to properly check for comprehension bound variables for planner

PiperOrigin-RevId: 892757311

Author: l46kok

runtime/src/main/java/dev/cel/runtime/planner/ActivationWrapper.java

@@ -19,4 +19,7 @@
 /** Identifies a resolver that can be unwrapped to bypass local variable state. */
 public interface ActivationWrapper extends GlobalResolver {
   GlobalResolver unwrap();
+
+  /** Returns true if the given name is bound by this local activation wrapper. */
+  boolean isLocallyBound(String name);
 }

runtime/src/main/java/dev/cel/runtime/planner/EvalFold.java

@@ -175,6 +175,11 @@ public GlobalResolver unwrap() {
       return resolver;
     }
 
+    @Override
+    public boolean isLocallyBound(String name) {
+      return name.equals(accuVar) || name.equals(iterVar) || name.equals(iterVar2);
+    }
+
     @Override
     public @Nullable Object resolve(String name) {
       if (name.equals(accuVar)) {

runtime/src/main/java/dev/cel/runtime/planner/NamespacedAttribute.java

@@ -75,7 +75,7 @@ public Object resolve(long exprId, GlobalResolver ctx, ExecutionFrame frame) {
 
       PartialVars partialVars = frame.partialVars().orElse(null);
 
-      if (partialVars != null) {
+      if (partialVars != null && !isLocallyBound(resolver, name)) {
         ImmutableList<CelAttributePattern> patterns = partialVars.unknowns();
         // Avoid enhanced for loop to prevent UnmodifiableIterator from being allocated
         for (int i = 0; i < qualifiers.size(); i++) {
@@ -151,6 +151,17 @@ private static Long getEnumValue(EnumType enumType, String field) {
                     String.format("Field %s was not found on enum %s", enumType.name(), field)));
   }
 
+  private boolean isLocallyBound(GlobalResolver resolver, String name) {
+    while (resolver instanceof ActivationWrapper) {
+      ActivationWrapper wrapper = (ActivationWrapper) resolver;
+      if (wrapper.isLocallyBound(name)) {
+        return true;
+      }
+      resolver = wrapper.unwrap();
+    }
+    return false;
+  }
+
   private GlobalResolver unwrapToNonLocal(GlobalResolver resolver) {
     while (resolver instanceof ActivationWrapper) {
       resolver = ((ActivationWrapper) resolver).unwrap();

runtime/src/test/java/dev/cel/runtime/PlannerInterpreterTest.java

@@ -283,6 +283,11 @@ public void planner_unknownResultSet_success() {
     declareVariable("unknown_list", ListType.create(SimpleType.INT));
     source = "unknown_list.map(x, x)";
     runTest(variables, CelAttributePattern.fromQualifiedIdentifier("unknown_list"));
+
+    clearAllDeclarations();
+    declareVariable("x", StructTypeReference.create(TestAllTypes.getDescriptor().getFullName()));
+    source = "cel.bind(x, [1, 2, 3], 1 in x)";
+    runTest(variables, CelAttributePattern.fromQualifiedIdentifier("x"));
   }
 
   @Test

runtime/src/test/resources/planner_unknownResultSet_success.baseline

@@ -458,4 +458,16 @@ single_timestamp {
   seconds: 15
 }
 , unknown_attributes=[unknown_list]}
-result:   CelUnknownSet{attributes=[unknown_list], unknownExprIds=[1]}
+result:   CelUnknownSet{attributes=[unknown_list], unknownExprIds=[1]}
+
+Source: cel.bind(x, [1, 2, 3], 1 in x)
+declare x {
+  value cel.expr.conformance.proto3.TestAllTypes
+}
+=====>
+bindings: {x=single_string: "test"
+single_timestamp {
+  seconds: 15
+}
+, unknown_attributes=[x]}
+result:   true

testing/src/main/java/dev/cel/testing/BUILD.bazel

@@ -87,6 +87,7 @@ java_library(
         "//common/types:message_type_provider",
         "//common/types:type_providers",
         "//common/values:cel_byte_string",
+        "//extensions",
         "//extensions:optional_library",
         "//runtime",
         "//runtime:function_binding",

testing/src/main/java/dev/cel/testing/BaseInterpreterTest.java

@@ -74,6 +74,7 @@
 import dev.cel.expr.conformance.proto3.TestAllTypes;
 import dev.cel.expr.conformance.proto3.TestAllTypes.NestedEnum;
 import dev.cel.expr.conformance.proto3.TestAllTypes.NestedMessage;
+import dev.cel.extensions.CelExtensions;
 import dev.cel.extensions.CelOptionalLibrary;
 import dev.cel.runtime.CelAttributePattern;
 import dev.cel.runtime.CelEvaluationException;
@@ -153,7 +154,7 @@ protected void prepareCompiler(CelTypeProvider typeProvider) {
     this.celCompiler =
         celCompiler
             .toCompilerBuilder()
-            .addLibraries(CelOptionalLibrary.INSTANCE)
+            .addLibraries(CelOptionalLibrary.INSTANCE, CelExtensions.bindings())
             .setOptions(celOptions)
             .build();
   }