Support JSON related WKTs in select optimized ASTs

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

@@ -98,10 +98,9 @@
  *
  * <p>The 3rd argument of {@code cel.@attribute} is a typed expression that binds the result type
  * {@code T}. For singular scalar leaves it is also the field's default value (e.g. {@code 0} or a
- * proto2 custom default). For repeated, map, and message leaves (including {@code
- * google.protobuf.Duration}, {@code google.protobuf.Timestamp}, and wrapper types) it is a
- * type-only dummy (e.g. {@code [0]}, {@code {"": 0}}, {@code Msg{}}) that the runtime never
- * evaluates.
+ * proto2 custom default). For repeated, map, and message leaves (including well-known types such as
+ * {@code google.protobuf.Duration} and {@code google.protobuf.Struct}) it is a type-only dummy
+ * (e.g. {@code [0]}, {@code {"": 0}}, {@code Msg{}}) that the runtime never evaluates.
  *
  * <p>Field presence paths ({@code cel.@hasField}) represent each step as a 2-tuple: {@code
  * [field_num, field_name]}.
@@ -280,8 +279,11 @@ private void rewriteSelectChain(
         throw new UnsupportedOperationException(
             "Optimization of Group fields is unsupported: " + field.getFullName());
       }
-      if (field.getType() == FieldDescriptor.Type.MESSAGE) {
-        checkUnsupportedMessageType(field.getMessageType().getFullName(), field.getFullName());
+      // TODO: Support Any.
+      if (field.getType() == FieldDescriptor.Type.MESSAGE
+          && field.getMessageType().getFullName().equals(CelTypes.ANY_MESSAGE)) {
+        throw new UnsupportedOperationException(
+            "Optimization of Any fields is currently unimplemented: " + field.getFullName());
       }
 
       CelMutableList qualifierElements =
@@ -345,26 +347,6 @@ private boolean isTopOfSelectChain(CelNavigableMutableAst navAst, CelNavigableMu
         && !node.parent().flatMap(parent -> getOptimizableField(navAst, parent)).isPresent();
   }
 
-  // TODO: Support the remaining well-known types below.
-  private static void checkUnsupportedMessageType(String messageFullName, String fieldFullName) {
-    if (messageFullName.equals(CelTypes.STRUCT_MESSAGE)) {
-      throw new UnsupportedOperationException(
-          "Optimization of Struct fields is currently unimplemented: " + fieldFullName);
-    }
-    if (messageFullName.equals(CelTypes.LIST_VALUE_MESSAGE)) {
-      throw new UnsupportedOperationException(
-          "Optimization of ListValue fields is currently unimplemented: " + fieldFullName);
-    }
-    if (messageFullName.equals(CelTypes.VALUE_MESSAGE)) {
-      throw new UnsupportedOperationException(
-          "Optimization of Value fields is currently unimplemented: " + fieldFullName);
-    }
-    if (messageFullName.equals(CelTypes.ANY_MESSAGE)) {
-      throw new UnsupportedOperationException(
-          "Optimization of Any fields is currently unimplemented: " + fieldFullName);
-    }
-  }
-
   private Optional<FieldDescriptor> getOptimizableField(
       CelNavigableMutableAst navAst, CelNavigableMutableExpr node) {
     if (node.getKind() != Kind.SELECT) {

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

@@ -51,7 +51,6 @@
 import dev.cel.common.navigation.CelNavigableMutableAst;
 import dev.cel.common.types.ListType;
 import dev.cel.common.types.MapType;
-import dev.cel.common.types.NullableType;
 import dev.cel.common.types.SimpleType;
 import dev.cel.common.types.StructTypeReference;
 import dev.cel.common.values.CelByteString;
@@ -309,6 +308,15 @@ private enum RewriteTestCase {
     PROTO3_WRAPPER(
         "msg.single_int64_wrapper",
         "cel.@attribute(msg, [[105, \"single_int64_wrapper\", 11]], google.protobuf.Int64Value{})"),
+    PROTO3_STRUCT(
+        "msg.single_struct.key",
+        "cel.@attribute(msg, [[103, \"single_struct\", 11]], google.protobuf.Struct{}).key"),
+    PROTO3_VALUE(
+        "msg.single_value",
+        "cel.@attribute(msg, [[104, \"single_value\", 11]], google.protobuf.Value{})"),
+    PROTO3_LIST_VALUE(
+        "msg.list_value[0]",
+        "cel.@attribute(msg, [[114, \"list_value\", 11]], google.protobuf.ListValue{})[0]"),
 
     // Map selects
     MAP_STRING_STRING(
@@ -795,6 +803,24 @@ private enum NativeSelectEvaluationTestCase {
         "msg.single_int64_wrapper",
         ImmutableMap.of("msg", TestAllTypes.getDefaultInstance()),
         NullValue.NULL_VALUE),
+    PROTO3_STRUCT_CONVERTS_TO_MAP(
+        "msg.single_struct",
+        ImmutableMap.of(
+            "msg",
+            TestAllTypes.newBuilder()
+                .setSingleStruct(
+                    Struct.newBuilder()
+                        .putFields("a", Value.newBuilder().setNumberValue(1.5).build()))
+                .build()),
+        ImmutableMap.of("a", 1.5)),
+    PROTO3_VALUE_UNSET_IS_NULL(
+        "msg.single_value",
+        ImmutableMap.of("msg", TestAllTypes.getDefaultInstance()),
+        NullValue.NULL_VALUE),
+    PROTO3_LIST_VALUE_UNSET_IS_EMPTY(
+        "msg.list_value",
+        ImmutableMap.of("msg", TestAllTypes.getDefaultInstance()),
+        ImmutableList.of()),
     PROTO3_MAP_OF_ANY_UNPACKS_VALUE(
         "msg.map_string_any['k']",
         ImmutableMap.of(
@@ -1083,31 +1109,18 @@ public void optimize_iterationLimitReached_throws() throws Exception {
     assertThat(e).hasMessageThat().isEqualTo("Max iteration count reached.");
   }
 
-  private enum UnsupportedWellKnownFieldTestCase {
-    STRUCT("msg.single_struct", "Optimization of Struct fields is currently unimplemented"),
-    LIST_VALUE("msg.list_value", "Optimization of ListValue fields is currently unimplemented"),
-    VALUE("msg.single_value", "Optimization of Value fields is currently unimplemented"),
-    ANY("msg.single_any", "Optimization of Any fields is currently unimplemented");
-
-    private final String expression;
-    private final String expectedMessageSubstring;
-
-    UnsupportedWellKnownFieldTestCase(String expression, String expectedMessageSubstring) {
-      this.expression = expression;
-      this.expectedMessageSubstring = expectedMessageSubstring;
-    }
-  }
-
   @Test
-  public void optimize_unsupportedWellKnownField_throwsUnsupportedOperationException(
-      @TestParameter UnsupportedWellKnownFieldTestCase testCase) throws Exception {
-    CelAbstractSyntaxTree ast = cel.compile(testCase.expression).getAst();
+  public void optimize_unsupportedWellKnownField_throwsUnsupportedOperationException()
+      throws Exception {
+    CelAbstractSyntaxTree ast = cel.compile("msg.single_any").getAst();
     SelectOptimizer optimizer = SelectOptimizer.newInstance(TestAllTypes.getDescriptor().getFile());
 
     UnsupportedOperationException e =
         assertThrows(UnsupportedOperationException.class, () -> optimizer.optimize(ast, cel));
 
-    assertThat(e).hasMessageThat().contains(testCase.expectedMessageSubstring);
+    assertThat(e)
+        .hasMessageThat()
+        .contains("Optimization of Any fields is currently unimplemented");
   }
 
   @Test
@@ -1349,12 +1362,20 @@ public void optimize_toParsedExpr_matchesExpectedSerializedProto() throws Except
   }
 
   @Test
-  public void optimize_wrapperFieldSelect_preservesNullableResultType() throws Exception {
-    CelAbstractSyntaxTree ast = cel.compile("msg.single_int64_wrapper").getAst();
+  public void optimize_wellKnownTypeFieldSelect_preservesResultType(
+      @TestParameter({
+            "msg.single_int64_wrapper",
+            "msg.single_struct",
+            "msg.single_value",
+            "msg.list_value"
+          })
+          String expression)
+      throws Exception {
+    CelAbstractSyntaxTree ast = cel.compile(expression).getAst();
 
     CelAbstractSyntaxTree optimizedAst = celOptimizer.optimize(ast);
 
-    assertThat(optimizedAst.getResultType()).isEqualTo(NullableType.create(SimpleType.INT));
+    assertThat(optimizedAst.getResultType()).isEqualTo(ast.getResultType());
   }
 
   @Test

runtime/src/main/java/dev/cel/runtime/planner/OptimizedSelectPlanner.java

@@ -18,7 +18,6 @@
 import static com.google.common.base.Preconditions.checkNotNull;
 
 import com.google.common.collect.ImmutableList;
-import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Iterables;
 import com.google.common.primitives.UnsignedLong;
 import com.google.errorprone.annotations.Immutable;
@@ -47,17 +46,6 @@ final class OptimizedSelectPlanner {
   static final String CEL_ATTRIBUTE_FUNCTION_NAME = "cel.@attribute";
   static final String CEL_HAS_FIELD_FUNCTION_NAME = "cel.@hasField";
 
-  /**
-   * Well-known message types whose CEL semantics (Any unpacking, JSON value conversion) are not
-   * implemented by the optimized traversal.
-   */
-  private static final ImmutableSet<String> UNSUPPORTED_WELL_KNOWN_TYPE_IDENTS =
-      ImmutableSet.of(
-          CelTypes.ANY_MESSAGE,
-          CelTypes.STRUCT_MESSAGE,
-          CelTypes.VALUE_MESSAGE,
-          CelTypes.LIST_VALUE_MESSAGE);
-
   private final AttributeFactory attributeFactory;
   private final CelValueConverter celValueConverter;
 
@@ -299,8 +287,9 @@ private static String parseMessageProtoTypeName(CelExpr expr) {
 
   private static String parseLeafMessageProtoTypeName(CelExpr expr) {
     String protoTypeName = parseMessageProtoTypeName(expr);
+    // TODO: Support Any.
     checkArgument(
-        !UNSUPPORTED_WELL_KNOWN_TYPE_IDENTS.contains(protoTypeName),
+        !protoTypeName.equals(CelTypes.ANY_MESSAGE),
         "Leaf well-known type '%s' is not supported by the select-optimized runtime",
         protoTypeName);
     return protoTypeName;

runtime/src/test/java/dev/cel/runtime/CelLiteRuntimeVersionSkewTest.java

@@ -27,7 +27,10 @@
 import com.google.protobuf.CodedOutputStream;
 import com.google.protobuf.ExtensionRegistryLite;
 import com.google.protobuf.Int64Value;
+import com.google.protobuf.NullValue;
+import com.google.protobuf.Struct;
 import com.google.protobuf.UnknownFieldSet;
+import com.google.protobuf.Value;
 import com.google.testing.junit.testparameterinjector.TestParameter;
 import com.google.testing.junit.testparameterinjector.TestParameterInjector;
 import dev.cel.bundle.Cel;
@@ -2201,6 +2204,16 @@ private enum DescriptorlessEvaluationTestCase {
             .addRepeatedInt64Wrapper(Int64Value.of(0))
             .build(),
         true),
+    JSON_TYPES(
+        "msg.single_struct.a == 1.0 && msg.single_value == null && has(msg.single_value)"
+            + " && msg.list_value == [] && !has(msg.list_value) && msg.repeated_value == [null]",
+        TestAllTypes.newBuilder()
+            .setSingleStruct(
+                Struct.newBuilder().putFields("a", Value.newBuilder().setNumberValue(1).build()))
+            .setSingleValue(Value.newBuilder().setNullValue(NullValue.NULL_VALUE))
+            .addRepeatedValue(Value.newBuilder().setNullValue(NullValue.NULL_VALUE))
+            .build(),
+        true),
     ;
 
     private final String expression;

runtime/src/test/java/dev/cel/runtime/planner/OptimizedSelectPlannerTest.java

@@ -837,14 +837,8 @@ public void plan_invalidAst_mapFieldWithEmptyMapDummy_throwsEvaluationException(
   }
 
   @Test
-  public void plan_invalidAst_repeatedMessageWithUnsupportedWellKnownType_throwsEvaluationException(
-      @TestParameter({
-            "google.protobuf.Any",
-            "google.protobuf.Struct",
-            "google.protobuf.Value",
-            "google.protobuf.ListValue"
-          })
-          String unsupportedProtoTypeName) {
+  public void
+      plan_invalidAst_repeatedMessageWithUnsupportedWellKnownType_throwsEvaluationException() {
     CelAbstractSyntaxTree ast =
         CelAbstractSyntaxTree.newParsedAst(
             CelExpr.ofCall(
@@ -867,7 +861,7 @@ public void plan_invalidAst_repeatedMessageWithUnsupportedWellKnownType_throwsEv
                     CelExpr.ofList(
                         8L,
                         ImmutableList.of(
-                            CelExpr.ofStruct(9L, unsupportedProtoTypeName, ImmutableList.of())),
+                            CelExpr.ofStruct(9L, "google.protobuf.Any", ImmutableList.of())),
                         ImmutableList.of()))),
             CelSource.newBuilder().build());
 
@@ -876,7 +870,7 @@ public void plan_invalidAst_repeatedMessageWithUnsupportedWellKnownType_throwsEv
     assertThat(e).hasCauseThat().isInstanceOf(IllegalArgumentException.class);
     assertThat(e)
         .hasMessageThat()
-        .contains("Leaf well-known type '" + unsupportedProtoTypeName + "' is not supported");
+        .contains("Leaf well-known type 'google.protobuf.Any' is not supported");
   }
 
   @Test