Fix a bug in the assertAstIdCorrectness method and avoid re-checking the AST in CelOptimizer if it was not modified.

PiperOrigin-RevId: 976199136

Author: Cristina Borza

optimizer/src/main/java/dev/cel/optimizer/CelOptimizerImpl.java

@@ -15,6 +15,7 @@
 package dev.cel.optimizer;
 
 import static com.google.common.base.Preconditions.checkNotNull;
+import static com.google.common.base.Preconditions.checkState;
 
 import com.google.common.collect.ImmutableSet;
 import dev.cel.bundle.Cel;
@@ -44,6 +45,7 @@ final class CelOptimizerImpl implements CelOptimizer {
   }
 
   @Override
+  @SuppressWarnings("ReferenceEquality")
   public CelAbstractSyntaxTree optimize(CelAbstractSyntaxTree ast) throws CelOptimizationException {
     if (!ast.isChecked()) {
       throw new IllegalArgumentException("AST must be type-checked.");
@@ -64,16 +66,18 @@ public CelAbstractSyntaxTree optimize(CelAbstractSyntaxTree ast) throws CelOptim
 
           OptimizationResult result = optimizer.optimize(optimizedAst, celOptimizerEnv);
 
-          if (!result.newFunctionDecls().isEmpty() || !result.newVarDecls().isEmpty()) {
-            celOptimizerEnv =
-                celOptimizerEnv
-                    .toCelBuilder()
-                    .addVarDeclarations(result.newVarDecls())
-                    .addFunctionDeclarations(result.newFunctionDecls())
-                    .build();
+          if (result.optimizedAst() != optimizedAst) {
+            if (!result.newFunctionDecls().isEmpty() || !result.newVarDecls().isEmpty()) {
+              celOptimizerEnv =
+                  celOptimizerEnv
+                      .toCelBuilder()
+                      .addVarDeclarations(result.newVarDecls())
+                      .addFunctionDeclarations(result.newFunctionDecls())
+                      .build();
+            }
+            optimizedAst = celOptimizerEnv.check(result.optimizedAst()).getAst();
+            assertAstIdCorrectness(optimizedAst);
           }
-          optimizedAst = celOptimizerEnv.check(result.optimizedAst()).getAst();
-          assertAstIdCorrectness(optimizedAst);
 
           for (CelOptimizerListener listener : listeners) {
             listener.onPassEnd(optimizer, preAst, optimizedAst);
@@ -131,19 +135,26 @@ private static void assertAstIdCorrectness(CelAbstractSyntaxTree ast) {
                   return;
                 }
 
-                if (astExpr.exprKind().getKind().equals(Kind.COMPREHENSION)) {
-                  if (!macroExpr.exprKind().getKind().equals(Kind.NOT_SET)) {
-                    throw new IllegalStateException(
-                        String.format(
-                            "Expected macro call node %d to be NOT_SET for comprehension, but"
-                                + " was %s.",
-                            macroExpr.id(), macroExpr.exprKind().getKind()));
-                  }
+                if (macroExpr.exprKind().getKind().equals(Kind.NOT_SET)) {
+                  // If a macro node is NOT_SET, its ID must be present in the main AST.
+                  checkState(
+                      ast.getSource().getMacroCalls().containsKey(macroExpr.id()),
+                      "Expected macro call node %s to be present in macro calls map, but was not.",
+                      macroExpr.id());
+                } else if (astExpr.exprKind().getKind().equals(Kind.COMPREHENSION)) {
+                  // We encountered something other than NOT_SET in macro source for comprehension
+                  // node. This is an error.
+                  throw new IllegalStateException(
+                      String.format(
+                          "Expected macro call node %d to be NOT_SET for comprehension, but was"
+                              + " %s.",
+                          macroExpr.id(), macroExpr.exprKind().getKind()));
                 } else if (!macroExpr.exprKind().getKind().equals(astExpr.exprKind().getKind())) {
+                  // Otherwise for all cases, the AST node should match exactly.
                   throw new IllegalStateException(
                       String.format(
-                          "Macro call node %d kind mismatch: expected %s (from AST), but was %s"
-                              + " (in macro call).",
+                          "Macro call node %d kind mismatch: expected %s (from AST), but was %s (in"
+                              + " macro call).",
                           macroExpr.id(),
                           astExpr.exprKind().getKind(),
                           macroExpr.exprKind().getKind()));

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

@@ -702,6 +702,66 @@ public void block_lazyEvaluationContainsError_cleansUpCycleState() throws Except
     assertThat(e).hasMessageThat().doesNotContain("Cycle detected");
   }
 
+  @Test
+  public void cse_nestedMacro_noOp_assertAstIdCorrectness() throws Exception {
+    Cel cel =
+        runtimeFlavor
+            .builder()
+            .addVar("x", SimpleType.DYN)
+            .setStandardMacros(CelStandardMacro.STANDARD_MACROS)
+            .setOptions(CelOptions.current().populateMacroCalls(true).build())
+            .addCompilerLibraries(CelExtensions.comprehensions())
+            .addRuntimeLibraries(CelExtensions.comprehensions())
+            .build();
+    CelOptimizer celOptimizer =
+        CelOptimizerFactory.standardCelOptimizerBuilder(cel)
+            .addAstOptimizers(SubexpressionOptimizer.getInstance())
+            .build();
+    CelAbstractSyntaxTree ast =
+        cel.compile("[{}, {\"a\": 1}, {\"b\": 2}].filter(m, has(x.a))").getAst();
+
+    CelAbstractSyntaxTree optimizedAst = celOptimizer.optimize(ast);
+
+    assertThat(CEL_UNPARSER.unparse(optimizedAst))
+        .isEqualTo("[{}, {\"a\": 1}, {\"b\": 2}].filter(m, has(x.a))");
+    assertThat(optimizedAst).isSameInstanceAs(ast);
+  }
+
+  @Test
+  public void cse_nestedMacro_withOptimization_assertAstIdCorrectness() throws Exception {
+    Cel cel =
+        runtimeFlavor
+            .builder()
+            .addVar("x", SimpleType.DYN)
+            .setStandardMacros(CelStandardMacro.STANDARD_MACROS)
+            .setOptions(
+                CelOptions.current()
+                    .populateMacroCalls(true)
+                    .enableHeterogeneousNumericComparisons(true)
+                    .build())
+            .addCompilerLibraries(CelExtensions.comprehensions())
+            .addRuntimeLibraries(CelExtensions.comprehensions())
+            .build();
+    CelOptimizer celOptimizer =
+        CelOptimizerFactory.standardCelOptimizerBuilder(cel)
+            .addAstOptimizers(
+                SubexpressionOptimizer.newInstance(
+                    SubexpressionOptimizerOptions.newBuilder().populateMacroCalls(true).build()))
+            .build();
+    CelAbstractSyntaxTree ast =
+        cel.compile(
+                "[{}, {\"a\": 1}, {\"b\": 2}].filter(m, has(x.a)) == [{}, {\"a\": 1}, {\"b\":"
+                    + " 2}].filter(m, has(x.a))")
+            .getAst();
+
+    CelAbstractSyntaxTree optimizedAst = celOptimizer.optimize(ast);
+
+    assertThat(CEL_UNPARSER.unparse(optimizedAst))
+        .isEqualTo(
+            "cel.@block([[{}, {\"a\": 1}, {\"b\": 2}].filter(@it:0:0, has(x.a))], @index0 =="
+                + " @index0)");
+  }
+
   /**
    * Converts AST containing cel.block related test functions to internal functions (e.g: cel.block
    * -> cel.@block)