Hoist loop-invariant subexpressions in SubexpressionOptimizer

optimizer/src/main/java/dev/cel/optimizer/optimizers/BUILD.bazel

@@ -60,11 +60,13 @@ java_library(
         ":default_optimizer_constants",
         "//:auto_value",
         "//bundle:cel",
+        "//checker:standard_decl",
         "//common:cel_ast",
         "//common:cel_source",
         "//common:compiler_common",
         "//common:mutable_ast",
         "//common:mutable_source",
+        "//common:operator",
         "//common/ast",
         "//common/ast:cel_block",
         "//common/ast:mutable_expr",

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

@@ -403,7 +403,17 @@ private enum CseTestCase {
         "pure_custom_func(msg.oneof_type.payload.single_int64) +"
             + " pure_custom_func(msg.oneof_type.payload.single_int32) +"
             + " pure_custom_func(msg.oneof_type.payload.single_int64) +"
-            + " pure_custom_func(msg.single_int64)");
+            + " pure_custom_func(msg.single_int64)"),
+    VALUE_NUMBERING_COMMUTATIVE_ADD_CONVERSE_RELATION_LOGICAL_AND(
+        "(msg.single_int64 + msg.single_int32 > 5 && msg.oneof_type.payload.single_int64 <"
+            + " msg.single_int64 + msg.single_int32) == (msg.single_int32 + msg.single_int64 >"
+            + " msg.oneof_type.payload.single_int64 && 5 < msg.single_int32 + msg.single_int64)"),
+    VALUE_NUMBERING_COMMUTATIVE_MULTIPLY_EQUALITY_LOGICAL_OR(
+        "(msg.oneof_type.payload.single_int32 * msg.single_int64 == 24 || msg.single_int32 !="
+            + " msg.oneof_type.payload.single_int64) && (msg.oneof_type.payload.single_int64 !="
+            + " msg.single_int32 || 24 == msg.single_int64 * msg.oneof_type.payload.single_int32)"),
+    ;
+
     private final String source;
 
     CseTestCase(String source) {

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

@@ -79,8 +79,9 @@ private static Cel setupCelEnv(CelBuilder celBuilder) {
                 .populateMacroCalls(true)
                 .enableHeterogeneousNumericComparisons(true)
                 .build())
-        .addCompilerLibraries(CelExtensions.bindings(), CelExtensions.strings())
-        .addRuntimeLibraries(CelExtensions.strings())
+        .addCompilerLibraries(
+            CelExtensions.bindings(), CelExtensions.strings(), CelExtensions.optional())
+        .addRuntimeLibraries(CelExtensions.strings(), CelExtensions.optional())
         .addFunctionDeclarations(
             CelFunctionDecl.newFunctionDeclaration(
                 "non_pure_custom_func",
@@ -756,12 +757,12 @@ public void cse_nestedMacro_noOp_assertAstIdCorrectness() throws Exception {
             .addAstOptimizers(SubexpressionOptimizer.getInstance())
             .build();
     CelAbstractSyntaxTree ast =
-        cel.compile("[{}, {\"a\": 1}, {\"b\": 2}].filter(m, has(x.a))").getAst();
+        cel.compile("[{}, {\"a\": 1}, {\"b\": 2}].filter(m, has(m.a))").getAst();
 
     CelAbstractSyntaxTree optimizedAst = celOptimizer.optimize(ast);
 
     assertThat(CEL_UNPARSER.unparse(optimizedAst))
-        .isEqualTo("[{}, {\"a\": 1}, {\"b\": 2}].filter(m, has(x.a))");
+        .isEqualTo("[{}, {\"a\": 1}, {\"b\": 2}].filter(m, has(m.a))");
     assertThat(optimizedAst).isSameInstanceAs(ast);
   }
 
@@ -796,8 +797,8 @@ public void cse_nestedMacro_withOptimization_assertAstIdCorrectness() throws Exc
 
     assertThat(CEL_UNPARSER.unparse(optimizedAst))
         .isEqualTo(
-            "cel.@block([[{}, {\"a\": 1}, {\"b\": 2}].filter(@it:0:0, has(x.a))], @index0 =="
-                + " @index0)");
+            "cel.@block([has(x.a), [{}, {\"a\": 1}, {\"b\": 2}].filter(@it:0:0, @index0)], @index1"
+                + " == @index1)");
   }
 
   @Test
@@ -822,6 +823,254 @@ public void optimize_afterSelectOptimizer_throwsException(
         .contains("SubexpressionOptimizer must run before SelectOptimizer.");
   }
 
+  private enum ValueNumberingTestCase {
+    COMMUTATIVE_INT_ADD(
+        "(msg.single_int64 + msg.single_int32) * (msg.single_int32 + msg.single_int64)",
+        "cel.@block([msg.single_int64 + msg.single_int32], @index0 * @index0)"),
+    COMMUTATIVE_UINT_MULTIPLY(
+        "(msg.single_uint64 * msg.single_uint32) + (msg.single_uint32 * msg.single_uint64)",
+        "cel.@block([msg.single_uint64 * msg.single_uint32], @index0 + @index0)"),
+    COMMUTATIVE_DOUBLE_ADD(
+        "(msg.single_double + 1.5) * (1.5 + msg.single_double)",
+        "cel.@block([msg.single_double + 1.5], @index0 * @index0)"),
+    COMMUTATIVE_DURATION_TIMESTAMP_ADD(
+        "(msg.single_duration + msg.single_timestamp) == (msg.single_timestamp +"
+            + " msg.single_duration)",
+        "cel.@block([msg.single_duration + msg.single_timestamp], @index0 == @index0)"),
+    NON_COMMUTATIVE_INT_SUBTRACT_PRESERVED(
+        "(msg.single_int64 - msg.single_int32) == (msg.single_int32 - msg.single_int64)",
+        "cel.@block([msg.single_int32, msg.single_int64], @index1 - @index0 == @index0 - @index1)"),
+    NON_COMMUTATIVE_STRING_CONCAT_PRESERVED(
+        "(msg.single_string + \"a\") == (\"a\" + msg.single_string)",
+        "cel.@block([msg.single_string], @index0 + \"a\" == \"a\" + @index0)"),
+    NON_COMMUTATIVE_DYN_ADD_PRESERVED(
+        "(dyn(msg.single_string) + dyn(\"a\")) == (dyn(\"a\") + dyn(msg.single_string))",
+        "cel.@block([dyn(\"a\"), dyn(msg.single_string)], @index1 + @index0 == @index0 +"
+            + " @index1)"),
+    COMMUTATIVE_EQUALS(
+        "(msg.single_int64 == 1) && (1 == msg.single_int64)",
+        "cel.@block([msg.single_int64 == 1], @index0 && @index0)"),
+    COMMUTATIVE_NOT_EQUALS(
+        "(msg.single_int64 != 1) || (1 != msg.single_int64)",
+        "cel.@block([msg.single_int64 != 1], @index0 || @index0)"),
+    COMMUTATIVE_LOGICAL_AND(
+        "((msg.single_int64 == 1) && (msg.single_int32 == 2)) || ((msg.single_int32 == 2) &&"
+            + " (msg.single_int64 == 1))",
+        "cel.@block([msg.single_int64 == 1 && msg.single_int32 == 2], @index0 || @index0)"),
+    COMMUTATIVE_LOGICAL_OR(
+        "((msg.single_int64 == 1) || (msg.single_int32 == 2)) && ((msg.single_int32 == 2) ||"
+            + " (msg.single_int64 == 1))",
+        "cel.@block([msg.single_int64 == 1 || msg.single_int32 == 2], @index0 && @index0)"),
+    RELATIONAL_CONVERSE_LESS_GREATER(
+        "(msg.single_int64 < msg.single_int32) && (msg.single_int32 > msg.single_int64)",
+        "cel.@block([msg.single_int64 < msg.single_int32], @index0 && @index0)"),
+    RELATIONAL_CONVERSE_LESS_EQUALS_GREATER_EQUALS(
+        "(msg.single_int64 <= msg.single_int32) && (msg.single_int32 >= msg.single_int64)",
+        "cel.@block([msg.single_int64 <= msg.single_int32], @index0 && @index0)"),
+    NOT_EQUALS_WITHOUT_SHARED_EQUALS_PRESERVED(
+        "[msg.single_int64 != 1, msg.single_int64 != 1, msg.single_int64 == 1]",
+        "cel.@block([msg.single_int64 != 1], [@index0, @index0, msg.single_int64 == 1])"),
+    OPTIONAL_LIST_AND_MAP_DISTINCT_VALUE_NUMBERS(
+        "([?optional.of(msg.single_int64)].size() == [optional.of(msg.single_int64)].size()) &&"
+            + " ({?\"k\": optional.of(msg.single_int64)}.size() == {\"k\":"
+            + " optional.of(msg.single_int64)}.size())",
+        "cel.@block([optional.of(msg.single_int64)], [?@index0].size() == [@index0].size() &&"
+            + " {?\"k\": @index0}.size() == {\"k\": @index0}.size())"),
+    ;
+
+    private final String source;
+    private final String expectedUnparsed;
+
+    ValueNumberingTestCase(String source, String expectedUnparsed) {
+      this.source = source;
+      this.expectedUnparsed = expectedUnparsed;
+    }
+  }
+
+  @Test
+  public void cse_valueNumbering(@TestParameter ValueNumberingTestCase testCase) throws Exception {
+    CelAbstractSyntaxTree ast = cel.compile(testCase.source).getAst();
+    CelOptimizer optimizer =
+        newCseOptimizer(
+            SubexpressionOptimizerOptions.newBuilder().populateMacroCalls(true).build());
+    ImmutableMap<String, Object> input =
+        ImmutableMap.of(
+            "msg",
+            TestAllTypes.newBuilder()
+                .setSingleInt64(1L)
+                .setSingleInt32(2)
+                .setSingleUint64(3L)
+                .setSingleUint32(4)
+                .setSingleDouble(2.5)
+                .setSingleString("hello")
+                .build());
+
+    CelAbstractSyntaxTree optimizedAst = optimizer.optimize(ast);
+
+    assertThat(CEL_UNPARSER.unparse(optimizedAst)).isEqualTo(testCase.expectedUnparsed);
+    assertThat(cel.createProgram(optimizedAst).eval(input))
+        .isEqualTo(cel.createProgram(ast).eval(input));
+  }
+
+  private enum HoistLoopInvariantsTestCase {
+    NESTED_COMPREHENSION(
+        "msg.repeated_int64.all(i, msg.repeated_int32.exists(j, j == 2))",
+        "cel.@block([msg.repeated_int32.exists(@it:0:0, @it:0:0 == 2)],"
+            + " msg.repeated_int64.all(@it:1:0, @index0))"),
+    INVARIANT_ARITHMETIC(
+        "msg.repeated_int64.exists(i, i == msg.single_int64 * 2 + msg.single_int32)",
+        "cel.@block([msg.single_int64 * 2 + msg.single_int32],"
+            + " msg.repeated_int64.exists(@it:0:0, @it:0:0 == @index0))"),
+    ZERO_TRIP_ERROR_NOT_EVALUATED(
+        "msg.repeated_nested_message.exists(m, m.bb == 1 / (msg.single_int32 - 2))",
+        "cel.@block([1 / (msg.single_int32 - 2)],"
+            + " msg.repeated_nested_message.exists(@it:0:0, @it:0:0.bb == @index0))"),
+    MAP_ACCUMULATION_PRESERVED(
+        "msg.repeated_int64.map(i, msg.single_int64 * 2)",
+        "cel.@block([msg.single_int64 * 2], msg.repeated_int64.map(@it:0:0, @index0))"),
+    LIST_LITERAL(
+        "msg.repeated_int64.exists(i, i in [msg.single_int64, msg.single_int32])",
+        "cel.@block([[msg.single_int64, msg.single_int32]],"
+            + " msg.repeated_int64.exists(@it:0:0, @it:0:0 in @index0))"),
+    LIST_APPENDED_TO_ITER_VAR(
+        "[[1], [2]].map(l, l + [msg.single_int64])",
+        "cel.@block([[msg.single_int64]], [[1], [2]].map(@it:0:0, @it:0:0 + @index0))"),
+    INVARIANT_INSIDE_DUPLICATE(
+        "msg.repeated_int64.exists(i, i in [msg.single_int64, msg.single_int32]) ||"
+            + " msg.repeated_int64.exists(i, i in [msg.single_int64, msg.single_int32])",
+        "cel.@block([[msg.single_int64, msg.single_int32], msg.repeated_int64.exists(@it:0:0,"
+            + " @it:0:0 in @index0)], @index1 || @index1)"),
+    INVARIANT_OF_NESTED_COMPREHENSION(
+        "msg.repeated_int64.all(i, msg.repeated_int32.exists(j, j == msg.single_int64 * 2))",
+        "cel.@block([msg.single_int64 * 2, msg.repeated_int32.exists(@it:0:0, @it:0:0 =="
+            + " @index0)], msg.repeated_int64.all(@it:1:0, @index1))"),
+    SHARED_TERMS_HOISTED_ONCE(
+        "msg.repeated_int64.exists(i, i == msg.single_int64 + msg.single_int32 +"
+            + " msg.single_sint64) && [msg.single_int64, msg.single_int32, msg.single_sint64]"
+            + ".size() > 0",
+        "cel.@block([msg.single_int64, msg.single_int32, msg.single_sint64, @index0 + @index1 +"
+            + " @index2], msg.repeated_int64.exists(@it:0:0, @it:0:0 == @index3) && [@index0,"
+            + " @index1, @index2].size() > 0)"),
+    SHARED_WITH_ACCUMULATED_LIST(
+        "msg.repeated_int64.exists(i, i in [msg.single_int64]) &&"
+            + " msg.repeated_int64.map(i, msg.single_int64).size() > 0",
+        "cel.@block([msg.single_int64, [@index0], msg.repeated_int64], @index2.exists(@it:0:0,"
+            + " @it:0:0 in @index1) && @index2.map(@it:0:1, @index0).size() > 0)"),
+    ;
+
+    private final String source;
+    private final String expectedUnparsed;
+
+    HoistLoopInvariantsTestCase(String source, String expectedUnparsed) {
+      this.source = source;
+      this.expectedUnparsed = expectedUnparsed;
+    }
+  }
+
+  @Test
+  public void cse_hoistLoopInvariants(@TestParameter HoistLoopInvariantsTestCase testCase)
+      throws Exception {
+    CelAbstractSyntaxTree ast = cel.compile(testCase.source).getAst();
+    CelOptimizer optimizer =
+        newCseOptimizer(
+            SubexpressionOptimizerOptions.newBuilder().populateMacroCalls(true).build());
+    ImmutableMap<String, Object> input =
+        ImmutableMap.of(
+            "msg",
+            TestAllTypes.newBuilder()
+                .setSingleInt64(1L)
+                .setSingleInt32(2)
+                .addAllRepeatedInt64(ImmutableList.of(1L, 2L, 3L))
+                .addRepeatedInt32(2)
+                .build());
+
+    CelAbstractSyntaxTree optimizedAst = optimizer.optimize(ast);
+
+    assertThat(CEL_UNPARSER.unparse(optimizedAst)).isEqualTo(testCase.expectedUnparsed);
+    assertThat(cel.createProgram(optimizedAst).eval(input))
+        .isEqualTo(cel.createProgram(ast).eval(input));
+  }
+
+  @Test
+  public void cse_customOperatorOverloads_preservesSemantics(
+      @TestParameter({
+            "(msg.single_string + msg.single_int64) == (msg.single_int64 + msg.single_string)",
+            "(msg.single_string > msg.single_int64) == (msg.single_int64 < msg.single_string)"
+          })
+          String source)
+      throws Exception {
+    Cel cel =
+        this.cel
+            .toCelBuilder()
+            .addFunctionDeclarations(
+                CelFunctionDecl.newFunctionDeclaration(
+                    "_+_",
+                    newGlobalOverload(
+                        "add_string_int", SimpleType.INT, SimpleType.STRING, SimpleType.INT),
+                    newGlobalOverload(
+                        "add_int_string", SimpleType.INT, SimpleType.INT, SimpleType.STRING)),
+                CelFunctionDecl.newFunctionDeclaration(
+                    "_>_",
+                    newGlobalOverload(
+                        "greater_string_int", SimpleType.BOOL, SimpleType.STRING, SimpleType.INT)),
+                CelFunctionDecl.newFunctionDeclaration(
+                    "_<_",
+                    newGlobalOverload(
+                        "less_int_string", SimpleType.BOOL, SimpleType.INT, SimpleType.STRING)))
+            .addFunctionBindings(
+                CelFunctionBinding.from(
+                    "add_string_int", String.class, Long.class, (s, i) -> s.length() + i),
+                CelFunctionBinding.from(
+                    "add_int_string", Long.class, String.class, (i, s) -> i * 100 + s.length()),
+                CelFunctionBinding.from(
+                    "greater_string_int", String.class, Long.class, (s, i) -> s.length() > i),
+                CelFunctionBinding.from(
+                    "less_int_string", Long.class, String.class, (i, s) -> i >= 0))
+            .build();
+    CelOptimizer optimizer =
+        CelOptimizerFactory.standardCelOptimizerBuilder(cel)
+            .addAstOptimizers(SubexpressionOptimizer.getInstance())
+            .build();
+    CelAbstractSyntaxTree ast = cel.compile(source).getAst();
+    ImmutableMap<String, Object> input =
+        ImmutableMap.of(
+            "msg", TestAllTypes.newBuilder().setSingleString("a").setSingleInt64(1L).build());
+
+    CelAbstractSyntaxTree optimizedAst = optimizer.optimize(ast);
+
+    assertThat(cel.createProgram(optimizedAst).eval(input))
+        .isEqualTo(cel.createProgram(ast).eval(input));
+  }
+
+  @Test
+  public void cse_oneSidedCustomOverloadOnDyn_notMerged(
+      @TestParameter({"(x < 5) == (5 > x)", "(x + 1) == (1 + x)"}) String source) throws Exception {
+    // With x: dyn, only the left operand also resolves to the custom overload, so the right one is
+    // standard-only and swappable. The two must not share a value number.
+    Cel cel =
+        this.cel
+            .toCelBuilder()
+            .addFunctionDeclarations(
+                CelFunctionDecl.newFunctionDeclaration(
+                    "_<_",
+                    newGlobalOverload(
+                        "less_string_int", SimpleType.BOOL, SimpleType.STRING, SimpleType.INT)),
+                CelFunctionDecl.newFunctionDeclaration(
+                    "_+_",
+                    newGlobalOverload(
+                        "add_string_int", SimpleType.INT, SimpleType.STRING, SimpleType.INT)))
+            .build();
+    CelOptimizer optimizer =
+        CelOptimizerFactory.standardCelOptimizerBuilder(cel)
+            .addAstOptimizers(SubexpressionOptimizer.getInstance())
+            .build();
+    CelAbstractSyntaxTree ast = cel.compile(source).getAst();
+
+    CelAbstractSyntaxTree optimizedAst = optimizer.optimize(ast);
+
+    assertThat(CEL_UNPARSER.unparse(optimizedAst)).isEqualTo(CEL_UNPARSER.unparse(ast));
+  }
+
   /**
    * Converts AST containing cel.block related test functions to internal functions (e.g: cel.block
    * -> cel.@block)

optimizer/src/test/resources/constfold_before_subexpression_unparsed.baseline

@@ -314,4 +314,16 @@ Test case: CUSTOM_FUNCTION_ELIMINABLE
 Source: pure_custom_func(msg.oneof_type.payload.single_int64) + pure_custom_func(msg.oneof_type.payload.single_int32) + pure_custom_func(msg.oneof_type.payload.single_int64) + pure_custom_func(msg.single_int64)
 =====>
 Result: 31
-Unparsed: cel.@block([msg.oneof_type.payload, pure_custom_func(@index0.single_int64)], @index1 + pure_custom_func(@index0.single_int32) + @index1 + pure_custom_func(msg.single_int64))
+Unparsed: cel.@block([msg.oneof_type.payload, pure_custom_func(@index0.single_int64)], @index1 + pure_custom_func(@index0.single_int32) + @index1 + pure_custom_func(msg.single_int64))
+
+Test case: VALUE_NUMBERING_COMMUTATIVE_ADD_CONVERSE_RELATION_LOGICAL_AND
+Source: (msg.single_int64 + msg.single_int32 > 5 && msg.oneof_type.payload.single_int64 < msg.single_int64 + msg.single_int32) == (msg.single_int32 + msg.single_int64 > msg.oneof_type.payload.single_int64 && 5 < msg.single_int32 + msg.single_int64)
+=====>
+Result: true
+Unparsed: cel.@block([msg.single_int64 + msg.single_int32, @index0 > 5 && msg.oneof_type.payload.single_int64 < @index0], @index1 == @index1)
+
+Test case: VALUE_NUMBERING_COMMUTATIVE_MULTIPLY_EQUALITY_LOGICAL_OR
+Source: (msg.oneof_type.payload.single_int32 * msg.single_int64 == 24 || msg.single_int32 != msg.oneof_type.payload.single_int64) && (msg.oneof_type.payload.single_int64 != msg.single_int32 || 24 == msg.single_int64 * msg.oneof_type.payload.single_int32)
+=====>
+Result: true
+Unparsed: cel.@block([msg.oneof_type.payload, @index0.single_int32 * msg.single_int64 == 24 || msg.single_int32 != @index0.single_int64], @index1 && @index1)