Add select-optimized AST conformance test targets for planner and lite runtimes

conformance/src/test/java/dev/cel/conformance/BUILD.bazel

@@ -41,6 +41,8 @@ java_library(
         "//extensions:optional_library",
         "//extensions:sets_runtime_library",
         "//extensions:strings_runtime_library",
+        "//optimizer",
+        "//optimizer/optimizers:select_optimizer",
         "//parser:macro",
         "//parser:parser_builder",
         "//parser:parser_factory",
@@ -81,6 +83,8 @@ java_library(
         "//:java_truth",
         "//common/values:proto_message_lite_value_provider",
         "//compiler:compiler_builder",
+        "//optimizer",
+        "//optimizer/optimizers:select_optimizer",
         "//parser:parser_factory",
         "//runtime:lite_runtime",
         "//runtime:lite_runtime_factory",
@@ -189,7 +193,13 @@ _TESTS_TO_SKIP_PLANNER = [
     "enums/strong_proto3",
 ]
 
-_TESTS_TO_SKIP_LITE = _TESTS_TO_SKIP_PLANNER + [
+_TESTS_TO_SKIP_PLANNER_SELECT_OPT = _TESTS_TO_SKIP_PLANNER + [
+    # TODO: Optional field selection on proto messages is not yet supported by
+    # SelectOptimizer.
+    "optionals",
+]
+
+_TESTS_TO_SKIP_LITE_SELECT_OPT = _TESTS_TO_SKIP_PLANNER + [
     # Optionals are not supported in CelLiteRuntime.
     "optionals",
     "block_ext/basic/optional_list,optional_map,optional_map_chained,optional_message",
@@ -199,18 +209,12 @@ _TESTS_TO_SKIP_LITE = _TESTS_TO_SKIP_PLANNER + [
     "proto2/extensions_has",
     "proto2/extensions_get",
 
-    # Custom proto2 default values on unset fields require SelectOptimization in CelLiteRuntime.
-    "proto2/empty_field/scalar_with_default",
-
     # TODO: Proto equality after dyn() unwrapping is not yet supported in
     # CelLiteRuntime.
     "comparisons/eq_literal/not_eq_dyn_proto2_msg_null,not_eq_dyn_proto3_msg_null",
     "comparisons/eq_wrapper/eq_proto_different_types",
     "comparisons/ne_literal/ne_proto_different_types",
 
-    # TODO: google.protobuf.Any unpacking is not yet supported in CelLiteRuntime.
-    "dynamic/any/literal_empty,var",
-
     # TODO: Message creation with prepopulated fields is not yet supported in
     # CelLiteRuntime.
     "comparisons/eq_wrapper/eq_bool,eq_bytes,eq_double,eq_float,eq_int32,eq_int64,eq_string,eq_uint32,eq_uint64",
@@ -259,6 +263,11 @@ _TESTS_TO_SKIP_LITE = _TESTS_TO_SKIP_PLANNER + [
     "wrappers",
 ]
 
+_TESTS_TO_SKIP_LITE = _TESTS_TO_SKIP_LITE_SELECT_OPT + [
+    # Custom proto2 default values on unset fields require SelectOptimization in CelLiteRuntime.
+    "proto2/empty_field/scalar_with_default",
+]
+
 conformance_test(
     name = "conformance",
     data = _ALL_TESTS,
@@ -285,9 +294,25 @@ conformance_test(
     skip_tests = _TESTS_TO_SKIP_PLANNER,
 )
 
+conformance_test(
+    name = "conformance_planner_select_opt",
+    data = _ALL_TESTS,
+    enable_select_optimization = True,
+    runtime = RUNTIME.PLANNER,
+    skip_tests = _TESTS_TO_SKIP_PLANNER_SELECT_OPT,
+)
+
 conformance_test(
     name = "conformance_lite",
     data = _ALL_TESTS,
     runtime = RUNTIME.LITE,
     skip_tests = _TESTS_TO_SKIP_LITE,
 )
+
+conformance_test(
+    name = "conformance_lite_select_opt",
+    data = _ALL_TESTS,
+    enable_select_optimization = True,
+    runtime = RUNTIME.LITE,
+    skip_tests = _TESTS_TO_SKIP_LITE_SELECT_OPT,
+)

conformance/src/test/java/dev/cel/conformance/ConformanceTest.java

@@ -62,6 +62,8 @@
 import dev.cel.extensions.CelOptionalLibrary;
 import dev.cel.extensions.CelSetsRuntimeLibrary;
 import dev.cel.extensions.CelStringRuntimeLibrary;
+import dev.cel.optimizer.CelOptimizerFactory;
+import dev.cel.optimizer.optimizers.SelectOptimizer;
 import dev.cel.parser.CelMacro;
 import dev.cel.parser.CelMacroExpander;
 import dev.cel.parser.CelMacroExprFactory;
@@ -72,6 +74,7 @@
 import dev.cel.runtime.CelEvaluationException;
 import dev.cel.runtime.CelLiteRuntimeFactory;
 import dev.cel.runtime.CelLiteRuntimeLibrary;
+import dev.cel.runtime.CelRuntime;
 import dev.cel.runtime.CelRuntimeBuilder;
 import dev.cel.runtime.CelRuntimeFactory;
 import dev.cel.runtime.CelRuntimeImpl;
@@ -145,6 +148,12 @@ public final class ConformanceTest extends Statement {
           .addFileDescriptors(CONFORMANCE_FILE_DESCRIPTORS)
           .build();
 
+  private static final SelectOptimizer SELECT_OPTIMIZER =
+      SelectOptimizer.newInstance(CONFORMANCE_FILE_DESCRIPTORS);
+
+  private static final CelRuntime UNUSED_RUNTIME =
+      CelRuntimeFactory.plannerRuntimeBuilder().build();
+
   enum RuntimeType {
     LEGACY,
     PLANNER,
@@ -230,6 +239,7 @@ private static SimpleTest defaultTestMatcherToTrueIfUnset(SimpleTest test) {
   private final SimpleTest test;
   private final boolean skip;
   private final RuntimeType runtimeType;
+  private final boolean enableSelectOptimization;
 
   public String getName() {
     return name;
@@ -241,23 +251,22 @@ public boolean shouldSkip() {
 
   @Override
   public void evaluate() throws Throwable {
-    CelValidationResult response = getParser(test).parse(test.getExpr(), test.getName());
-    assertThat(response.getErrors()).isEmpty();
-    if (!test.getDisableCheck()) {
-      response = getChecker(test).check(response.getAst());
-    }
+    CelParser parser = getParser(test);
+    CelValidationResult response = parser.parse(test.getExpr(), test.getName());
     assertThat(response.getErrors()).isEmpty();
     CelAbstractSyntaxTree ast = response.getAst();
-    if (ast.getExpr().callOrDefault().function().equals(CelBlock.FUNCTION_NAME)) {
-      ast =
-          CelAbstractSyntaxTree.newCheckedAst(
-              ast.getExpr(),
-              ast.getSource().toBuilder()
-                  .addAllExtensions(
-                      Extension.create("cel_block", Version.of(1, 1), Component.COMPONENT_RUNTIME))
-                  .build(),
-              ast.getReferenceMap(),
-              ast.getTypeMap());
+    if (!test.getDisableCheck()) {
+      CelChecker checker = getChecker(test);
+      response = checker.check(ast);
+      assertThat(response.getErrors()).isEmpty();
+      ast = response.getAst();
+      if (enableSelectOptimization) {
+        ast =
+            CelOptimizerFactory.standardCelOptimizerBuilder(parser, checker, UNUSED_RUNTIME)
+                .addAstOptimizers(SELECT_OPTIMIZER)
+                .build()
+                .optimize(ast);
+      }
     }
     Type resultType = CelProtoTypes.celTypeToType(ast.getResultType());
 
@@ -272,6 +281,18 @@ public void evaluate() throws Throwable {
       return;
     }
 
+    if (ast.getExpr().callOrDefault().function().equals(CelBlock.FUNCTION_NAME)) {
+      ast =
+          CelAbstractSyntaxTree.newCheckedAst(
+              ast.getExpr(),
+              ast.getSource().toBuilder()
+                  .addAllExtensions(
+                      Extension.create("cel_block", Version.of(1, 1), Component.COMPONENT_RUNTIME))
+                  .build(),
+              ast.getReferenceMap(),
+              ast.getTypeMap());
+    }
+
     ExprValue result = null;
     CelEvaluationException error = null;
     try {
@@ -343,11 +364,17 @@ private Program createProgram(CelAbstractSyntaxTree ast) throws CelEvaluationExc
     return builder.build().createProgram(ast);
   }
 
-  ConformanceTest(String name, SimpleTest test, boolean skip, RuntimeType runtimeType) {
+  ConformanceTest(
+      String name,
+      SimpleTest test,
+      boolean skip,
+      RuntimeType runtimeType,
+      boolean enableSelectOptimization) {
     this.name = checkNotNull(name);
     this.test = defaultTestMatcherToTrueIfUnset(checkNotNull(test));
     this.skip = skip;
     this.runtimeType = checkNotNull(runtimeType);
+    this.enableSelectOptimization = enableSelectOptimization;
   }
 
   /**

conformance/src/test/java/dev/cel/conformance/ConformanceTestRunner.java

@@ -44,6 +44,7 @@ public final class ConformanceTestRunner extends ParentRunner<ConformanceTest> {
   private final ImmutableSortedMap<String, SimpleTestFile> testFiles;
   private final ImmutableList<String> testsToSkip;
   private final RuntimeType runtimeType;
+  private final boolean enableSelectOptimization;
 
   private static ImmutableSortedMap<String, SimpleTestFile> loadTestFiles() {
     List<String> testPaths =
@@ -89,7 +90,9 @@ protected List<ConformanceTest> getChildren() {
         for (SimpleTest test : testSection.getTestList()) {
           String name =
               String.format("%s/%s/%s", testFile.getName(), testSection.getName(), test.getName());
-          tests.add(new ConformanceTest(name, test, shouldSkipTest(name), runtimeType));
+          tests.add(
+              new ConformanceTest(
+                  name, test, shouldSkipTest(name), runtimeType, enableSelectOptimization));
         }
       }
     }
@@ -128,5 +131,9 @@ public ConformanceTestRunner(Class<?> clazz) throws InitializationError {
     runtimeType =
         RuntimeType.valueOf(
             System.getProperty("dev.cel.conformance.ConformanceTests.runtime", "LEGACY"));
+    enableSelectOptimization =
+        Boolean.parseBoolean(
+            System.getProperty(
+                "dev.cel.conformance.ConformanceTests.enable_select_optimization", "false"));
   }
 }

conformance/src/test/java/dev/cel/conformance/conformance_test.bzl

@@ -38,11 +38,12 @@ def _expand_tests_to_skip(tests_to_skip):
             result.append(test_to_skip[0:slash] + part)
     return result
 
-def _conformance_test_args(data, skip_tests, runtime):
+def _conformance_test_args(data, skip_tests, runtime, enable_select_optimization):
     return [
         "-Ddev.cel.conformance.ConformanceTests.skip_tests={}".format(",".join(_expand_tests_to_skip(skip_tests))),
         "-Ddev.cel.conformance.ConformanceTests.tests={}".format(",".join(["$(location {})".format(t) for t in data])),
         "-Ddev.cel.conformance.ConformanceTests.runtime={}".format(runtime),
+        "-Ddev.cel.conformance.ConformanceTests.enable_select_optimization={}".format(enable_select_optimization),
     ]
 
 MODE = struct(
@@ -60,7 +61,13 @@ RUNTIME = struct(
     LITE = "LITE",
 )
 
-def conformance_test(name, data, mode = MODE.TEST, skip_tests = [], runtime = RUNTIME.LEGACY):
+def conformance_test(
+        name,
+        data,
+        mode = MODE.TEST,
+        skip_tests = [],
+        runtime = RUNTIME.LEGACY,
+        enable_select_optimization = False):
     """Executes conformance tests
 
     Args:
@@ -74,11 +81,12 @@ def conformance_test(name, data, mode = MODE.TEST, skip_tests = [], runtime = RU
                    exclude from the run.
        runtime: The CEL runtime implementation to evaluate against (`RUNTIME.LEGACY`,
                 `RUNTIME.PLANNER`, or `RUNTIME.LITE`).
+       enable_select_optimization: Whether to optimize checked ASTs with `SelectOptimizer`.
     """
     if mode == MODE.DASHBOARD:
         java_test(
             name = "_" + name,
-            jvm_flags = _conformance_test_args(data, skip_tests, runtime),
+            jvm_flags = _conformance_test_args(data, skip_tests, runtime, enable_select_optimization),
             data = data,
             size = "small",
             test_class = "dev.cel.conformance.ConformanceTests",
@@ -104,7 +112,7 @@ def conformance_test(name, data, mode = MODE.TEST, skip_tests = [], runtime = RU
     elif mode == MODE.TEST:
         java_test(
             name = name,
-            jvm_flags = _conformance_test_args(data, skip_tests, runtime),
+            jvm_flags = _conformance_test_args(data, skip_tests, runtime, enable_select_optimization),
             data = data,
             size = "small",
             test_class = "dev.cel.conformance.ConformanceTests",
@@ -113,7 +121,7 @@ def conformance_test(name, data, mode = MODE.TEST, skip_tests = [], runtime = RU
     elif mode == MODE.MAVEN_TEST:
         java_test(
             name = name,
-            jvm_flags = _conformance_test_args(data, skip_tests, runtime),
+            jvm_flags = _conformance_test_args(data, skip_tests, runtime, enable_select_optimization),
             data = data,
             size = "small",
             test_class = "dev.cel.conformance.ConformanceTests",