Improve CSE for comprehensions by taking into their nesting level and types into account

PiperOrigin-RevId: 664956665

Author: l46kok

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

@@ -21,7 +21,9 @@
 import com.google.auto.value.AutoValue;
 import com.google.common.base.Preconditions;
 import com.google.common.base.Strings;
+import com.google.common.collect.HashBasedTable;
 import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.Table;
 import com.google.errorprone.annotations.Immutable;
 import dev.cel.common.CelAbstractSyntaxTree;
 import dev.cel.common.CelMutableAst;
@@ -270,17 +272,22 @@ public CelMutableAst renumberIdsConsecutively(CelMutableAst mutableAst) {
    * @param ast AST containing type-checked references
    * @param newIterVarPrefix Prefix to use for new iteration variable identifier name. For example,
    *     providing @c will produce @c0:0, @c0:1, @c1:0, @c2:0... as new names.
-   * @param newResultPrefix Prefix to use for new comprehensin result identifier names.
+   * @param newAccuVarPrefix Prefix to use for new accumulation variable identifier name.
+   * @param incrementSerially If true, indices for the mangled variables are incremented serially
+   *     per occurrence regardless of their nesting level or its types.
    */
   public MangledComprehensionAst mangleComprehensionIdentifierNames(
-      CelMutableAst ast, String newIterVarPrefix, String newResultPrefix) {
+      CelMutableAst ast,
+      String newIterVarPrefix,
+      String newAccuVarPrefix,
+      boolean incrementSerially) {
     CelNavigableMutableAst navigableMutableAst = CelNavigableMutableAst.fromAst(ast);
     Predicate<CelNavigableMutableExpr> comprehensionIdentifierPredicate = x -> true;
     comprehensionIdentifierPredicate =
         comprehensionIdentifierPredicate
             .and(node -> node.getKind().equals(Kind.COMPREHENSION))
             .and(node -> !node.expr().comprehension().iterVar().startsWith(newIterVarPrefix))
-            .and(node -> !node.expr().comprehension().accuVar().startsWith(newResultPrefix));
+            .and(node -> !node.expr().comprehension().accuVar().startsWith(newAccuVarPrefix));
 
     LinkedHashMap<CelNavigableMutableExpr, MangledComprehensionType> comprehensionsToMangle =
         navigableMutableAst
@@ -352,18 +359,43 @@ public MangledComprehensionAst mangleComprehensionIdentifierNames(
     // The map that we'll eventually return to the caller.
     HashMap<MangledComprehensionName, MangledComprehensionType> mangledIdentNamesToType =
         new HashMap<>();
+    // Intermediary table used for the purposes of generating a unique mangled variable name.
+    Table<Integer, MangledComprehensionType, MangledComprehensionName> comprehensionLevelToType =
+        HashBasedTable.create();
     CelMutableExpr mutatedComprehensionExpr = navigableMutableAst.getAst().expr();
     CelMutableSource newSource = navigableMutableAst.getAst().source();
     int iterCount = 0;
     for (Entry<CelNavigableMutableExpr, MangledComprehensionType> comprehensionEntry :
         comprehensionsToMangle.entrySet()) {
-      String mangledIterVarName = newIterVarPrefix + ":" + iterCount;
-      String mangledResultName = newResultPrefix + ":" + iterCount;
-      MangledComprehensionName mangledComprehensionName =
-          MangledComprehensionName.of(mangledIterVarName, mangledResultName);
-      mangledIdentNamesToType.put(mangledComprehensionName, comprehensionEntry.getValue());
+      CelNavigableMutableExpr comprehensionNode = comprehensionEntry.getKey();
+      MangledComprehensionType comprehensionEntryType = comprehensionEntry.getValue();
+
+      CelMutableExpr comprehensionExpr = comprehensionNode.expr();
+      MangledComprehensionName mangledComprehensionName;
+      if (incrementSerially) {
+        // In case of applying CSE via cascaded cel.binds, not only is mangling based on level/types
+        // meaningless (because all comprehensions are nested anyways, thus all indices would be
+        // uinque),
+        // it can lead to an erroneous result due to extracting a common subexpr with accu_var at
+        // the wrong scope.
+        // Example: "[1].exists(k, k > 1) && [2].exists(l, l > 1). The loop step for both branches
+        // are identical, but shouldn't be extracted.
+        String mangledIterVarName = newIterVarPrefix + ":" + iterCount;
+        String mangledResultName = newAccuVarPrefix + ":" + iterCount;
+        mangledComprehensionName =
+            MangledComprehensionName.of(mangledIterVarName, mangledResultName);
+        mangledIdentNamesToType.put(mangledComprehensionName, comprehensionEntry.getValue());
+      } else {
+        mangledComprehensionName =
+            getMangledComprehensionName(
+                newIterVarPrefix,
+                newAccuVarPrefix,
+                comprehensionNode,
+                comprehensionLevelToType,
+                comprehensionEntryType);
+      }
+      mangledIdentNamesToType.put(mangledComprehensionName, comprehensionEntryType);
 
-      CelMutableExpr comprehensionExpr = comprehensionEntry.getKey().expr();
       String iterVar = comprehensionExpr.comprehension().iterVar();
       String accuVar = comprehensionExpr.comprehension().accuVar();
       mutatedComprehensionExpr =
@@ -396,6 +428,45 @@ public MangledComprehensionAst mangleComprehensionIdentifierNames(
         ImmutableMap.copyOf(mangledIdentNamesToType));
   }
 
+  private static MangledComprehensionName getMangledComprehensionName(
+      String newIterVarPrefix,
+      String newResultPrefix,
+      CelNavigableMutableExpr comprehensionNode,
+      Table<Integer, MangledComprehensionType, MangledComprehensionName> comprehensionLevelToType,
+      MangledComprehensionType comprehensionEntryType) {
+    MangledComprehensionName mangledComprehensionName;
+    int comprehensionNestingLevel = countComprehensionNestingLevel(comprehensionNode);
+    if (comprehensionLevelToType.contains(comprehensionNestingLevel, comprehensionEntryType)) {
+      mangledComprehensionName =
+          comprehensionLevelToType.get(comprehensionNestingLevel, comprehensionEntryType);
+    } else {
+      // First time encountering the pair of <ComprehensionLevel, CelType>. Generate a unique
+      // mangled variable name for this.
+      int uniqueTypeIdx = comprehensionLevelToType.row(comprehensionNestingLevel).size();
+      String mangledIterVarName =
+          newIterVarPrefix + ":" + comprehensionNestingLevel + ":" + uniqueTypeIdx;
+      String mangledResultName =
+          newResultPrefix + ":" + comprehensionNestingLevel + ":" + uniqueTypeIdx;
+      mangledComprehensionName = MangledComprehensionName.of(mangledIterVarName, mangledResultName);
+      comprehensionLevelToType.put(
+          comprehensionNestingLevel, comprehensionEntryType, mangledComprehensionName);
+    }
+    return mangledComprehensionName;
+  }
+
+  private static int countComprehensionNestingLevel(CelNavigableMutableExpr comprehensionExpr) {
+    int nestedLevel = 0;
+    Optional<CelNavigableMutableExpr> maybeParent = comprehensionExpr.parent();
+    while (maybeParent.isPresent()) {
+      if (maybeParent.get().getKind().equals(Kind.COMPREHENSION)) {
+        nestedLevel++;
+      }
+
+      maybeParent = maybeParent.get().parent();
+    }
+    return nestedLevel;
+  }
+
   /**
    * Replaces a subtree in the given expression node. This operation is intended for AST
    * optimization purposes.

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

@@ -92,8 +92,8 @@ public class SubexpressionOptimizer implements CelAstOptimizer {
   private static final SubexpressionOptimizer INSTANCE =
       new SubexpressionOptimizer(SubexpressionOptimizerOptions.newBuilder().build());
   private static final String BIND_IDENTIFIER_PREFIX = "@r";
-  private static final String MANGLED_COMPREHENSION_IDENTIFIER_PREFIX = "@c";
-  private static final String MANGLED_COMPREHENSION_RESULT_PREFIX = "@x";
+  private static final String MANGLED_COMPREHENSION_ITER_VAR_PREFIX = "@it";
+  private static final String MANGLED_COMPREHENSION_ACCU_VAR_PREFIX = "@ac";
   private static final String CEL_BLOCK_FUNCTION = "cel.@block";
   private static final String BLOCK_INDEX_PREFIX = "@index";
   private static final Extension CEL_BLOCK_AST_EXTENSION_TAG =
@@ -138,8 +138,9 @@ private OptimizationResult optimizeUsingCelBlock(CelAbstractSyntaxTree ast, Cel
     MangledComprehensionAst mangledComprehensionAst =
         astMutator.mangleComprehensionIdentifierNames(
             astToModify,
-            MANGLED_COMPREHENSION_IDENTIFIER_PREFIX,
-            MANGLED_COMPREHENSION_RESULT_PREFIX);
+            MANGLED_COMPREHENSION_ITER_VAR_PREFIX,
+            MANGLED_COMPREHENSION_ACCU_VAR_PREFIX,
+            /* incrementSerially= */ false);
     astToModify = mangledComprehensionAst.mutableAst();
     CelMutableSource sourceToModify = astToModify.source();
 
@@ -339,8 +340,9 @@ private OptimizationResult optimizeUsingCelBind(CelAbstractSyntaxTree ast) {
         astMutator
             .mangleComprehensionIdentifierNames(
                 astToModify,
-                MANGLED_COMPREHENSION_IDENTIFIER_PREFIX,
-                MANGLED_COMPREHENSION_RESULT_PREFIX)
+                MANGLED_COMPREHENSION_ITER_VAR_PREFIX,
+                MANGLED_COMPREHENSION_ACCU_VAR_PREFIX,
+                /* incrementSerially= */ true)
             .mutableAst();
     CelMutableSource sourceToModify = astToModify.source();