From 9828b360bf8286bd1adc0be838d31653ab558067 Mon Sep 17 00:00:00 2001
From: Sean Huh <sokwhan@google.com>
Date: Thu, 29 Oct 2026 14:55:55 -0700
Subject: [PATCH] Keep SelectOptimizer from reusing macro call node IDs

PiperOrigin-RevId: 1009987752
---
 .../optimizer/optimizers/SelectOptimizer.java | 10 +++++++--
 .../optimizers/SelectOptimizerTest.java       | 21 +++++++++++++++++++
 2 files changed, 29 insertions(+), 2 deletions(-)

diff --git a/optimizer/src/main/java/dev/cel/optimizer/optimizers/SelectOptimizer.java b/optimizer/src/main/java/dev/cel/optimizer/optimizers/SelectOptimizer.java
index 5cd15f818..79dfb8fed 100644
--- a/optimizer/src/main/java/dev/cel/optimizer/optimizers/SelectOptimizer.java
+++ b/optimizer/src/main/java/dev/cel/optimizer/optimizers/SelectOptimizer.java
@@ -19,6 +19,7 @@
 import static com.google.common.collect.ImmutableList.toImmutableList;
 import static dev.cel.optimizer.optimizers.DefaultOptimizerConstants.CEL_ATTRIBUTE_FUNCTION_NAME;
 import static dev.cel.optimizer.optimizers.DefaultOptimizerConstants.CEL_HAS_FIELD_FUNCTION_NAME;
+import static java.lang.Math.max;
 
 import com.google.auto.value.AutoValue;
 import com.google.common.annotations.VisibleForTesting;
@@ -244,8 +245,13 @@ public OptimizationResult optimize(CelAbstractSyntaxTree ast, Cel cel) {
       return OptimizationResult.create(ast);
     }
 
-    MonotonicIdGenerator idGenerator =
-        CelExprIdGeneratorFactory.newMonotonicIdGenerator(navAst.getRoot().maxId());
+    // Macro calls may hold nodes absent from the expression (e.g. comprehension variables) whose
+    // IDs exceed its own, so new IDs must start past them too.
+    long maxId = navAst.getRoot().maxId();
+    for (CelMutableExpr macroCall : astToModify.source().getMacroCalls().values()) {
+      maxId = max(maxId, CelNavigableMutableExpr.fromExpr(macroCall).maxId());
+    }
+    MonotonicIdGenerator idGenerator = CelExprIdGeneratorFactory.newMonotonicIdGenerator(maxId);
 
     int iterationCount = 0;
     for (CelNavigableMutableExpr topNode : topOfChainSelects) {
diff --git a/optimizer/src/test/java/dev/cel/optimizer/optimizers/SelectOptimizerTest.java b/optimizer/src/test/java/dev/cel/optimizer/optimizers/SelectOptimizerTest.java
index 85dfcb613..10c2fad43 100644
--- a/optimizer/src/test/java/dev/cel/optimizer/optimizers/SelectOptimizerTest.java
+++ b/optimizer/src/test/java/dev/cel/optimizer/optimizers/SelectOptimizerTest.java
@@ -1727,6 +1727,27 @@ public void optimize_optionalMapKeySelection_rewritesMapField() throws Exception
             "cel.@attribute(msg, [[61, \"map_string_string\", -1, [9, 9]]], {\"\": \"\"}).?key");
   }
 
+  @Test
+  public void optimize_afterConstantFolding_preservesMacroCalls() throws Exception {
+    CelOptimizer optimizer =
+        CelOptimizerFactory.standardCelOptimizerBuilder(cel)
+            .addAstOptimizers(
+                ConstantFoldingOptimizer.getInstance(),
+                SelectOptimizer.newInstance(
+                    PROTO2_TEST_ALL_TYPES_DESCRIPTOR.getFile(),
+                    NestedTestAllTypes.getDescriptor().getFile()))
+            .build();
+    CelAbstractSyntaxTree ast =
+        cel.compile("[nested_msg].map(n, n.child.payload.single_int64 + (1 + 1))").getAst();
+
+    CelAbstractSyntaxTree optimizedAst = optimizer.optimize(ast);
+
+    assertThat(CEL_UNPARSER.unparse(optimizedAst))
+        .isEqualTo(
+            "[nested_msg].map(n, cel.@attribute(n, [[1, \"child\", 11], [2, \"payload\", 11],"
+                + " [2, \"single_int64\", 3]], -64) + 2)");
+  }
+
   private static CelOptimizer newSelectOptimizer(Cel cel) {
     return CelOptimizerFactory.standardCelOptimizerBuilder(cel)
         .addAstOptimizers(