From 5c7efafd3c5abc611fb07f941dbcaec109865340 Mon Sep 17 00:00:00 2001
From: Sokwhan Huh <sokwhan@google.com>
Date: Tue, 5 Mar 2024 13:44:00 -0800
Subject: [PATCH] Fix constant folding when ternary arg is replaced with a
 comprehension

PiperOrigin-RevId: 612955389
---
 optimizer/src/main/java/dev/cel/optimizer/MutableAst.java   | 6 +++++-
 .../optimizer/optimizers/ConstantFoldingOptimizerTest.java  | 5 ++++-
 2 files changed, 9 insertions(+), 2 deletions(-)

diff --git a/optimizer/src/main/java/dev/cel/optimizer/MutableAst.java b/optimizer/src/main/java/dev/cel/optimizer/MutableAst.java
index 181dda25c..bf813eef4 100644
--- a/optimizer/src/main/java/dev/cel/optimizer/MutableAst.java
+++ b/optimizer/src/main/java/dev/cel/optimizer/MutableAst.java
@@ -130,7 +130,11 @@ public CelAbstractSyntaxTree replaceSubtree(
       CelAbstractSyntaxTree ast, CelExpr newExpr, long exprIdToReplace) {
     return replaceSubtreeWithNewAst(
         ast,
-        CelAbstractSyntaxTree.newParsedAst(newExpr, CelSource.newBuilder().build()),
+        CelAbstractSyntaxTree.newParsedAst(
+            newExpr,
+            // Copy the macro call information to the new AST such that macro call map can be
+            // normalized post-replacement.
+            CelSource.newBuilder().addAllMacroCalls(ast.getSource().getMacroCalls()).build()),
         exprIdToReplace);
   }
 
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 5ccc41388..ddadd5a20 100644
--- a/optimizer/src/test/java/dev/cel/optimizer/optimizers/ConstantFoldingOptimizerTest.java
+++ b/optimizer/src/test/java/dev/cel/optimizer/optimizers/ConstantFoldingOptimizerTest.java
@@ -204,6 +204,7 @@ public void constantFold_success(String source, String expected) throws Exceptio
   @TestParameters(
       "{source: 'cel.bind(r0, [1, 2, 3], cel.bind(r1, 1 in r0 && 2 in x, r1))', expected:"
           + " 'cel.bind(r0, [1, 2, 3], cel.bind(r1, 1 in r0 && 2 in x, r1))'}")
+  @TestParameters("{source: 'false ? false : cel.bind(a, x, a)', expected: 'cel.bind(a, x, a)'}")
   public void constantFold_macros_macroCallMetadataPopulated(String source, String expected)
       throws Exception {
     Cel cel =
@@ -247,6 +248,8 @@ public void constantFold_macros_macroCallMetadataPopulated(String source, String
   @TestParameters(
       "{source: '[{}, {\"a\": 1}, {\"b\": 2}].filter(m, has({\"a\": true}.a)) == "
           + " [{}, {\"a\": 1}, {\"b\": 2}]'}")
+  @TestParameters("{source: 'cel.bind(r0, [1, 2, 3], cel.bind(r1, 1 in r0 && 2 in r0, r1))'}")
+  @TestParameters("{source: 'false ? false : cel.bind(a, true, a)'}")
   public void constantFold_macros_withoutMacroCallMetadata(String source) throws Exception {
     Cel cel =
         CelFactory.standardCelBuilder()
@@ -255,7 +258,7 @@ public void constantFold_macros_withoutMacroCallMetadata(String source) throws E
             .addMessageTypes(TestAllTypes.getDescriptor())
             .setStandardMacros(CelStandardMacro.STANDARD_MACROS)
             .setOptions(CelOptions.current().populateMacroCalls(false).build())
-            .addCompilerLibraries(CelOptionalLibrary.INSTANCE)
+            .addCompilerLibraries(CelExtensions.bindings(), CelOptionalLibrary.INSTANCE)
             .addRuntimeLibraries(CelOptionalLibrary.INSTANCE)
             .build();
     CelOptimizer celOptimizer =