Avoid unnecessary AST renumbering in CEL optimizers when no changes are made.

PiperOrigin-RevId: 976189920

Author: CEL Dev Team

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

@@ -39,6 +39,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.");
@@ -49,16 +50,18 @@ public CelAbstractSyntaxTree optimize(CelAbstractSyntaxTree ast) throws CelOptim
     try {
       for (CelAstOptimizer optimizer : astOptimizers) {
         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);
       }
     } catch (CelValidationException e) {
       throw new CelOptimizationException(

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

@@ -121,6 +121,7 @@ private static CelMutableExpr newOptionalNoneExpr() {
   }
 
   @Override
+  @SuppressWarnings("ReferenceEquality")
   public OptimizationResult optimize(CelAbstractSyntaxTree ast, Cel cel)
       throws CelOptimizationException {
     CelBuilder builder = cel.toCelBuilder();
@@ -134,12 +135,17 @@ public OptimizationResult optimize(CelAbstractSyntaxTree ast, Cel cel)
     // Override the environment's expected type to generally allow all subtrees to be folded.
     Cel optimizerEnv = builder.setResultType(SimpleType.DYN).build();
 
-    CelMutableAst mutableAst = CelMutableAst.fromCelAst(ast);
-    ImmutableMap<String, CelType> identTypes = precomputeIdentTypes(mutableAst);
+    CelMutableAst initialMutableAst = CelMutableAst.fromCelAst(ast);
+    ImmutableMap<String, CelType> identTypes = precomputeIdentTypes(initialMutableAst);
 
-    mutableAst = foldConstants(optimizerEnv, valueProvider, identTypes, mutableAst);
+    CelMutableAst mutableAst =
+        foldConstants(optimizerEnv, valueProvider, identTypes, initialMutableAst);
     mutableAst = pruneOptionalElements(mutableAst);
 
+    if (mutableAst == initialMutableAst) {
+      return OptimizationResult.create(ast);
+    }
+
     return OptimizationResult.create(astMutator.renumberIdsConsecutively(mutableAst).toParsedAst());
   }
 
@@ -735,10 +741,21 @@ private CelMutableAst pruneOptionalListElements(CelMutableAst mutableAst, CelMut
       updatedIndicesBuilder.add(newOptIndex);
     }
 
+    // An optional list is modified if:
+    // 1. An optional.none() was dropped - it this case, the updatedElements.size() decreases.
+    // 2. An optional.of(literal) was unwrapped into a regular element - in this case,
+    // updatedIndices.size() decreases.
+    // If both counts are unchanged, neither case occurred, and we can return the original AST.
+    ImmutableList<CelMutableExpr> updatedElements = updatedElemBuilder.build();
+    ImmutableList<Integer> updatedIndices = updatedIndicesBuilder.build();
+    if (updatedElements.size() == list.elements().size()
+        && updatedIndices.size() == list.optionalIndices().size()) {
+      return mutableAst;
+    }
+
     return astMutator.replaceSubtree(
         mutableAst,
-        CelMutableExpr.ofList(
-            CelMutableList.create(updatedElemBuilder.build(), updatedIndicesBuilder.build())),
+        CelMutableExpr.ofList(CelMutableList.create(updatedElements, updatedIndices)),
         expr.id());
   }
 

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

@@ -101,8 +101,10 @@ public static InliningOptimizer newInstance(
   }
 
   @Override
+  @SuppressWarnings("ReferenceEquality")
   public OptimizationResult optimize(CelAbstractSyntaxTree ast, Cel cel) {
-    CelMutableAst mutableAst = CelMutableAst.fromCelAst(ast);
+    CelMutableAst initialMutableAst = CelMutableAst.fromCelAst(ast);
+    CelMutableAst mutableAst = initialMutableAst;
     for (InlineVariable inlineVariable : inlineVariables) {
       mutableAst =
           astMutator.mutateUntilFixedPoint(
@@ -125,6 +127,10 @@ public OptimizationResult optimize(CelAbstractSyntaxTree ast, Cel cel) {
               });
     }
 
+    if (mutableAst == initialMutableAst) {
+      return OptimizationResult.create(ast);
+    }
+
     return OptimizationResult.create(astMutator.renumberIdsConsecutively(mutableAst).toParsedAst());
   }
 

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

@@ -702,6 +702,31 @@ 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);
+  }
+
   /**
    * Converts AST containing cel.block related test functions to internal functions (e.g: cel.block
    * -> cel.@block)