Change AST Mutator to accept mutable ASTs

common/src/main/java/dev/cel/common/CelSource.java

@@ -279,6 +279,10 @@ public Builder clearMacroCall(long exprId) {
       return this;
     }
 
+    public ImmutableSet<Extension> getExtensions() {
+      return extensions.build();
+    }
+
     /**
      * Adds one or more {@link Extension}s to the source information. Extensions implement set
      * semantics and deduped if same ones are provided.

common/src/main/java/dev/cel/common/navigation/BUILD.bazel

@@ -57,6 +57,7 @@ java_library(
         ":common",
         "//:auto_value",
         "//common/ast:mutable_ast",
+        "//common/types:type_providers",
         "@maven//:com_google_errorprone_error_prone_annotations",
         "@maven//:com_google_guava_guava",
         "@maven//:org_jspecify_jspecify",

common/src/main/java/dev/cel/common/navigation/CelNavigableMutableAst.java

@@ -15,6 +15,8 @@
 package dev.cel.common.navigation;
 
 import dev.cel.common.ast.CelMutableAst;
+import dev.cel.common.types.CelType;
+import java.util.Optional;
 
 /**
  * Decorates a {@link CelMutableAst} with navigational properties. This allows us to visit a node's
@@ -44,4 +46,15 @@ public CelNavigableMutableExpr getRoot() {
   public CelMutableAst getAst() {
     return ast;
   }
+
+  /**
+   * Returns the type of the expression node for a type-checked AST. This simply proxies down the
+   * call to {@link CelMutableAst#getType(long)}.
+   *
+   * @return Optional of {@link CelType} or {@link Optional#empty} if the type does not exist at the
+   *     ID.
+   */
+  public Optional<CelType> getType(long exprId) {
+    return ast.getType(exprId);
+  }
 }

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

@@ -93,6 +93,7 @@ java_library(
         "//common/ast:mutable_ast",
         "//common/navigation",
         "//common/navigation:common",
+        "//common/navigation:mutable_navigation",
         "//common/types:type_providers",
         "@maven//:com_google_errorprone_error_prone_annotations",
         "@maven//:com_google_guava_guava",

optimizer/src/main/java/dev/cel/optimizer/optimizers/SubexpressionOptimizer.java

@@ -222,12 +222,7 @@ private OptimizationResult optimizeUsingCelBlock(CelNavigableAst navigableAst, C
                       iterVarType ->
                           newVarDecls.add(
                               CelVarDecl.newVarDeclaration(name.iterVarName(), iterVarType)));
-              type.resultType()
-                  .ifPresent(
-                      comprehensionResultType ->
-                          newVarDecls.add(
-                              CelVarDecl.newVarDeclaration(
-                                  name.resultName(), comprehensionResultType)));
+              newVarDecls.add(CelVarDecl.newVarDeclaration(name.resultName(), type.resultType()));
             });
 
     // Type-check all sub-expressions then create new block index identifiers.

optimizer/src/test/java/dev/cel/optimizer/BUILD.bazel

@@ -13,7 +13,9 @@ java_library(
         "//common:compiler_common",
         "//common:options",
         "//common/ast",
+        "//common/ast:mutable_ast",
         "//common/navigation",
+        "//common/navigation:mutable_navigation",
         "//common/resources/testdata/proto3:test_all_types_java_proto",
         "//common/types",
         "//compiler",