Override environment's expected result type during constant folding

PiperOrigin-RevId: 807744200

Author: l46kok

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

@@ -26,6 +26,7 @@ java_library(
         "//common/ast",
         "//common/ast:mutable_expr",
         "//common/navigation:mutable_navigation",
+        "//common/types",
         "//extensions:optional_library",
         "//optimizer:ast_optimizer",
         "//optimizer:mutable_ast",

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

@@ -38,6 +38,7 @@
 import dev.cel.common.ast.CelMutableExprConverter;
 import dev.cel.common.navigation.CelNavigableMutableAst;
 import dev.cel.common.navigation.CelNavigableMutableExpr;
+import dev.cel.common.types.SimpleType;
 import dev.cel.extensions.CelOptionalLibrary.Function;
 import dev.cel.optimizer.AstMutator;
 import dev.cel.optimizer.CelAstOptimizer;
@@ -88,6 +89,9 @@ private static CelMutableExpr newOptionalNoneExpr() {
   @Override
   public OptimizationResult optimize(CelAbstractSyntaxTree ast, Cel cel)
       throws CelOptimizationException {
+    // Override the environment's expected type to generally allow all subtrees to be folded.
+    Cel optimizerEnv = cel.toCelBuilder().setResultType(SimpleType.DYN).build();
+
     CelMutableAst mutableAst = CelMutableAst.fromCelAst(ast);
     int iterCount = 0;
     boolean continueFolding = true;
@@ -112,7 +116,7 @@ public OptimizationResult optimize(CelAbstractSyntaxTree ast, Cel cel)
         mutatedResult = maybePruneBranches(mutableAst, foldableExpr.expr());
         if (!mutatedResult.isPresent()) {
           // Evaluate the call then fold
-          mutatedResult = maybeFold(cel, mutableAst, foldableExpr);
+          mutatedResult = maybeFold(optimizerEnv, mutableAst, foldableExpr);
         }
 
         if (!mutatedResult.isPresent()) {

optimizer/src/test/java/dev/cel/optimizer/optimizers/ConstantFoldingOptimizerTest.java

@@ -347,6 +347,20 @@ public void constantFold_addFoldableFunction_success() throws Exception {
     assertThat(CEL_UNPARSER.unparse(optimizedAst)).isEqualTo("true");
   }
 
+  @Test
+  public void constantFold_withExpectedResultTypeSet_success() throws Exception {
+    Cel cel = CelFactory.standardCelBuilder().setResultType(SimpleType.STRING).build();
+    CelOptimizer optimizer =
+        CelOptimizerFactory.standardCelOptimizerBuilder(cel)
+            .addAstOptimizers(ConstantFoldingOptimizer.getInstance())
+            .build();
+    CelAbstractSyntaxTree ast = cel.compile("string(!true)").getAst();
+
+    CelAbstractSyntaxTree optimizedAst = optimizer.optimize(ast);
+
+    assertThat(CEL_UNPARSER.unparse(optimizedAst)).isEqualTo("\"false\"");
+  }
+
   @Test
   public void constantFold_withMacroCallPopulated_comprehensionsAreReplacedWithNotSet()
       throws Exception {