Namespace resolution fix for planner

PiperOrigin-RevId: 857005456

Author: l46kok

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

@@ -0,0 +1,22 @@
+// Copyright 2026 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package dev.cel.runtime.planner;
+
+import dev.cel.runtime.GlobalResolver;
+
+/** Identifies a resolver that can be unwrapped to bypass local variable state. */
+public interface ActivationWrapper extends GlobalResolver {
+  GlobalResolver unwrap();
+}

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

@@ -29,19 +29,22 @@ final class AttributeFactory {
   private final CelValueConverter celValueConverter;
 
   NamespacedAttribute newAbsoluteAttribute(String... names) {
-    return new NamespacedAttribute(typeProvider, celValueConverter, ImmutableSet.copyOf(names));
+    return NamespacedAttribute.create(typeProvider, celValueConverter, ImmutableSet.copyOf(names));
   }
 
   RelativeAttribute newRelativeAttribute(PlannedInterpretable operand) {
     return new RelativeAttribute(operand, celValueConverter);
   }
 
   MaybeAttribute newMaybeAttribute(String name) {
+    // When there's a single name with a dot prefix, it indicates that the 'maybe' attribute is a
+    // globally namespaced identifier.
+    // Otherwise, the candidate names resolved from the container should be inferred.
+    ImmutableSet<String> names =
+        name.startsWith(".") ? ImmutableSet.of(name) : container.resolveCandidateNames(name);
+
     return new MaybeAttribute(
-        this,
-        ImmutableList.of(
-            new NamespacedAttribute(
-                typeProvider, celValueConverter, container.resolveCandidateNames(name))));
+        this, ImmutableList.of(NamespacedAttribute.create(typeProvider, celValueConverter, names)));
   }
 
   static AttributeFactory newAttributeFactory(

runtime/src/main/java/dev/cel/runtime/planner/BUILD.bazel

@@ -114,6 +114,7 @@ java_library(
         "RelativeAttribute.java",
     ],
     deps = [
+        ":activation_wrapper",
         ":eval_helpers",
         ":execution_frame",
         ":planned_interpretable",
@@ -130,6 +131,14 @@ java_library(
     ],
 )
 
+java_library(
+    name = "activation_wrapper",
+    srcs = ["ActivationWrapper.java"],
+    deps = [
+        "//runtime:interpretable",
+    ],
+)
+
 java_library(
     name = "qualifier",
     srcs = ["Qualifier.java"],
@@ -333,6 +342,7 @@ java_library(
     name = "eval_fold",
     srcs = ["EvalFold.java"],
     deps = [
+        ":activation_wrapper",
         ":execution_frame",
         ":planned_interpretable",
         "//runtime:concatenated_list_view",

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

@@ -97,7 +97,7 @@ private Object evalMap(Map<?, ?> iterRange, Folder folder, ExecutionFrame frame)
       if (!iterVar2.isEmpty()) {
         folder.iterVar2Val = entry.getValue();
       }
-      
+
       boolean cond = (boolean) condition.eval(folder, frame);
       if (!cond) {
         return result.eval(folder, frame);
@@ -149,7 +149,7 @@ private static Object maybeUnwrapAccumulator(Object val) {
     return val;
   }
 
-  private static class Folder implements GlobalResolver {
+  private static class Folder implements ActivationWrapper {
     private final GlobalResolver resolver;
     private final String accuVar;
     private final String iterVar;
@@ -166,6 +166,11 @@ private Folder(GlobalResolver resolver, String accuVar, String iterVar, String i
       this.iterVar2 = iterVar2;
     }
 
+    @Override
+    public GlobalResolver unwrap() {
+      return resolver;
+    }
+
     @Override
     public @Nullable Object resolve(String name) {
       if (name.equals(accuVar)) {

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

@@ -28,15 +28,30 @@
 
 @Immutable
 final class NamespacedAttribute implements Attribute {
+  private final ImmutableSet<Integer> disambiguateNames;
   private final ImmutableSet<String> namespacedNames;
   private final ImmutableList<Qualifier> qualifiers;
   private final CelValueConverter celValueConverter;
   private final CelTypeProvider typeProvider;
 
   @Override
   public Object resolve(GlobalResolver ctx, ExecutionFrame frame) {
+    GlobalResolver inputVars = ctx;
+    // Unwrap any local activations to ensure that we reach the variables provided as input
+    // to the expression in the event that we need to disambiguate between global and local
+    // variables.
+    if (!disambiguateNames.isEmpty()) {
+      inputVars = unwrapToRoot(ctx);
+    }
+
+    int i = 0;
     for (String name : namespacedNames) {
-      Object value = ctx.resolve(name);
+      GlobalResolver resolver = ctx;
+      if (disambiguateNames.contains(i)) {
+        resolver = inputVars;
+      }
+
+      Object value = resolver.resolve(name);
       if (value != null) {
         if (!qualifiers.isEmpty()) {
           return applyQualifiers(value, celValueConverter, qualifiers);
@@ -69,6 +84,7 @@ public Object resolve(GlobalResolver ctx, ExecutionFrame frame) {
         throw new IllegalStateException(
             "Unexpected type resolution when there were remaining qualifiers: " + type.name());
       }
+      i++;
     }
 
     return MissingAttribute.newMissingAttribute(namespacedNames);
@@ -82,12 +98,20 @@ ImmutableSet<String> candidateVariableNames() {
     return namespacedNames;
   }
 
+  private GlobalResolver unwrapToRoot(GlobalResolver resolver) {
+    while (resolver instanceof ActivationWrapper) {
+      resolver = ((ActivationWrapper) resolver).unwrap();
+    }
+    return resolver;
+  }
+
   @Override
   public NamespacedAttribute addQualifier(Qualifier qualifier) {
     return new NamespacedAttribute(
         typeProvider,
         celValueConverter,
         namespacedNames,
+        disambiguateNames,
         ImmutableList.<Qualifier>builder().addAll(qualifiers).add(qualifier).build());
   }
 
@@ -106,21 +130,40 @@ private static Object applyQualifiers(
     return obj;
   }
 
-  NamespacedAttribute(
+  static NamespacedAttribute create(
       CelTypeProvider typeProvider,
       CelValueConverter celValueConverter,
       ImmutableSet<String> namespacedNames) {
-    this(typeProvider, celValueConverter, namespacedNames, ImmutableList.of());
+    ImmutableSet.Builder<String> namesBuilder = ImmutableSet.builder();
+    ImmutableSet.Builder<Integer> indicesBuilder = ImmutableSet.builder();
+    int i = 0;
+    for (String name : namespacedNames) {
+      if (name.startsWith(".")) {
+        indicesBuilder.add(i);
+        namesBuilder.add(name.substring(1));
+      } else {
+        namesBuilder.add(name);
+      }
+      i++;
+    }
+    return new NamespacedAttribute(
+        typeProvider,
+        celValueConverter,
+        namesBuilder.build(),
+        indicesBuilder.build(),
+        ImmutableList.of());
   }
 
-  private NamespacedAttribute(
+  NamespacedAttribute(
       CelTypeProvider typeProvider,
       CelValueConverter celValueConverter,
       ImmutableSet<String> namespacedNames,
+      ImmutableSet<Integer> disambiguateNames,
       ImmutableList<Qualifier> qualifiers) {
     this.typeProvider = typeProvider;
     this.celValueConverter = celValueConverter;
     this.namespacedNames = namespacedNames;
+    this.disambiguateNames = disambiguateNames;
     this.qualifiers = qualifiers;
   }
 }

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

@@ -15,6 +15,7 @@
 package dev.cel.runtime.planner;
 
 import com.google.auto.value.AutoValue;
+import com.google.common.base.Strings;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
@@ -48,6 +49,7 @@
 import dev.cel.runtime.CelResolvedOverload;
 import dev.cel.runtime.DefaultDispatcher;
 import dev.cel.runtime.Program;
+import java.util.HashMap;
 import java.util.NoSuchElementException;
 import java.util.Optional;
 
@@ -161,8 +163,12 @@ 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 identName = celExpr.ident().name();
+    if (ctx.isLocalVar(identName)) {
+      return EvalAttribute.create(celExpr.id(), attributeFactory.newAbsoluteAttribute(identName));
+    }
+
+    return EvalAttribute.create(celExpr.id(), attributeFactory.newMaybeAttribute(identName));
   }
 
   private PlannedInterpretable planCheckedIdent(
@@ -314,10 +320,18 @@ private PlannedInterpretable planComprehension(CelExpr expr, PlannerContext ctx)
 
     PlannedInterpretable accuInit = plan(comprehension.accuInit(), ctx);
     PlannedInterpretable iterRange = plan(comprehension.iterRange(), ctx);
+
+    ctx.pushLocalVars(comprehension.accuVar(), comprehension.iterVar(), comprehension.iterVar2());
+
     PlannedInterpretable loopCondition = plan(comprehension.loopCondition(), ctx);
     PlannedInterpretable loopStep = plan(comprehension.loopStep(), ctx);
+
+    ctx.popLocalVars(comprehension.iterVar(), comprehension.iterVar2());
+
     PlannedInterpretable result = plan(comprehension.result(), ctx);
 
+    ctx.popLocalVars(comprehension.accuVar());
+
     return EvalFold.create(
         expr.id(),
         comprehension.accuVar(),
@@ -460,15 +474,57 @@ private static Builder newBuilder() {
     }
   }
 
-  @AutoValue
-  abstract static class PlannerContext {
+  static final class PlannerContext {
+    private final ImmutableMap<Long, CelReference> referenceMap;
+    private final ImmutableMap<Long, CelType> typeMap;
+    private final HashMap<String, Integer> localVars = new HashMap<>();
+
+    ImmutableMap<Long, CelReference> referenceMap() {
+      return referenceMap;
+    }
 
-    abstract ImmutableMap<Long, CelReference> referenceMap();
+    ImmutableMap<Long, CelType> typeMap() {
+      return typeMap;
+    }
 
-    abstract ImmutableMap<Long, CelType> typeMap();
+    private void pushLocalVars(String... names) {
+      for (String name : names) {
+        if (Strings.isNullOrEmpty(name)) {
+          continue;
+        }
+        localVars.merge(name, 1, Integer::sum);
+      }
+    }
+
+    private void popLocalVars(String... names) {
+      for (String name : names) {
+        if (Strings.isNullOrEmpty(name)) {
+          continue;
+        }
+        Integer count = localVars.get(name);
+        if (count != null) {
+          if (count == 1) {
+            localVars.remove(name);
+          } else {
+            localVars.put(name, count - 1);
+          }
+        }
+      }
+    }
+
+    /** Checks if the given name is a local variable in the current scope. */
+    private boolean isLocalVar(String name) {
+      return localVars.containsKey(name);
+    }
+
+    private PlannerContext(
+        ImmutableMap<Long, CelReference> referenceMap, ImmutableMap<Long, CelType> typeMap) {
+      this.referenceMap = referenceMap;
+      this.typeMap = typeMap;
+    }
 
-    private static PlannerContext create(CelAbstractSyntaxTree ast) {
-      return new AutoValue_ProgramPlanner_PlannerContext(ast.getReferenceMap(), ast.getTypeMap());
+    static PlannerContext create(CelAbstractSyntaxTree ast) {
+      return new PlannerContext(ast.getReferenceMap(), ast.getTypeMap());
     }
   }