Fail fast on optional field selection in SelectOptimizer

PiperOrigin-RevId: 998708843

Author: l46kok

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

@@ -128,6 +128,7 @@ java_library(
         "//common:cel_source",
         "//common:compiler_common",
         "//common:mutable_ast",
+        "//common:operator",
         "//common/ast",
         "//common/ast:expr_factory",
         "//common/ast:mutable_expr",

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

@@ -41,6 +41,7 @@
 import dev.cel.common.CelSource.Extension;
 import dev.cel.common.CelSource.Extension.Component;
 import dev.cel.common.CelSource.Extension.Version;
+import dev.cel.common.Operator;
 import dev.cel.common.ast.CelConstant;
 import dev.cel.common.ast.CelExpr.ExprKind.Kind;
 import dev.cel.common.ast.CelExprIdGeneratorFactory;
@@ -59,8 +60,10 @@
 import dev.cel.common.navigation.CelNavigableMutableExpr;
 import dev.cel.common.navigation.TraversalOrder;
 import dev.cel.common.types.CelKind;
+import dev.cel.common.types.CelType;
 import dev.cel.common.types.CelTypes;
 import dev.cel.common.types.ListType;
+import dev.cel.common.types.OptionalType;
 import dev.cel.common.types.SimpleType;
 import dev.cel.common.types.TypeParamType;
 import dev.cel.common.values.CelByteString;
@@ -122,7 +125,7 @@
  *
  * <p>Map indexing and non-protobuf selects pass through untouched.
  *
- * <p><b>Rename Resilience & Dynamic Type Limitations:</b>
+ * <p><b>Limitations:</b>
  *
  * <ul>
  *   <li><b>Protobuf Extensions:</b> Extension fields are not rename-resilient; runtime lookup
@@ -131,6 +134,11 @@
  *       from {@code google.protobuf.Any}, or evaluated via classless wire payloads ({@code
  *       RawProtoMessageLiteValue}), no cross-check between the embedded field number and field name
  *       is performed at runtime.
+ *   <li><b>Oneofs:</b> The lite runtime reads fields from the wire, so it reports every serialized
+ *       member of a oneof as present. Protobuf serializes at most one, unless an older-schema
+ *       binary merges or mutates the message, or the bytes are concatenated or altered by hand.
+ *   <li><b>Optionals:</b> Optional field selection on a message ({@code msg.?field}, or a field of
+ *       an optional message) is not yet supported and throws during optimization.
  * </ul>
  */
 public final class SelectOptimizer implements CelAstOptimizer {
@@ -213,6 +221,7 @@ public OptimizationResult optimize(CelAbstractSyntaxTree ast, Cel cel) {
 
     CelMutableAst astToModify = CelMutableAst.fromCelAst(ast);
     CelNavigableMutableAst navAst = CelNavigableMutableAst.fromAst(astToModify);
+    navAst.getRoot().allNodes().forEach(node -> checkNotOptionalFieldSelection(navAst, node));
     ImmutableList<CelNavigableMutableExpr> topOfChainSelects =
         navAst
             .getRoot()
@@ -259,7 +268,6 @@ private void rewriteSelectChain(
                 () -> new IllegalStateException("Expected optimizable field on select node"));
     fields.add(topField);
 
-    // TODO: Support optional field selection (_?._) once integrated with lite runtime.
     CelMutableExpr currentExpr = topNode.expr().select().operand();
     while (currentExpr.getKind() == Kind.SELECT) {
       CelMutableSelect select = currentExpr.select();
@@ -357,11 +365,43 @@ private Optional<FieldDescriptor> getOptimizableField(
 
   private Optional<FieldDescriptor> getOptimizableFieldForExpr(
       CelNavigableMutableAst navAst, CelMutableSelect select) {
-    return navAst
-        .getType(select.operand().id())
-        .filter(type -> type.kind() == CelKind.STRUCT)
-        .flatMap(type -> descriptorPool.findDescriptor(type.name()))
-        .map(desc -> desc.findFieldByName(select.field()));
+    return navAst.getType(select.operand().id()).flatMap(type -> findField(type, select.field()));
+  }
+
+  private Optional<FieldDescriptor> findField(CelType type, String field) {
+    if (type.kind() != CelKind.STRUCT) {
+      return Optional.empty();
+    }
+    return descriptorPool.findDescriptor(type.name()).map(desc -> desc.findFieldByName(field));
+  }
+
+  // TODO: Support optional field selection once integrated with lite runtime.
+  private void checkNotOptionalFieldSelection(
+      CelNavigableMutableAst navAst, CelNavigableMutableExpr node) {
+    Optional<FieldDescriptor> field = Optional.empty();
+    if (node.getKind() == Kind.SELECT) {
+      // A field of an optional message, e.g. `.bb` in `msg.repeated_nested_message[?0].bb`
+      CelMutableSelect select = node.expr().select();
+      field =
+          navAst
+              .getType(select.operand().id())
+              .filter(OptionalType.class::isInstance)
+              .flatMap(type -> findField(type.parameters().get(0), select.field()));
+    } else if (node.getKind() == Kind.CALL
+        && node.expr().call().function().equals(Operator.OPTIONAL_SELECT.getFunction())) {
+      // `msg.?field`, where `msg` may itself be optional
+      List<CelMutableExpr> args = node.expr().call().args();
+      field =
+          navAst
+              .getType(args.get(0).id())
+              .map(type -> type instanceof OptionalType ? type.parameters().get(0) : type)
+              .flatMap(type -> findField(type, args.get(1).constant().stringValue()));
+    }
+    if (field.isPresent()) {
+      throw new UnsupportedOperationException(
+          "Optimization of optional field selection is currently unimplemented: "
+              + field.get().getFullName());
+    }
   }
 
   private static CelMutableExpr resolveDummyOrDefExpr(

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

@@ -367,6 +367,10 @@ private enum RewriteTestCase {
             + "cel.@attribute(msg, [[227, \"map_string_message\", -1, [9, 11]]],"
             + " {\"\": cel.expr.conformance.proto3.TestAllTypes.NestedMessage{}}).key, "
             + "[[1, \"bb\"]])"),
+    PROTO_MAP_FIELD_SELECT_STOPS_AT_ANY_BOUNDARY(
+        "msg.map_string_any['k'].value",
+        "cel.@attribute(msg, [[311, \"map_string_any\", -1, [9, 11]]],"
+            + " {\"\": google.protobuf.Any{}})[\"k\"].value"),
 
     MIXED_BOOLEAN_EXPRESSION(
         "msg.single_int64 > 0 && has(msg.single_nested_message)",
@@ -1469,24 +1473,37 @@ public void optimize_comprehensionOverRepeatedMessageField_preservesStaticElemen
   }
 
   @Test
-  public void optimize_optionalSelect_passesThroughUntouched() throws Exception {
-    Cel celWithOptional =
-        cel.toCelBuilder()
-            .addCompilerLibraries(CelExtensions.optional())
-            .addRuntimeLibraries(CelExtensions.optional())
-            .build();
-    CelOptimizer optimizer =
-        CelOptimizerFactory.standardCelOptimizerBuilder(celWithOptional)
-            .addAstOptimizers(
-                SelectOptimizer.newInstance(
-                    SelectOptimizerOptions.newBuilder().build(),
-                    TestAllTypes.getDescriptor().getFile()))
-            .build();
-    CelAbstractSyntaxTree ast = celWithOptional.compile("msg.?single_nested_message.bb").getAst();
+  public void optimize_optionalFieldSelection_throwsUnsupportedOperationException(
+      @TestParameter({
+            "msg.?single_nested_message",
+            "msg.repeated_nested_message[?0].?bb",
+            "msg.repeated_nested_message[?0].bb"
+          })
+          String expression)
+      throws Exception {
+    Cel celWithOptional = cel.toCelBuilder().addCompilerLibraries(CelExtensions.optional()).build();
+    CelAbstractSyntaxTree ast = celWithOptional.compile(expression).getAst();
+    SelectOptimizer optimizer = SelectOptimizer.newInstance(TestAllTypes.getDescriptor().getFile());
 
-    CelAbstractSyntaxTree optimizedAst = optimizer.optimize(ast);
+    UnsupportedOperationException e =
+        assertThrows(
+            UnsupportedOperationException.class, () -> optimizer.optimize(ast, celWithOptional));
+
+    assertThat(e)
+        .hasMessageThat()
+        .contains("Optimization of optional field selection is currently unimplemented");
+  }
 
-    assertThat(optimizedAst.getExpr()).isEqualTo(ast.getExpr());
+  @Test
+  public void optimize_optionalMapKeySelection_rewritesMapField() throws Exception {
+    Cel celWithOptional = cel.toCelBuilder().addCompilerLibraries(CelExtensions.optional()).build();
+    CelAbstractSyntaxTree ast = celWithOptional.compile("msg.map_string_string.?key").getAst();
+
+    CelAbstractSyntaxTree optimizedAst = newSelectOptimizer(celWithOptional).optimize(ast);
+
+    assertThat(CEL_UNPARSER.unparse(optimizedAst))
+        .isEqualTo(
+            "cel.@attribute(msg, [[61, \"map_string_string\", -1, [9, 9]]], {\"\": \"\"}).?key");
   }
 
   private static CelOptimizer newSelectOptimizer(Cel cel) {