Fix planner to search local scopes for identifiers before container resolution

PiperOrigin-RevId: 852950877

Author: l46kok

checker/src/main/java/dev/cel/checker/Env.java

@@ -454,35 +454,66 @@ public Env add(String name, Type type) {
    * <p>Returns {@code null} if the function cannot be found.
    */
   public @Nullable CelIdentDecl tryLookupCelIdent(CelContainer container, String name) {
+    // Attempt to find the decl with just the ident name to account for shadowed variables
+    CelIdentDecl decl = tryLookupCelIdentFromLocalScopes(name);
+    if (decl != null) {
+      return decl;
+    }
+
     for (String cand : container.resolveCandidateNames(name)) {
-      // First determine whether we know this name already.
-      CelIdentDecl decl = findIdentDecl(cand);
+      decl = tryLookupCelIdent(cand);
       if (decl != null) {
         return decl;
       }
+    }
 
-      // Next try to import the name as a reference to a message type.
-      // This is done via the type provider.
-      Optional<CelType> type = typeProvider.lookupCelType(cand);
-      if (type.isPresent()) {
-        decl = CelIdentDecl.newIdentDeclaration(cand, type.get());
-        decls.get(0).putIdent(decl);
-        return decl;
-      }
+    return null;
+  }
 
-      // Next try to import this as an enum value by splitting the name in a type prefix and
-      // the enum inside.
-      Integer enumValue = typeProvider.lookupEnumValue(cand);
-      if (enumValue != null) {
-        decl =
-            CelIdentDecl.newBuilder()
-                .setName(cand)
-                .setType(SimpleType.INT)
-                .setConstant(CelConstant.ofValue(enumValue))
-                .build();
+  private @Nullable CelIdentDecl tryLookupCelIdent(String cand) {
+    // First determine whether we know this name already.
+    CelIdentDecl decl = findIdentDecl(cand);
+    if (decl != null) {
+      return decl;
+    }
 
-        decls.get(0).putIdent(decl);
-        return decl;
+    // Next try to import the name as a reference to a message type.
+    // This is done via the type provider.
+    Optional<CelType> type = typeProvider.lookupCelType(cand);
+    if (type.isPresent()) {
+      decl = CelIdentDecl.newIdentDeclaration(cand, type.get());
+      decls.get(0).putIdent(decl);
+      return decl;
+    }
+
+    // Next try to import this as an enum value by splitting the name in a type prefix and
+    // the enum inside.
+    Integer enumValue = typeProvider.lookupEnumValue(cand);
+    if (enumValue != null) {
+      decl =
+          CelIdentDecl.newBuilder()
+              .setName(cand)
+              .setType(SimpleType.INT)
+              .setConstant(CelConstant.ofValue(enumValue))
+              .build();
+
+      decls.get(0).putIdent(decl);
+      return decl;
+    }
+    return null;
+  }
+
+  private @Nullable CelIdentDecl tryLookupCelIdentFromLocalScopes(String name) {
+    int firstUserSpaceScope = 2;
+    // Iterate from the top of the stack down to the first local scope.
+    // Note that:
+    // Scope 0: Standard environment
+    // Scope 1: User defined environment
+    // Scope 2 and onwards: comprehension scopes
+    for (int i = decls.size() - 1; i >= firstUserSpaceScope; i--) {
+      CelIdentDecl ident = decls.get(i).getIdent(name);
+      if (ident != null) {
+        return ident;
       }
     }
     return null;

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

@@ -17,6 +17,7 @@
 import com.google.auto.value.AutoValue;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.ImmutableSet;
 import com.google.errorprone.annotations.CheckReturnValue;
 import javax.annotation.concurrent.ThreadSafe;
 import dev.cel.common.CelAbstractSyntaxTree;
@@ -64,7 +65,6 @@ public final class ProgramPlanner {
   private final CelContainer container;
   private final CelOptions options;
 
-
   /**
    * Plans a {@link Program} from the provided parsed-only or type-checked {@link
    * CelAbstractSyntaxTree}.
@@ -155,8 +155,15 @@ private PlannedInterpretable planIdent(CelExpr celExpr, PlannerContext ctx) {
       return planCheckedIdent(celExpr.id(), ref, ctx.typeMap());
     }
 
-    return EvalAttribute.create(
-        celExpr.id(), attributeFactory.newMaybeAttribute(celExpr.ident().name()));
+    String name = celExpr.ident().name();
+    // If the identifier matches a variable that's available in local scope (from a comprehension)
+    // force it to resolve as an absolute attribute (simple name only), bypassing container
+    // resolution.
+    if (ctx.scopedVariables().contains(name)) {
+      return EvalAttribute.create(celExpr.id(), attributeFactory.newAbsoluteAttribute(name));
+    }
+
+    return EvalAttribute.create(celExpr.id(), attributeFactory.newMaybeAttribute(name));
   }
 
   private PlannedInterpretable planCheckedIdent(
@@ -291,9 +298,19 @@ private PlannedInterpretable planComprehension(CelExpr expr, PlannerContext ctx)
 
     PlannedInterpretable accuInit = plan(comprehension.accuInit(), ctx);
     PlannedInterpretable iterRange = plan(comprehension.iterRange(), ctx);
-    PlannedInterpretable loopCondition = plan(comprehension.loopCondition(), ctx);
-    PlannedInterpretable loopStep = plan(comprehension.loopStep(), ctx);
-    PlannedInterpretable result = plan(comprehension.result(), ctx);
+
+    PlannerContext innerCtx;
+    if (comprehension.iterVar2().isEmpty()) {
+      innerCtx = ctx.addScopedIdentifiers(comprehension.iterVar(), comprehension.accuVar());
+    } else {
+      innerCtx =
+          ctx.addScopedIdentifiers(
+              comprehension.iterVar(), comprehension.iterVar2(), comprehension.accuVar());
+    }
+
+    PlannedInterpretable loopCondition = plan(comprehension.loopCondition(), innerCtx);
+    PlannedInterpretable loopStep = plan(comprehension.loopStep(), innerCtx);
+    PlannedInterpretable result = plan(comprehension.result(), innerCtx);
 
     return EvalFold.create(
         expr.id(),
@@ -444,8 +461,23 @@ abstract static class PlannerContext {
 
     abstract ImmutableMap<Long, CelType> typeMap();
 
+    /** Tracks variables that are defined in the current scope (e.g. comprehensions). */
+    abstract ImmutableSet<String> scopedVariables();
+
     private static PlannerContext create(CelAbstractSyntaxTree ast) {
-      return new AutoValue_ProgramPlanner_PlannerContext(ast.getReferenceMap(), ast.getTypeMap());
+      return new AutoValue_ProgramPlanner_PlannerContext(
+          ast.getReferenceMap(), ast.getTypeMap(), ImmutableSet.of());
+    }
+
+    /** Returns a new context with the provided variables added to the local scope. */
+    PlannerContext addScopedIdentifiers(String... variables) {
+      ImmutableSet.Builder<String> newLocals =
+          ImmutableSet.builderWithExpectedSize(scopedVariables().size() + variables.length);
+      newLocals.addAll(scopedVariables());
+      newLocals.add(variables);
+
+      return new AutoValue_ProgramPlanner_PlannerContext(
+          referenceMap(), typeMap(), newLocals.build());
     }
   }
 

runtime/src/test/java/dev/cel/runtime/planner/ProgramPlannerTest.java

@@ -851,13 +851,69 @@ public void plan_comprehension_iterationLimit_success() throws Exception {
                 ImmutableList.of(2L, 3L), ImmutableList.of(3L, 4L), ImmutableList.of(4L, 5L)));
   }
 
+  @Test
+  public void plan_comprehension_withContainer_comprehensionVariableShadowed() throws Exception {
+    ProgramPlanner planner =
+        ProgramPlanner.newPlanner(
+            TYPE_PROVIDER,
+            ProtoMessageValueProvider.newInstance(CEL_OPTIONS, DYNAMIC_PROTO),
+            newDispatcher(),
+            CEL_VALUE_CONVERTER,
+            CelContainer.ofName("com"),
+            CEL_OPTIONS);
+    CelAbstractSyntaxTree ast = compile("[0].exists(int_var, int_var == 0)");
+
+    Program program = planner.plan(ast);
+
+    boolean result = (boolean) program.eval(ImmutableMap.of("com.int_var", 1));
+
+    assertThat(result).isTrue();
+  }
+
+  @Test
+  public void plan_ident_withContainer_resolutionOrderHolds() throws Exception {
+    CelContainer container = CelContainer.ofName("foo");
+    CelCompiler celCompiler =
+        CelCompilerFactory.standardCelCompilerBuilder()
+            .setContainer(container)
+            .addVar("bar", SimpleType.INT)
+            .addVar("foo.bar", SimpleType.INT)
+            .build();
+    ProgramPlanner planner =
+        ProgramPlanner.newPlanner(
+            TYPE_PROVIDER,
+            ProtoMessageValueProvider.newInstance(CEL_OPTIONS, DYNAMIC_PROTO),
+            newDispatcher(),
+            CEL_VALUE_CONVERTER,
+            container,
+            CEL_OPTIONS);
+    CelAbstractSyntaxTree ast = compile(celCompiler, "bar");
+
+    Program program = planner.plan(ast);
+
+    Long result =
+        (Long)
+            program.eval(
+                ImmutableMap.of(
+                    "bar", 1,
+                    "foo.bar", 2));
+
+    // When there's no shadowed variables, container resolution should always take precedence (e.g:
+    // `foo.bar`)
+    assertThat(result).isEqualTo(2);
+  }
+
   private CelAbstractSyntaxTree compile(String expression) throws Exception {
-    CelAbstractSyntaxTree ast = CEL_COMPILER.parse(expression).getAst();
+    return compile(CEL_COMPILER, expression);
+  }
+
+  private CelAbstractSyntaxTree compile(CelCompiler compiler, String expression) throws Exception {
+    CelAbstractSyntaxTree ast = compiler.parse(expression).getAst();
     if (isParseOnly) {
       return ast;
     }
 
-    return CEL_COMPILER.check(ast).getAst();
+    return compiler.check(ast).getAst();
   }
 
   private static CelByteString concatenateByteArrays(CelByteString bytes1, CelByteString bytes2) {

runtime/src/test/resources/comprehension.baseline

@@ -12,3 +12,11 @@ Source: [0, 1, 2].exists(x, x > 2)
 =====>
 bindings: {}
 result:   false
+
+Source: [0].exists(x, x == 0)
+declare com.x {
+  value int
+}
+=====>
+bindings: {com.x=1}
+result:   true

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

@@ -858,6 +858,11 @@ public void comprehension() throws Exception {
 
     source = "[0, 1, 2].exists(x, x > 2)";
     runTest();
+
+    declareVariable("com.x", SimpleType.INT);
+    container = CelContainer.ofName("com");
+    source = "[0].exists(x, x == 0)";
+    runTest(ImmutableMap.of("com.x", 1));
   }
 
   @Test