From 2f87ed9436e15261ec43e5d82c582ce17f34cce1 Mon Sep 17 00:00:00 2001
From: Sokwhan Huh <sokwhan@google.com>
Date: Tue, 16 Sep 2025 11:22:30 -0700
Subject: [PATCH] Override environment's expected result type during constant
 folding

PiperOrigin-RevId: 807773409
---
 .../java/dev/cel/optimizer/optimizers/BUILD.bazel  |  1 +
 .../optimizers/ConstantFoldingOptimizer.java       |  6 +++++-
 .../optimizers/ConstantFoldingOptimizerTest.java   | 14 ++++++++++++++
 3 files changed, 20 insertions(+), 1 deletion(-)

diff --git a/optimizer/src/main/java/dev/cel/optimizer/optimizers/BUILD.bazel b/optimizer/src/main/java/dev/cel/optimizer/optimizers/BUILD.bazel
index 7728c0ac8..7674870a8 100644
--- a/optimizer/src/main/java/dev/cel/optimizer/optimizers/BUILD.bazel
+++ b/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",
diff --git a/optimizer/src/main/java/dev/cel/optimizer/optimizers/ConstantFoldingOptimizer.java b/optimizer/src/main/java/dev/cel/optimizer/optimizers/ConstantFoldingOptimizer.java
index e46ee2e98..fd52f4138 100644
--- a/optimizer/src/main/java/dev/cel/optimizer/optimizers/ConstantFoldingOptimizer.java
+++ b/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()) {
diff --git a/optimizer/src/test/java/dev/cel/optimizer/optimizers/ConstantFoldingOptimizerTest.java b/optimizer/src/test/java/dev/cel/optimizer/optimizers/ConstantFoldingOptimizerTest.java
index 66c5d86d5..ea1e77edb 100644
--- a/optimizer/src/test/java/dev/cel/optimizer/optimizers/ConstantFoldingOptimizerTest.java
+++ b/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 {