Change comprehension variable mangling logic to always generate a unique index per comprehension expr

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

@@ -20,9 +20,7 @@
 
 import com.google.auto.value.AutoValue;
 import com.google.common.base.Preconditions;
-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;
@@ -261,43 +259,22 @@ public MangledComprehensionAst mangleComprehensionIdentifierNames(
                           "Unexpected CelNavigableMutableExpr collision");
                     },
                     LinkedHashMap::new));
-    int iterCount = 0;
 
     // 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()) {
-      iterCount++;
-      CelNavigableMutableExpr comprehensionNode = comprehensionEntry.getKey();
-      MangledComprehensionType comprehensionEntryType = comprehensionEntry.getValue();
-
-      CelMutableExpr comprehensionExpr = comprehensionNode.expr();
-      int comprehensionNestingLevel = countComprehensionNestingLevel(comprehensionNode);
-      MangledComprehensionName mangledComprehensionName;
-      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);
-      }
-      mangledIdentNamesToType.put(mangledComprehensionName, comprehensionEntryType);
+      String mangledIterVarName = newIterVarPrefix + ":" + iterCount;
+      String mangledResultName = newResultPrefix + ":" + iterCount;
+      MangledComprehensionName mangledComprehensionName =
+          MangledComprehensionName.of(mangledIterVarName, mangledResultName);
+      mangledIdentNamesToType.put(mangledComprehensionName, comprehensionEntry.getValue());
 
+      CelMutableExpr comprehensionExpr = comprehensionEntry.getKey().expr();
       String iterVar = comprehensionExpr.comprehension().iterVar();
       String accuVar = comprehensionExpr.comprehension().accuVar();
       mutatedComprehensionExpr =
@@ -315,6 +292,7 @@ public MangledComprehensionAst mangleComprehensionIdentifierNames(
               iterVar,
               mangledComprehensionName,
               comprehensionExpr.id());
+      iterCount++;
     }
 
     if (iterCount >= iterationLimit) {
@@ -822,19 +800,6 @@ private static long getMaxId(CelMutableExpr mutableExpr) {
         .orElseThrow(NoSuchElementException::new);
   }
 
-  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;
-  }
-
   /**
    * Intermediate value class to store the mangled identifiers for iteration variable and the
    * comprehension result.

optimizer/src/test/java/dev/cel/optimizer/AstMutatorTest.java

@@ -706,7 +706,7 @@ public void list_replaceElement() throws Exception {
   }
 
   @Test
-  public void createStruct_replaceValue() throws Exception {
+  public void struct_replaceValue() throws Exception {
     // Tree shape (brackets are expr IDs):
     //        TestAllTypes [1]
     //             single_int64 [2]
@@ -722,7 +722,7 @@ public void createStruct_replaceValue() throws Exception {
   }
 
   @Test
-  public void createMap_replaceKey() throws Exception {
+  public void map_replaceKey() throws Exception {
     // Tree shape (brackets are expr IDs):
     //        map [1]
     //       map_entry [2]
@@ -737,7 +737,7 @@ public void createMap_replaceKey() throws Exception {
   }
 
   @Test
-  public void createMap_replaceValue() throws Exception {
+  public void map_replaceValue() throws Exception {
     // Tree shape (brackets are expr IDs):
     //        map [1]
     //       map_entry [2]
@@ -808,15 +808,15 @@ public void mangleComprehensionVariable_singleMacro() throws Exception {
     assertThat(mangledAst.getExpr().toString())
         .isEqualTo(
             "COMPREHENSION [13] {\n"
-                + "  iter_var: @c0:0\n"
+                + "  iter_var: @c:0\n"
                 + "  iter_range: {\n"
                 + "    LIST [1] {\n"
                 + "      elements: {\n"
                 + "        CONSTANT [2] { value: false }\n"
                 + "      }\n"
                 + "    }\n"
                 + "  }\n"
-                + "  accu_var: @x0:0\n"
+                + "  accu_var: @x:0\n"
                 + "  accu_init: {\n"
                 + "    CONSTANT [6] { value: false }\n"
                 + "  }\n"
@@ -828,7 +828,7 @@ public void mangleComprehensionVariable_singleMacro() throws Exception {
                 + "          function: !_\n"
                 + "          args: {\n"
                 + "            IDENT [7] {\n"
-                + "              name: @x0:0\n"
+                + "              name: @x:0\n"
                 + "            }\n"
                 + "          }\n"
                 + "        }\n"
@@ -840,21 +840,22 @@ public void mangleComprehensionVariable_singleMacro() throws Exception {
                 + "      function: _||_\n"
                 + "      args: {\n"
                 + "        IDENT [10] {\n"
-                + "          name: @x0:0\n"
+                + "          name: @x:0\n"
                 + "        }\n"
                 + "        IDENT [5] {\n"
-                + "          name: @c0:0\n"
+                + "          name: @c:0\n"
                 + "        }\n"
                 + "      }\n"
                 + "    }\n"
                 + "  }\n"
                 + "  result: {\n"
                 + "    IDENT [12] {\n"
-                + "      name: @x0:0\n"
+                + "      name: @x:0\n"
                 + "    }\n"
                 + "  }\n"
                 + "}");
-    assertThat(CEL_UNPARSER.unparse(mangledAst)).isEqualTo("[false].exists(@c0:0, @c0:0)");
+
+    assertThat(CEL_UNPARSER.unparse(mangledAst)).isEqualTo("[false].exists(@c:0, @c:0)");
     assertThat(CEL.createProgram(CEL.check(mangledAst).getAst()).eval()).isEqualTo(false);
     assertConsistentMacroCalls(ast);
   }
@@ -891,7 +892,7 @@ public void mangleComprehensionVariable_nestedMacroWithShadowedVariables() throw
     assertThat(mangledAst.getExpr().toString())
         .isEqualTo(
             "COMPREHENSION [27] {\n"
-                + "  iter_var: @c0:0\n"
+                + "  iter_var: @c:1\n"
                 + "  iter_range: {\n"
                 + "    LIST [1] {\n"
                 + "      elements: {\n"
@@ -901,7 +902,7 @@ public void mangleComprehensionVariable_nestedMacroWithShadowedVariables() throw
                 + "      }\n"
                 + "    }\n"
                 + "  }\n"
-                + "  accu_var: @x0:0\n"
+                + "  accu_var: @x:1\n"
                 + "  accu_init: {\n"
                 + "    CONSTANT [20] { value: false }\n"
                 + "  }\n"
@@ -913,7 +914,7 @@ public void mangleComprehensionVariable_nestedMacroWithShadowedVariables() throw
                 + "          function: !_\n"
                 + "          args: {\n"
                 + "            IDENT [21] {\n"
-                + "              name: @x0:0\n"
+                + "              name: @x:1\n"
                 + "            }\n"
                 + "          }\n"
                 + "        }\n"
@@ -925,20 +926,20 @@ public void mangleComprehensionVariable_nestedMacroWithShadowedVariables() throw
                 + "      function: _||_\n"
                 + "      args: {\n"
                 + "        IDENT [24] {\n"
-                + "          name: @x0:0\n"
+                + "          name: @x:1\n"
                 + "        }\n"
                 + "        COMPREHENSION [19] {\n"
-                + "          iter_var: @c1:0\n"
+                + "          iter_var: @c:0\n"
                 + "          iter_range: {\n"
                 + "            LIST [5] {\n"
                 + "              elements: {\n"
                 + "                IDENT [6] {\n"
-                + "                  name: @c0:0\n"
+                + "                  name: @c:1\n"
                 + "                }\n"
                 + "              }\n"
                 + "            }\n"
                 + "          }\n"
-                + "          accu_var: @x1:0\n"
+                + "          accu_var: @x:0\n"
                 + "          accu_init: {\n"
                 + "            CONSTANT [12] { value: false }\n"
                 + "          }\n"
@@ -950,7 +951,7 @@ public void mangleComprehensionVariable_nestedMacroWithShadowedVariables() throw
                 + "                  function: !_\n"
                 + "                  args: {\n"
                 + "                    IDENT [13] {\n"
-                + "                      name: @x1:0\n"
+                + "                      name: @x:0\n"
                 + "                    }\n"
                 + "                  }\n"
                 + "                }\n"
@@ -962,13 +963,13 @@ public void mangleComprehensionVariable_nestedMacroWithShadowedVariables() throw
                 + "              function: _||_\n"
                 + "              args: {\n"
                 + "                IDENT [16] {\n"
-                + "                  name: @x1:0\n"
+                + "                  name: @x:0\n"
                 + "                }\n"
                 + "                CALL [10] {\n"
                 + "                  function: _==_\n"
                 + "                  args: {\n"
                 + "                    IDENT [9] {\n"
-                + "                      name: @c1:0\n"
+                + "                      name: @c:0\n"
                 + "                    }\n"
                 + "                    CONSTANT [11] { value: 1 }\n"
                 + "                  }\n"
@@ -978,7 +979,7 @@ public void mangleComprehensionVariable_nestedMacroWithShadowedVariables() throw
                 + "          }\n"
                 + "          result: {\n"
                 + "            IDENT [18] {\n"
-                + "              name: @x1:0\n"
+                + "              name: @x:0\n"
                 + "            }\n"
                 + "          }\n"
                 + "        }\n"
@@ -987,12 +988,13 @@ public void mangleComprehensionVariable_nestedMacroWithShadowedVariables() throw
                 + "  }\n"
                 + "  result: {\n"
                 + "    IDENT [26] {\n"
-                + "      name: @x0:0\n"
+                + "      name: @x:1\n"
                 + "    }\n"
                 + "  }\n"
                 + "}");
+
     assertThat(CEL_UNPARSER.unparse(mangledAst))
-        .isEqualTo("[x].exists(@c0:0, [@c0:0].exists(@c1:0, @c1:0 == 1))");
+        .isEqualTo("[x].exists(@c:1, [@c:1].exists(@c:0, @c:0 == 1))");
     assertThat(CEL.createProgram(CEL.check(mangledAst).getAst()).eval(ImmutableMap.of("x", 1)))
         .isEqualTo(true);
     assertConsistentMacroCalls(ast);

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

@@ -456,6 +456,9 @@ private enum CseTestCase {
     MULTIPLE_MACROS_2(
         "[[1].exists(i, i > 0)] + [[1].exists(j, j > 0)] + [['a'].exists(k, k == 'a')] +"
             + " [['a'].exists(l, l == 'a')] == [true, true, true, true]"),
+    MULTIPLE_MACROS_3(
+        "[1].exists(i, i > 0) && [1].exists(j, j > 0) && [1].exists(k, k > 1) && [2].exists(l, l >"
+            + " 1)"),
     NESTED_MACROS("[1,2,3].map(i, [1, 2, 3].map(i, i + 1)) == [[2, 3, 4], [2, 3, 4], [2, 3, 4]]"),
     NESTED_MACROS_2("[1, 2].map(y, [1, 2, 3].filter(x, x == y)) == [[1], [2]]"),
     INCLUSION_LIST("1 in [1,2,3] && 2 in [1,2,3] && 3 in [3, [1,2,3]] && 1 in [1,2,3]"),

optimizer/src/test/resources/constfold_before_subexpression_unparsed.baseline

@@ -302,6 +302,22 @@ Result: true
 [BLOCK_RECURSION_DEPTH_8]: true
 [BLOCK_RECURSION_DEPTH_9]: true
 
+Test case: MULTIPLE_MACROS_3
+Source: [1].exists(i, i > 0) && [1].exists(j, j > 0) && [1].exists(k, k > 1) && [2].exists(l, l > 1)
+=====>
+Result: false
+[CASCADED_BINDS]: false
+[BLOCK_COMMON_SUBEXPR_ONLY]: false
+[BLOCK_RECURSION_DEPTH_1]: false
+[BLOCK_RECURSION_DEPTH_2]: false
+[BLOCK_RECURSION_DEPTH_3]: false
+[BLOCK_RECURSION_DEPTH_4]: false
+[BLOCK_RECURSION_DEPTH_5]: false
+[BLOCK_RECURSION_DEPTH_6]: false
+[BLOCK_RECURSION_DEPTH_7]: false
+[BLOCK_RECURSION_DEPTH_8]: false
+[BLOCK_RECURSION_DEPTH_9]: false
+
 Test case: NESTED_MACROS
 Source: [1,2,3].map(i, [1, 2, 3].map(i, i + 1)) == [[2, 3, 4], [2, 3, 4], [2, 3, 4]]
 =====>
@@ -386,17 +402,17 @@ Test case: MACRO_SHADOWED_VARIABLE
 Source: [x - 1 > 3 ? x - 1 : 5].exists(x, x - 1 > 3) || x - 1 > 3
 =====>
 Result: true
-[CASCADED_BINDS]: cel.bind(@r0, x - 1, cel.bind(@r1, @r0 > 3, [@r1 ? @r0 : 5].exists(@c0:0, @c0:0 - 1 > 3) || @r1))
-[BLOCK_COMMON_SUBEXPR_ONLY]: cel.@block([x - 1, @index0 > 3], [@index1 ? @index0 : 5].exists(@c0:0, @c0:0 - 1 > 3) || @index1)
-[BLOCK_RECURSION_DEPTH_1]: cel.@block([x - 1, @index0 > 3, @index1 ? @index0 : 5, [@index2], @c0:0 - 1, @index4 > 3, @x0:0 || @index5], @index3.exists(@c0:0, @index5) || @index1)
-[BLOCK_RECURSION_DEPTH_2]: cel.@block([x - 1 > 3, @index0 ? (x - 1) : 5, @c0:0 - 1 > 3, [@index1], @x0:0 || @index2], @index3.exists(@c0:0, @index2) || @index0)
-[BLOCK_RECURSION_DEPTH_3]: cel.@block([x - 1 > 3, [@index0 ? (x - 1) : 5], @x0:0 || @c0:0 - 1 > 3, @index1.exists(@c0:0, @c0:0 - 1 > 3)], @index3 || @index0)
-[BLOCK_RECURSION_DEPTH_4]: cel.@block([x - 1 > 3, [@index0 ? (x - 1) : 5].exists(@c0:0, @c0:0 - 1 > 3)], @index1 || @index0)
-[BLOCK_RECURSION_DEPTH_5]: cel.@block([x - 1 > 3], [@index0 ? (x - 1) : 5].exists(@c0:0, @c0:0 - 1 > 3) || @index0)
-[BLOCK_RECURSION_DEPTH_6]: cel.@block([x - 1 > 3], [@index0 ? (x - 1) : 5].exists(@c0:0, @c0:0 - 1 > 3) || @index0)
-[BLOCK_RECURSION_DEPTH_7]: cel.@block([x - 1 > 3], [@index0 ? (x - 1) : 5].exists(@c0:0, @c0:0 - 1 > 3) || @index0)
-[BLOCK_RECURSION_DEPTH_8]: cel.@block([x - 1 > 3], [@index0 ? (x - 1) : 5].exists(@c0:0, @c0:0 - 1 > 3) || @index0)
-[BLOCK_RECURSION_DEPTH_9]: cel.@block([x - 1 > 3], [@index0 ? (x - 1) : 5].exists(@c0:0, @c0:0 - 1 > 3) || @index0)
+[CASCADED_BINDS]: cel.bind(@r0, x - 1, cel.bind(@r1, @r0 > 3, [@r1 ? @r0 : 5].exists(@c:0, @c:0 - 1 > 3) || @r1))
+[BLOCK_COMMON_SUBEXPR_ONLY]: cel.@block([x - 1, @index0 > 3], [@index1 ? @index0 : 5].exists(@c:0, @c:0 - 1 > 3) || @index1)
+[BLOCK_RECURSION_DEPTH_1]: cel.@block([x - 1, @index0 > 3, @index1 ? @index0 : 5, [@index2], @c:0 - 1, @index4 > 3, @x:0 || @index5], @index3.exists(@c:0, @index5) || @index1)
+[BLOCK_RECURSION_DEPTH_2]: cel.@block([x - 1 > 3, @index0 ? (x - 1) : 5, @c:0 - 1 > 3, [@index1], @x:0 || @index2], @index3.exists(@c:0, @index2) || @index0)
+[BLOCK_RECURSION_DEPTH_3]: cel.@block([x - 1 > 3, [@index0 ? (x - 1) : 5], @x:0 || @c:0 - 1 > 3, @index1.exists(@c:0, @c:0 - 1 > 3)], @index3 || @index0)
+[BLOCK_RECURSION_DEPTH_4]: cel.@block([x - 1 > 3, [@index0 ? (x - 1) : 5].exists(@c:0, @c:0 - 1 > 3)], @index1 || @index0)
+[BLOCK_RECURSION_DEPTH_5]: cel.@block([x - 1 > 3], [@index0 ? (x - 1) : 5].exists(@c:0, @c:0 - 1 > 3) || @index0)
+[BLOCK_RECURSION_DEPTH_6]: cel.@block([x - 1 > 3], [@index0 ? (x - 1) : 5].exists(@c:0, @c:0 - 1 > 3) || @index0)
+[BLOCK_RECURSION_DEPTH_7]: cel.@block([x - 1 > 3], [@index0 ? (x - 1) : 5].exists(@c:0, @c:0 - 1 > 3) || @index0)
+[BLOCK_RECURSION_DEPTH_8]: cel.@block([x - 1 > 3], [@index0 ? (x - 1) : 5].exists(@c:0, @c:0 - 1 > 3) || @index0)
+[BLOCK_RECURSION_DEPTH_9]: cel.@block([x - 1 > 3], [@index0 ? (x - 1) : 5].exists(@c:0, @c:0 - 1 > 3) || @index0)
 
 Test case: MACRO_SHADOWED_VARIABLE_2
 Source: ["foo", "bar"].map(x, [x + x, x + x]).map(x, [x + x, x + x])
@@ -668,4 +684,4 @@ Result: 31
 [BLOCK_RECURSION_DEPTH_6]: cel.@block([pure_custom_func(msg.oneof_type.payload.single_int64), @index0 + pure_custom_func(msg.oneof_type.payload.single_int32) + @index0], @index1 + pure_custom_func(msg.single_int64))
 [BLOCK_RECURSION_DEPTH_7]: cel.@block([pure_custom_func(msg.oneof_type.payload.single_int64)], @index0 + pure_custom_func(msg.oneof_type.payload.single_int32) + @index0 + pure_custom_func(msg.single_int64))
 [BLOCK_RECURSION_DEPTH_8]: cel.@block([pure_custom_func(msg.oneof_type.payload.single_int64)], @index0 + pure_custom_func(msg.oneof_type.payload.single_int32) + @index0 + pure_custom_func(msg.single_int64))
-[BLOCK_RECURSION_DEPTH_9]: cel.@block([pure_custom_func(msg.oneof_type.payload.single_int64)], @index0 + pure_custom_func(msg.oneof_type.payload.single_int32) + @index0 + pure_custom_func(msg.single_int64))
+[BLOCK_RECURSION_DEPTH_9]: cel.@block([pure_custom_func(msg.oneof_type.payload.single_int64)], @index0 + pure_custom_func(msg.oneof_type.payload.single_int32) + @index0 + pure_custom_func(msg.single_int64))