Support wrappers in select optimized ASTs

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

@@ -99,8 +99,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} and {@code google.protobuf.Timestamp}) it is a type-only dummy (e.g.
- * {@code [0]}, {@code {"": 0}}, {@code Msg{}}) that the runtime never evaluates.
+ * 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.
  *
  * <p>Field presence paths ({@code cel.@hasField}) represent each step as a 2-tuple: {@code
  * [field_num, field_name]}.
@@ -344,8 +345,7 @@ private boolean isTopOfSelectChain(CelNavigableMutableAst navAst, CelNavigableMu
         && !node.parent().flatMap(parent -> getOptimizableField(navAst, parent)).isPresent();
   }
 
-  // TODO: Support STRUCT_MESSAGE, LIST_VALUE_MESSAGE, VALUE_MESSAGE, ANY_MESSAGE,
-  // and wrapper types.
+  // 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(
@@ -363,10 +363,6 @@ private static void checkUnsupportedMessageType(String messageFullName, String f
       throw new UnsupportedOperationException(
           "Optimization of Any fields is currently unimplemented: " + fieldFullName);
     }
-    if (CelTypes.isWrapperType(messageFullName)) {
-      throw new UnsupportedOperationException(
-          "Optimization of wrapper fields is currently unimplemented: " + fieldFullName);
-    }
   }
 
   private Optional<FieldDescriptor> getOptimizableField(

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

@@ -22,6 +22,7 @@ java_library(
         "//common/ast",
         "//common/navigation:mutable_navigation",
         "//common/types",
+        "//common/values",
         "//common/values:cel_byte_string",
         "//extensions",
         "//extensions:optional_library",

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

@@ -51,9 +51,11 @@
 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;
+import dev.cel.common.values.NullValue;
 import dev.cel.expr.conformance.proto2.NestedTestAllTypes;
 import dev.cel.expr.conformance.proto2.TestAllTypesProto;
 import dev.cel.expr.conformance.proto3.TestAllTypes;
@@ -304,6 +306,9 @@ private enum RewriteTestCase {
         "msg.single_duration == duration(\"1h\")",
         "cel.@attribute(msg, [[101, \"single_duration\", 11]], google.protobuf.Duration{}) =="
             + " duration(\"1h\")"),
+    PROTO3_WRAPPER(
+        "msg.single_int64_wrapper",
+        "cel.@attribute(msg, [[105, \"single_int64_wrapper\", 11]], google.protobuf.Int64Value{})"),
 
     // Map selects
     MAP_STRING_STRING(
@@ -781,6 +786,15 @@ private enum NativeSelectEvaluationTestCase {
         "msg.single_timestamp",
         ImmutableMap.of("msg", TestAllTypes.getDefaultInstance()),
         Instant.EPOCH),
+    PROTO3_WRAPPER_SET_TO_ZERO_UNWRAPS(
+        "msg.single_int64_wrapper",
+        ImmutableMap.of(
+            "msg", TestAllTypes.newBuilder().setSingleInt64Wrapper(Int64Value.of(0)).build()),
+        0L),
+    PROTO3_WRAPPER_UNSET_IS_NULL(
+        "msg.single_int64_wrapper",
+        ImmutableMap.of("msg", TestAllTypes.getDefaultInstance()),
+        NullValue.NULL_VALUE),
     PROTO3_MAP_OF_ANY_UNPACKS_VALUE(
         "msg.map_string_any['k']",
         ImmutableMap.of(
@@ -847,6 +861,11 @@ private enum NativeSelectEvaluationTestCase {
         true),
     HAS_FIELD_PROTO3_OPTIONAL_SCALAR_UNSET(
         "has(msg.optional_bool)", ImmutableMap.of("msg", TestAllTypes.getDefaultInstance()), false),
+    HAS_FIELD_PROTO3_WRAPPER_SET_TO_ZERO(
+        "has(msg.single_int64_wrapper)",
+        ImmutableMap.of(
+            "msg", TestAllTypes.newBuilder().setSingleInt64Wrapper(Int64Value.of(0)).build()),
+        true),
     HAS_FIELD_PROTO2_SCALAR_SET_TO_DEFAULT(
         "has(proto2_msg.single_int32)",
         ImmutableMap.of(
@@ -1068,9 +1087,7 @@ 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"),
-    WRAPPER(
-        "msg.single_int64_wrapper", "Optimization of wrapper fields is currently unimplemented");
+    ANY("msg.single_any", "Optimization of Any fields is currently unimplemented");
 
     private final String expression;
     private final String expectedMessageSubstring;
@@ -1331,6 +1348,15 @@ public void optimize_toParsedExpr_matchesExpectedSerializedProto() throws Except
         .isEqualTo("cel.@attribute(msg, [[14, \"single_string\", 9]], \"\") + \"suffix\"");
   }
 
+  @Test
+  public void optimize_wrapperFieldSelect_preservesNullableResultType() throws Exception {
+    CelAbstractSyntaxTree ast = cel.compile("msg.single_int64_wrapper").getAst();
+
+    CelAbstractSyntaxTree optimizedAst = celOptimizer.optimize(ast);
+
+    assertThat(optimizedAst.getResultType()).isEqualTo(NullableType.create(SimpleType.INT));
+  }
+
   @Test
   public void optimize_resultFunctionDeclarations_containsOnlySingularAttributeAndHasField()
       throws Exception {

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

@@ -49,8 +49,7 @@ final class OptimizedSelectPlanner {
 
   /**
    * Well-known message types whose CEL semantics (Any unpacking, JSON value conversion) are not
-   * implemented by the optimized traversal. Wrapper types are rejected via {@link
-   * CelTypes#isWrapperType}.
+   * implemented by the optimized traversal.
    */
   private static final ImmutableSet<String> UNSUPPORTED_WELL_KNOWN_TYPE_IDENTS =
       ImmutableSet.of(
@@ -167,8 +166,8 @@ private static SelectField parseMapLeaf(
     CelMap.Entry entry = dummyExpr.map().entries().get(0);
     validateScalarDummy(mapEntrySpec.keyTypeCode(), entry.key());
     if (mapEntrySpec.valueTypeCode() == SelectField.MESSAGE_TYPE_CODE) {
-      // Unlike leaf fields, map values of Any, Struct, Value, and wrapper types are permitted:
-      // the traversal only produces the map, and indexing it applies standard CEL conversion.
+      // Unlike leaf fields, map values of unsupported well-known types are permitted: the
+      // traversal only produces the map, and indexing it applies standard CEL conversion.
       return SelectField.createMap(
           fieldNumber, fieldName, mapEntrySpec, parseMessageProtoTypeName(entry.value()));
     }
@@ -301,8 +300,7 @@ private static String parseMessageProtoTypeName(CelExpr expr) {
   private static String parseLeafMessageProtoTypeName(CelExpr expr) {
     String protoTypeName = parseMessageProtoTypeName(expr);
     checkArgument(
-        !CelTypes.isWrapperType(protoTypeName)
-            && !UNSUPPORTED_WELL_KNOWN_TYPE_IDENTS.contains(protoTypeName),
+        !UNSUPPORTED_WELL_KNOWN_TYPE_IDENTS.contains(protoTypeName),
         "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

@@ -26,6 +26,7 @@
 import com.google.protobuf.ByteString;
 import com.google.protobuf.CodedOutputStream;
 import com.google.protobuf.ExtensionRegistryLite;
+import com.google.protobuf.Int64Value;
 import com.google.protobuf.UnknownFieldSet;
 import com.google.testing.junit.testparameterinjector.TestParameter;
 import com.google.testing.junit.testparameterinjector.TestParameterInjector;
@@ -2191,6 +2192,15 @@ private enum DescriptorlessEvaluationTestCase {
             + " && msg.map_string_duration['d'] == duration('5m')",
         POPULATED_SERVER_MESSAGE,
         true),
+    WRAPPERS(
+        "msg.single_int64_wrapper == 0 && has(msg.single_int64_wrapper)"
+            + " && msg.single_string_wrapper == null && msg.repeated_int64_wrapper == [1, 0]",
+        TestAllTypes.newBuilder()
+            .setSingleInt64Wrapper(Int64Value.of(0))
+            .addRepeatedInt64Wrapper(Int64Value.of(1))
+            .addRepeatedInt64Wrapper(Int64Value.of(0))
+            .build(),
+        true),
     ;
 
     private final String expression;

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

@@ -551,21 +551,20 @@ public void plan_invalidAst_unsupportedWellKnownType_throwsEvaluationException()
                             CelExpr.ofList(
                                 4L,
                                 ImmutableList.of(
-                                    CelExpr.ofConstant(5L, CelConstant.ofValue(105L)),
-                                    CelExpr.ofConstant(
-                                        6L, CelConstant.ofValue("single_int64_wrapper")),
+                                    CelExpr.ofConstant(5L, CelConstant.ofValue(100L)),
+                                    CelExpr.ofConstant(6L, CelConstant.ofValue("single_any")),
                                     CelExpr.ofConstant(7L, CelConstant.ofValue(11L))),
                                 ImmutableList.of())),
                         ImmutableList.of()),
-                    CelExpr.ofStruct(8L, "google.protobuf.Int64Value", ImmutableList.of()))),
+                    CelExpr.ofStruct(8L, "google.protobuf.Any", ImmutableList.of()))),
             CelSource.newBuilder().build());
 
     CelEvaluationException e = assertThrows(CelEvaluationException.class, () -> PLANNER.plan(ast));
 
     assertThat(e).hasCauseThat().isInstanceOf(IllegalArgumentException.class);
     assertThat(e)
         .hasMessageThat()
-        .contains("Leaf well-known type 'google.protobuf.Int64Value' is not supported");
+        .contains("Leaf well-known type 'google.protobuf.Any' is not supported");
   }
 
   @Test
@@ -839,7 +838,12 @@ public void plan_invalidAst_mapFieldWithEmptyMapDummy_throwsEvaluationException(
 
   @Test
   public void plan_invalidAst_repeatedMessageWithUnsupportedWellKnownType_throwsEvaluationException(
-      @TestParameter({"google.protobuf.Int64Value", "google.protobuf.Any"})
+      @TestParameter({
+            "google.protobuf.Any",
+            "google.protobuf.Struct",
+            "google.protobuf.Value",
+            "google.protobuf.ListValue"
+          })
           String unsupportedProtoTypeName) {
     CelAbstractSyntaxTree ast =
         CelAbstractSyntaxTree.newParsedAst(