From fa6b4b0977f92699062546c44ca47a70b91cae5c Mon Sep 17 00:00:00 2001 From: Sean Huh Date: Thu, 8 Oct 2026 16:37:08 -0700 Subject: [PATCH] Support wrappers in select optimized ASTs PiperOrigin-RevId: 996146653 --- .../optimizer/optimizers/SelectOptimizer.java | 12 +++---- .../dev/cel/optimizer/optimizers/BUILD.bazel | 1 + .../optimizers/SelectOptimizerTest.java | 32 +++++++++++++++++-- .../planner/OptimizedSelectPlanner.java | 10 +++--- .../CelLiteRuntimeVersionSkewTest.java | 10 ++++++ .../planner/OptimizedSelectPlannerTest.java | 16 ++++++---- 6 files changed, 58 insertions(+), 23 deletions(-) diff --git a/optimizer/src/main/java/dev/cel/optimizer/optimizers/SelectOptimizer.java b/optimizer/src/main/java/dev/cel/optimizer/optimizers/SelectOptimizer.java index 3efac648b..1e067f03c 100644 --- a/optimizer/src/main/java/dev/cel/optimizer/optimizers/SelectOptimizer.java +++ b/optimizer/src/main/java/dev/cel/optimizer/optimizers/SelectOptimizer.java @@ -99,8 +99,9 @@ *

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. * *

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 getOptimizableField( diff --git a/optimizer/src/test/java/dev/cel/optimizer/optimizers/BUILD.bazel b/optimizer/src/test/java/dev/cel/optimizer/optimizers/BUILD.bazel index f7c2ef401..3f67c49c6 100644 --- a/optimizer/src/test/java/dev/cel/optimizer/optimizers/BUILD.bazel +++ b/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", diff --git a/optimizer/src/test/java/dev/cel/optimizer/optimizers/SelectOptimizerTest.java b/optimizer/src/test/java/dev/cel/optimizer/optimizers/SelectOptimizerTest.java index fd1ad8787..2b3808a30 100644 --- a/optimizer/src/test/java/dev/cel/optimizer/optimizers/SelectOptimizerTest.java +++ b/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 { diff --git a/runtime/src/main/java/dev/cel/runtime/planner/OptimizedSelectPlanner.java b/runtime/src/main/java/dev/cel/runtime/planner/OptimizedSelectPlanner.java index a1687652f..2691c171b 100644 --- a/runtime/src/main/java/dev/cel/runtime/planner/OptimizedSelectPlanner.java +++ b/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 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; diff --git a/runtime/src/test/java/dev/cel/runtime/CelLiteRuntimeVersionSkewTest.java b/runtime/src/test/java/dev/cel/runtime/CelLiteRuntimeVersionSkewTest.java index c2d4631cf..176703cac 100644 --- a/runtime/src/test/java/dev/cel/runtime/CelLiteRuntimeVersionSkewTest.java +++ b/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; diff --git a/runtime/src/test/java/dev/cel/runtime/planner/OptimizedSelectPlannerTest.java b/runtime/src/test/java/dev/cel/runtime/planner/OptimizedSelectPlannerTest.java index 2a1b647b9..8428091e2 100644 --- a/runtime/src/test/java/dev/cel/runtime/planner/OptimizedSelectPlannerTest.java +++ b/runtime/src/test/java/dev/cel/runtime/planner/OptimizedSelectPlannerTest.java @@ -551,13 +551,12 @@ 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)); @@ -565,7 +564,7 @@ public void plan_invalidAst_unsupportedWellKnownType_throwsEvaluationException() 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(