copybara-service · Oct 10, 2026
diff --git a/common/ast/BUILD.bazel b/common/ast/BUILD.bazel
@@ -48,6 +48,11 @@ java_library(
     exports = ["//common/src/main/java/dev/cel/common/ast:cel_expr_visitor"],
 )
 
+cel_android_library(
+    name = "cel_expr_visitor_android",
+    exports = ["//common/src/main/java/dev/cel/common/ast:cel_expr_visitor_android"],
+)
+
 java_library(
     name = "expr_factory",
     exports = ["//common/src/main/java/dev/cel/common/ast:expr_factory"],

diff --git a/common/src/main/java/dev/cel/common/BUILD.bazel b/common/src/main/java/dev/cel/common/BUILD.bazel
@@ -350,6 +350,7 @@ java_library(
         "//common/ast",
         "//common/ast:mutable_expr",
         "//common/types:type_providers",
+        "@maven//:com_google_guava_guava",
     ],
 )
 

diff --git a/common/src/main/java/dev/cel/common/CelMutableAst.java b/common/src/main/java/dev/cel/common/CelMutableAst.java
@@ -14,11 +14,11 @@
 
 package dev.cel.common;
 
+import com.google.common.collect.ImmutableMap;
 import dev.cel.common.ast.CelMutableExpr;
 import dev.cel.common.ast.CelMutableExprConverter;
 import dev.cel.common.ast.CelReference;
 import dev.cel.common.types.CelType;
-import java.util.HashMap;
 import java.util.Map;
 import java.util.Optional;
 
@@ -32,8 +32,8 @@
 public final class CelMutableAst {
   private final CelMutableExpr mutatedExpr;
   private final CelMutableSource source;
-  private final Map<Long, CelReference> references;
-  private final Map<Long, CelType> types;
+  private final ImmutableMap<Long, CelReference> references;
+  private final ImmutableMap<Long, CelType> types;
 
   /** Returns the underlying {@link CelMutableExpr} representation of the abstract syntax tree. */
   public CelMutableExpr expr() {
@@ -68,6 +68,11 @@ public Optional<CelType> getType(long exprId) {
     return Optional.ofNullable(types.get(exprId));
   }
 
+  /** Returns the type map of the abstract syntax tree. */
+  public ImmutableMap<Long, CelType> getTypeMap() {
+    return types;
+  }
+
   /** Converts this mutable AST into a parsed {@link CelAbstractSyntaxTree}. */
   public CelAbstractSyntaxTree toParsedAst() {
     return toParsedAst(false);
@@ -102,11 +107,16 @@ public static CelMutableAst fromCelAst(CelAbstractSyntaxTree ast) {
    * builder.
    */
   public static CelMutableAst of(CelMutableExpr mutableExpr, CelMutableSource mutableSource) {
-    return new CelMutableAst(mutableExpr, mutableSource);
+    return new CelMutableAst(mutableExpr, mutableSource, ImmutableMap.of(), ImmutableMap.of());
   }
 
-  private CelMutableAst(CelMutableExpr mutatedExpr, CelMutableSource mutableSource) {
-    this(mutatedExpr, mutableSource, new HashMap<>(), new HashMap<>());
+  /**
+   * Constructs an instance of {@link CelMutableAst} with the mutable expression, its source
+   * builder, and node types.
+   */
+  public static CelMutableAst of(
+      CelMutableExpr mutableExpr, CelMutableSource mutableSource, Map<Long, CelType> types) {
+    return new CelMutableAst(mutableExpr, mutableSource, ImmutableMap.of(), types);
   }
 
   private CelMutableAst(
@@ -116,7 +126,7 @@ private CelMutableAst(
       Map<Long, CelType> types) {
     this.mutatedExpr = mutatedExpr;
     this.source = mutableSource;
-    this.references = new HashMap<>(references);
-    this.types = new HashMap<>(types);
+    this.references = ImmutableMap.copyOf(references);
+    this.types = ImmutableMap.copyOf(types);
   }
 }
diff --git a/common/src/main/java/dev/cel/common/ast/BUILD.bazel b/common/src/main/java/dev/cel/common/ast/BUILD.bazel
@@ -7,6 +7,7 @@ package(
     ],
     default_visibility = [
         "//common/ast:__pkg__",
+        "//publish:__pkg__",
     ],
 )
 
@@ -64,9 +65,9 @@ java_library(
     ],
     deps = [
         ":ast",
+        ":cel_expr_visitor",
         "//common:cel_ast",
         "//common/annotations",
-        "//common/navigation",
         "@maven//:com_google_guava_guava",
     ],
 )
@@ -78,9 +79,9 @@ cel_android_library(
     ],
     deps = [
         ":ast_android",
+        ":cel_expr_visitor_android",
         "//common:cel_ast_android",
         "//common/annotations",
-        "//common/navigation:navigation_android",
         "@maven_android//:com_google_guava_guava",
     ],
 )
@@ -143,6 +144,17 @@ java_library(
     ],
 )
 
+cel_android_library(
+    name = "cel_expr_visitor_android",
+    srcs = ["CelExprVisitor.java"],
+    tags = [
+    ],
+    deps = [
+        ":ast_android",
+        "//common:cel_ast_android",
+    ],
+)
+
 java_library(
     name = "expr_factory",
     srcs = EXPR_FACTORY_SOURCES,

diff --git a/common/src/main/java/dev/cel/common/ast/CelBlock.java b/common/src/main/java/dev/cel/common/ast/CelBlock.java
@@ -14,13 +14,14 @@
 
 package dev.cel.common.ast;
 
-import static com.google.common.collect.ImmutableList.toImmutableList;
+import static com.google.common.base.Preconditions.checkArgument;
 
-import com.google.common.base.Preconditions;
 import com.google.common.collect.ImmutableList;
 import dev.cel.common.CelAbstractSyntaxTree;
 import dev.cel.common.annotations.Internal;
-import dev.cel.common.navigation.CelNavigableExpr;
+import dev.cel.common.ast.CelExpr.CelCall;
+import dev.cel.common.ast.CelExpr.CelIdent;
+import dev.cel.common.ast.CelExpr.ExprKind.Kind;
 import java.util.Optional;
 
 /**
@@ -36,10 +37,6 @@ public final class CelBlock {
 
   private final CelExpr blockExpr;
 
-  private CelBlock(CelExpr blockExpr) {
-    this.blockExpr = blockExpr;
-  }
-
   public ImmutableList<CelExpr> indices() {
     return blockExpr.call().args().get(0).list().elements();
   }
@@ -61,84 +58,79 @@ public CelExpr expr() {
    * @throws IllegalArgumentException if the block is malformed or its indices are invalid.
    */
   public static Optional<CelBlock> extract(CelAbstractSyntaxTree ast) {
-    CelNavigableExpr celNavigableExpr = CelNavigableExpr.fromExpr(ast.getExpr());
-
-    ImmutableList<CelExpr> allCelBlocks =
-        celNavigableExpr
-            .allNodes()
-            .map(CelNavigableExpr::expr)
-            .filter(expr -> expr.callOrDefault().function().equals(FUNCTION_NAME))
-            .collect(toImmutableList());
-    if (allCelBlocks.isEmpty()) {
+    CelExpr root = ast.getExpr();
+    BlockValidator validator = new BlockValidator(root);
+    if (!isBlockCall(root)) {
+      validator.validate(root, 0);
       return Optional.empty();
     }
 
-    Preconditions.checkArgument(
-        allCelBlocks.size() == 1,
-        "Expected 1 cel.block function to be present but found %s",
-        allCelBlocks.size());
-    Preconditions.checkArgument(
-        celNavigableExpr.expr().equals(allCelBlocks.get(0)),
-        "Expected cel.block to be present at root");
-
-    return Optional.of(fromExpr(allCelBlocks.get(0)));
-  }
-
-  /**
-   * Constructs a {@link CelBlock} from a {@link CelExpr}.
-   *
-   * @throws IllegalArgumentException if the expression is not a valid block.
-   */
-  private static CelBlock fromExpr(CelExpr expr) {
-    Preconditions.checkArgument(
-        expr.exprKind().getKind() == CelExpr.ExprKind.Kind.CALL,
-        "Expected cel.@block to be a call expression");
-    Preconditions.checkArgument(
-        expr.call().function().equals(FUNCTION_NAME), "Expected function to be cel.@block");
-    Preconditions.checkArgument(
-        expr.call().args().size() == 2, "Expected exactly 2 arguments for cel.@block");
-    Preconditions.checkArgument(
-        expr.call().args().get(0).exprKind().getKind() == CelExpr.ExprKind.Kind.LIST,
-        "Expected first argument of cel.@block to be a list");
-
-    CelBlock block = new CelBlock(expr);
-
-    // Assert correctness on block indices used in subexpressions
+    CelBlock block = new CelBlock(root);
     ImmutableList<CelExpr> subexprs = block.indices();
     for (int i = 0; i < subexprs.size(); i++) {
-      verifyBlockIndex(subexprs.get(i), i, expr);
+      validator.validate(subexprs.get(i), i);
     }
 
-    // Assert correctness on block indices used in block result
-    CelExpr blockResult = block.result();
-    verifyBlockIndex(blockResult, subexprs.size(), expr);
-    boolean resultHasAtLeastOneBlockIndex =
-        CelNavigableExpr.fromExpr(blockResult)
-            .allNodes()
-            .map(CelNavigableExpr::expr)
-            .anyMatch(e -> e.identOrDefault().name().startsWith(INDEX_PREFIX));
-    Preconditions.checkArgument(
-        resultHasAtLeastOneBlockIndex,
+    checkArgument(
+        validator.validate(block.result(), subexprs.size()),
         "Expected at least one reference of index in cel.block result");
 
-    return block;
+    return Optional.of(block);
+  }
+
+  private static boolean isBlockCall(CelExpr expr) {
+    return expr.getKind().equals(Kind.CALL) && expr.call().function().equals(FUNCTION_NAME);
+  }
+
+  private static final class BlockValidator extends CelExprVisitor {
+    private final CelExpr root;
+    private int maxIndexValue;
+    private boolean hasBlockIndex;
+
+    private boolean validate(CelExpr expr, int maxIndexValue) {
+      this.maxIndexValue = maxIndexValue;
+      this.hasBlockIndex = false;
+      visit(expr);
+      return hasBlockIndex;
+    }
+
+    @Override
+    public void visit(CelExpr expr) {
+      if (!expr.getKind().equals(Kind.NOT_SET)) {
+        super.visit(expr);
+      }
+    }
+
+    @Override
+    protected void visit(CelExpr expr, CelCall call) {
+      if (call.function().equals(FUNCTION_NAME)) {
+        throw new IllegalArgumentException(
+            isBlockCall(root)
+                ? "Expected 1 cel.block function to be present but found 2"
+                : "Expected cel.block to be present at root");
+      }
+      super.visit(expr, call);
+    }
+
+    @Override
+    protected void visit(CelExpr expr, CelIdent ident) {
+      if (ident.name().startsWith(INDEX_PREFIX)) {
+        hasBlockIndex = true;
+        int indexValue = Integer.parseInt(ident.name().substring(INDEX_PREFIX.length()));
+        checkArgument(
+            indexValue < maxIndexValue,
+            "Illegal block index found. The index value must be less than %s. Expr: %s",
+            maxIndexValue,
+            root);
+      }
+    }
+
+    private BlockValidator(CelExpr root) {
+      this.root = root;
+    }
   }
 
-  private static void verifyBlockIndex(CelExpr celExpr, int maxIndexValue, CelExpr rootBlock) {
-    boolean areAllIndicesValid =
-        CelNavigableExpr.fromExpr(celExpr)
-            .allNodes()
-            .map(CelNavigableExpr::expr)
-            .filter(expr -> expr.identOrDefault().name().startsWith(INDEX_PREFIX))
-            .map(CelExpr::ident)
-            .allMatch(
-                blockIdent ->
-                    Integer.parseInt(blockIdent.name().substring(INDEX_PREFIX.length()))
-                        < maxIndexValue);
-    Preconditions.checkArgument(
-        areAllIndicesValid,
-        "Illegal block index found. The index value must be less than %s. Expr: %s",
-        maxIndexValue,
-        rootBlock);
+  private CelBlock(CelExpr blockExpr) {
+    this.blockExpr = blockExpr;
   }
 }
diff --git a/optimizer/src/main/java/dev/cel/optimizer/AstMutator.java b/optimizer/src/main/java/dev/cel/optimizer/AstMutator.java
@@ -16,6 +16,7 @@
 
 import static com.google.common.collect.ImmutableList.toImmutableList;
 import static com.google.common.collect.ImmutableMap.toImmutableMap;
+import static com.google.common.collect.Streams.stream;
 import static java.lang.Math.max;
 import static java.util.stream.Collectors.toCollection;
 
@@ -26,7 +27,6 @@
 import com.google.common.collect.HashBasedTable;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
-import com.google.common.collect.Streams;
 import com.google.common.collect.Table;
 import com.google.errorprone.annotations.Immutable;
 import dev.cel.common.CelAbstractSyntaxTree;
@@ -283,7 +283,7 @@ public MangledComprehensionAst mangleComprehensionIdentifierNames(
     }
 
     return MangledComprehensionAst.of(
-        CelMutableAst.of(mutatedComprehensionExpr, newSource),
+        CelMutableAst.of(mutatedComprehensionExpr, newSource, ast.getTypeMap()),
         ImmutableMap.copyOf(mangledIdentNamesToType));
   }
 
@@ -506,6 +506,11 @@ public CelMutableAst replaceSubtree(
 
     CelMutableExpr mutatedRoot =
         mutateExpr(stableIdGenerator::renumberId, ast.expr(), newAst.expr(), exprIdToReplace);
+    ImmutableMap<Long, CelType> remappedTypes =
+        ast.getTypeMap().entrySet().stream()
+            .filter(e -> stableIdGenerator.hasId(e.getKey()))
+            .collect(
+                toImmutableMap(e -> stableIdGenerator.renumberId(e.getKey()), Entry::getValue));
     CelMutableSource newAstSource =
         CelMutableSource.newInstance().setDescription(ast.source().getDescription());
     if (!ast.source().getMacroCalls().isEmpty()) {
@@ -521,7 +526,7 @@ public CelMutableAst replaceSubtree(
     newAstSource =
         normalizeMacroSource(
             newAstSource, exprIdToReplace, mutatedRoot, stableIdGenerator::renumberId);
-    return CelMutableAst.of(mutatedRoot, newAstSource);
+    return CelMutableAst.of(mutatedRoot, newAstSource, remappedTypes);
   }
 
   /**
@@ -625,7 +630,7 @@ public CelMutableAst mutateUntilFixedPoint(
             navAst
                 .getRoot()
                 .allNodes(traversalOrder)
-                .flatMap(node -> Streams.stream(nodeRewriter.apply(node)))
+                .flatMap(node -> stream(nodeRewriter.apply(node)))
                 .findFirst());
   }
 

diff --git a/optimizer/src/main/java/dev/cel/optimizer/optimizers/BUILD.bazel b/optimizer/src/main/java/dev/cel/optimizer/optimizers/BUILD.bazel
@@ -71,8 +71,6 @@ java_library(
         "//common/navigation",
         "//common/navigation:common",
         "//common/navigation:mutable_navigation",
-        "//common/types",
-        "//common/types:type_providers",
         "//extensions:bindings",
         "//optimizer:ast_optimizer",
         "//optimizer:mutable_ast",