Skip to content

Commit 7298cee

Browse files
l46kokcopybara-github
authored andcommitted
Optimize field access in lite messages by skipping unread fields
PiperOrigin-RevId: 993996287
1 parent dab6df3 commit 7298cee

12 files changed

Lines changed: 1760 additions & 409 deletions

‎common/src/main/java/dev/cel/common/values/BUILD.bazel‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -280,7 +280,9 @@ java_library(
280280
],
281281
deps = [
282282
":base_proto_cel_value_converter",
283+
":optimized_selectable",
283284
":preadapted_list",
285+
":select_field",
284286
":values",
285287
"//:auto_value",
286288
"//common:options",
@@ -320,6 +322,7 @@ java_library(
320322
"ProtoLiteCelValueConverter.java",
321323
"ProtoMessageLiteValue.java",
322324
"RawProtoMessageLiteValue.java",
325+
"WireMessageLite.java",
323326
],
324327
tags = [
325328
],
@@ -350,6 +353,7 @@ cel_android_library(
350353
"ProtoLiteCelValueConverter.java",
351354
"ProtoMessageLiteValue.java",
352355
"RawProtoMessageLiteValue.java",
356+
"WireMessageLite.java",
353357
],
354358
tags = [
355359
],

‎common/src/main/java/dev/cel/common/values/ProtoLiteCelValueConverter.java‎

Lines changed: 294 additions & 159 deletions
Large diffs are not rendered by default.

‎common/src/main/java/dev/cel/common/values/ProtoMessageLiteValue.java‎

Lines changed: 112 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -21,12 +21,14 @@
2121
import com.google.common.collect.ImmutableListMultimap;
2222
import com.google.common.collect.ImmutableMap;
2323
import com.google.errorprone.annotations.Immutable;
24+
import com.google.protobuf.ByteString;
2425
import com.google.protobuf.MessageLite;
2526
import dev.cel.common.types.CelType;
2627
import dev.cel.common.types.StructTypeReference;
2728
import dev.cel.common.values.ProtoLiteCelValueConverter.MessageFields;
2829
import dev.cel.protobuf.CelLiteDescriptor.FieldLiteDescriptor;
2930
import java.io.IOException;
31+
import java.util.Objects;
3032
import java.util.Optional;
3133
import org.jspecify.annotations.Nullable;
3234

@@ -38,41 +40,72 @@
3840
* <p>If the codebase has access to full protobuf messages with descriptors, use {@code
3941
* ProtoMessageValue} instead.
4042
*
43+
* <p>An instance is backed by either a materialized {@link #rawValue()} or unparsed {@link
44+
* #wireBytes()} (exactly one is non-null). Wire-backed instances decode selected fields directly
45+
* from {@link #wireBytes()} and lazily parse the full {@link MessageLite} only if {@link #value()}
46+
* is invoked.
47+
*
4148
* <p>Implements {@link OptimizedSelectable} so that select chains can address fields by number:
4249
*
4350
* <ul>
4451
* <li><b>Field renames:</b> If a protobuf field is renamed in schema after an AST was compiled,
4552
* resolving by {@link SelectField#fieldNumber()} maps the number to the runtime descriptor's
4653
* current field name, preventing {@code CelAttributeNotFoundException}.
4754
* <li><b>Version skew / unknown fields:</b> When evaluating payloads serialized by a newer binary
48-
* containing fields absent from the local {@code CelLiteDescriptor}, the unknown wire bytes
49-
* are preserved in {@link #unknownFields()} and decoded on demand using the compile-time wire
50-
* type and default metadata in {@link SelectField}.
55+
* containing fields absent from the local {@code CelLiteDescriptor}, the unknown fields are
56+
* decoded on demand directly from the message's wire bytes using the compile-time wire type
57+
* and default metadata in {@link SelectField}.
5158
* </ul>
5259
*/
5360
@AutoValue
5461
@Immutable
55-
public abstract class ProtoMessageLiteValue extends StructValue<String, MessageLite>
62+
abstract class ProtoMessageLiteValue extends StructValue<String, MessageLite>
5663
implements OptimizedSelectable {
5764

58-
@Override
59-
public abstract MessageLite value();
65+
// Populated when wrapping an already-materialized root MessageLite (e.g., from activation).
66+
abstract @Nullable MessageLite rawValue();
67+
68+
// Populated when slicing a nested submessage from parent wire bytes to avoid deserializing and
69+
// re-serializing intermediate hops; lazily parsed into a MessageLite only if value() is called.
70+
abstract @Nullable ByteString wireBytes();
6071

6172
@Override
6273
public abstract CelType celType();
6374

6475
abstract ProtoLiteCelValueConverter protoLiteCelValueConverter();
6576

77+
@Memoized
78+
@Override
79+
public MessageLite value() {
80+
MessageLite msg = rawValue();
81+
if (msg != null) {
82+
return msg;
83+
}
84+
return protoLiteCelValueConverter()
85+
.parseMessageLite(checkNotNull(wireBytes()), celType().name());
86+
}
87+
88+
@Memoized
89+
ByteString serializedRawValue() {
90+
return checkNotNull(rawValue()).toByteString();
91+
}
92+
93+
ByteString toByteString() {
94+
ByteString bytes = wireBytes();
95+
return bytes != null ? bytes : serializedRawValue();
96+
}
97+
6698
@Memoized
6799
MessageFields messageFields() {
68100
try {
69-
return protoLiteCelValueConverter().readMessageFields(value(), celType().name());
101+
return protoLiteCelValueConverter().readAllFields(toByteString(), celType().name());
70102
} catch (IOException e) {
71-
throw new IllegalStateException("Unable to read message fields for " + celType().name(), e);
103+
throw new IllegalArgumentException(
104+
"Failed to decode proto message of type: " + celType().name(), e);
72105
}
73106
}
74107

75-
ImmutableMap<String, Object> fieldValues() {
108+
private ImmutableMap<String, Object> fieldValues() {
76109
return messageFields().values();
77110
}
78111

@@ -82,9 +115,30 @@ ImmutableListMultimap<Integer, Object> unknownFields() {
82115

83116
@Override
84117
public boolean isZeroValue() {
118+
ByteString bytes = wireBytes();
119+
if (bytes != null && bytes.isEmpty()) {
120+
return true;
121+
}
85122
return value().getDefaultInstanceForType().equals(value());
86123
}
87124

125+
@Override
126+
public final boolean equals(Object other) {
127+
if (other == this) {
128+
return true;
129+
}
130+
if (!(other instanceof ProtoMessageLiteValue)) {
131+
return false;
132+
}
133+
ProtoMessageLiteValue that = (ProtoMessageLiteValue) other;
134+
return this.celType().equals(that.celType()) && this.value().equals(that.value());
135+
}
136+
137+
@Override
138+
public final int hashCode() {
139+
return Objects.hash(value(), celType());
140+
}
141+
88142
@Override
89143
public Object select(String field) {
90144
return find(field)
@@ -93,16 +147,15 @@ public Object select(String field) {
93147

94148
@Override
95149
public Optional<Object> find(String field) {
96-
Object fieldValue = fieldValues().get(field);
97-
return Optional.ofNullable(fieldValue)
98-
.map(value -> protoLiteCelValueConverter().toRuntimeValue(fieldValue));
150+
return Optional.ofNullable(fieldValues().get(field))
151+
.map(protoLiteCelValueConverter()::toRuntimeValue);
99152
}
100153

101154
@Override
102155
public Object selectByFieldNumber(SelectField field) {
103156
FieldLiteDescriptor fd = findFieldDescriptor(field);
104157
if (fd != null) {
105-
Object known = fieldValues().get(fd.getFieldName());
158+
Object known = readField(fd);
106159
if (known != null) {
107160
return protoLiteCelValueConverter().toRuntimeValue(known);
108161
}
@@ -112,28 +165,48 @@ public Object selectByFieldNumber(SelectField field) {
112165
return protoLiteCelValueConverter().getDefaultCelValue(fd);
113166
}
114167
return RawProtoMessageLiteValue.selectWireOrDefault(
115-
field, unknownFields().get(field.fieldNumber()), protoLiteCelValueConverter());
168+
field,
169+
RawProtoMessageLiteValue.readWireField(toByteString(), field.fieldNumber()),
170+
protoLiteCelValueConverter());
116171
}
117172

118173
@Override
119174
public boolean hasFieldByNumber(SelectField field) {
120175
FieldLiteDescriptor fd = findFieldDescriptor(field);
121176
if (fd != null) {
122-
return fieldValues().containsKey(fd.getFieldName());
177+
return hasField(fd);
123178
}
124-
return RawProtoMessageLiteValue.isPresentInWire(
125-
field, unknownFields().get(field.fieldNumber()));
179+
return RawProtoMessageLiteValue.isPresentInWire(toByteString(), field);
126180
}
127181

128182
@Override
129183
public Optional<Object> findByFieldNumber(SelectField field) {
130184
FieldLiteDescriptor fd = findFieldDescriptor(field);
131185
if (fd != null) {
132-
return Optional.ofNullable(fieldValues().get(fd.getFieldName()))
133-
.map(value -> protoLiteCelValueConverter().toRuntimeValue(value));
186+
return Optional.ofNullable(readField(fd)).map(protoLiteCelValueConverter()::toRuntimeValue);
134187
}
135188
return RawProtoMessageLiteValue.navigateWire(
136-
field, unknownFields().get(field.fieldNumber()), protoLiteCelValueConverter());
189+
field,
190+
RawProtoMessageLiteValue.readWireField(toByteString(), field.fieldNumber()),
191+
protoLiteCelValueConverter());
192+
}
193+
194+
private @Nullable Object readField(FieldLiteDescriptor fd) {
195+
try {
196+
return protoLiteCelValueConverter().readSingleField(toByteString(), fd);
197+
} catch (IOException e) {
198+
throw new IllegalArgumentException(
199+
"Failed to decode proto message of type: " + celType().name(), e);
200+
}
201+
}
202+
203+
private boolean hasField(FieldLiteDescriptor fd) {
204+
try {
205+
return protoLiteCelValueConverter().hasSingleField(toByteString(), fd);
206+
} catch (IOException e) {
207+
throw new IllegalArgumentException(
208+
"Failed to decode proto message of type: " + celType().name(), e);
209+
}
137210
}
138211

139212
private @Nullable FieldLiteDescriptor findFieldDescriptor(SelectField field) {
@@ -142,13 +215,30 @@ public Optional<Object> findByFieldNumber(SelectField field) {
142215
.orElse(null);
143216
}
144217

145-
public static ProtoMessageLiteValue create(
218+
static ProtoMessageLiteValue create(
146219
MessageLite value, String typeName, ProtoLiteCelValueConverter protoLiteCelValueConverter) {
147220
checkNotNull(value);
148221
checkNotNull(typeName);
149222
checkNotNull(protoLiteCelValueConverter);
150223
return new AutoValue_ProtoMessageLiteValue(
151-
value, StructTypeReference.create(typeName), protoLiteCelValueConverter);
224+
value,
225+
/* wireBytes= */ null,
226+
StructTypeReference.create(typeName),
227+
protoLiteCelValueConverter);
228+
}
229+
230+
static ProtoMessageLiteValue create(
231+
ByteString wireBytes,
232+
String typeName,
233+
ProtoLiteCelValueConverter protoLiteCelValueConverter) {
234+
checkNotNull(wireBytes);
235+
checkNotNull(typeName);
236+
checkNotNull(protoLiteCelValueConverter);
237+
return new AutoValue_ProtoMessageLiteValue(
238+
/* rawValue= */ null,
239+
wireBytes,
240+
StructTypeReference.create(typeName),
241+
protoLiteCelValueConverter);
152242
}
153243

154244
ProtoMessageLiteValue() {}

‎common/src/main/java/dev/cel/common/values/ProtoMessageValue.java‎

Lines changed: 60 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -14,8 +14,9 @@
1414

1515
package dev.cel.common.values;
1616

17+
import static com.google.common.base.Preconditions.checkNotNull;
18+
1719
import com.google.auto.value.AutoValue;
18-
import com.google.common.base.Preconditions;
1920
import com.google.errorprone.annotations.Immutable;
2021
import com.google.protobuf.Descriptors.Descriptor;
2122
import com.google.protobuf.Descriptors.FieldDescriptor;
@@ -28,7 +29,8 @@
2829
/** ProtoMessageValue is a struct value with protobuf support. */
2930
@AutoValue
3031
@Immutable
31-
public abstract class ProtoMessageValue extends StructValue<String, Message> {
32+
public abstract class ProtoMessageValue extends StructValue<String, Message>
33+
implements OptimizedSelectable {
3234

3335
@Override
3436
public abstract Message value();
@@ -60,28 +62,38 @@ public Optional<Object> find(String field) {
6062
FieldDescriptor fieldDescriptor =
6163
findField(celDescriptorPool(), value().getDescriptorForType(), field);
6264

63-
// Selecting a field on a protobuf message yields a default value even if the field is not
64-
// declared. Therefore, we must exhaustively test whether they are actually declared.
65-
if (fieldDescriptor.isRepeated()) {
66-
if (value().getRepeatedFieldCount(fieldDescriptor) == 0) {
67-
return Optional.empty();
68-
}
69-
} else if (!value().hasField(fieldDescriptor)) {
70-
return Optional.empty();
71-
}
65+
return findFieldValue(fieldDescriptor);
66+
}
7267

73-
return Optional.of(
74-
protoCelValueConverter().fromProtoMessageFieldToCelValue(value(), fieldDescriptor));
68+
@Override
69+
public Object selectByFieldNumber(SelectField field) {
70+
FieldDescriptor fieldDescriptor = findFieldByNumber(value().getDescriptorForType(), field);
71+
72+
return protoCelValueConverter().fromProtoMessageFieldToCelValue(value(), fieldDescriptor);
73+
}
74+
75+
@Override
76+
public boolean hasFieldByNumber(SelectField field) {
77+
FieldDescriptor fieldDescriptor = findFieldByNumber(value().getDescriptorForType(), field);
78+
79+
return isFieldPresent(fieldDescriptor);
80+
}
81+
82+
@Override
83+
public Optional<Object> findByFieldNumber(SelectField field) {
84+
FieldDescriptor fieldDescriptor = findFieldByNumber(value().getDescriptorForType(), field);
85+
86+
return findFieldValue(fieldDescriptor);
7587
}
7688

7789
public static ProtoMessageValue create(
7890
Message value,
7991
CelDescriptorPool celDescriptorPool,
8092
ProtoCelValueConverter protoCelValueConverter,
8193
boolean enableJsonFieldNames) {
82-
Preconditions.checkNotNull(value);
83-
Preconditions.checkNotNull(celDescriptorPool);
84-
Preconditions.checkNotNull(protoCelValueConverter);
94+
checkNotNull(value);
95+
checkNotNull(celDescriptorPool);
96+
checkNotNull(protoCelValueConverter);
8597
return new AutoValue_ProtoMessageValue(
8698
value,
8799
StructTypeReference.create(value.getDescriptorForType().getFullName()),
@@ -90,6 +102,36 @@ public static ProtoMessageValue create(
90102
enableJsonFieldNames);
91103
}
92104

105+
private Optional<Object> findFieldValue(FieldDescriptor fieldDescriptor) {
106+
if (!isFieldPresent(fieldDescriptor)) {
107+
return Optional.empty();
108+
}
109+
110+
return Optional.of(
111+
protoCelValueConverter().fromProtoMessageFieldToCelValue(value(), fieldDescriptor));
112+
}
113+
114+
private boolean isFieldPresent(FieldDescriptor fieldDescriptor) {
115+
// Selecting a field on a protobuf message yields a default value even if the field is not
116+
// declared. Therefore, we must exhaustively test whether they are actually declared.
117+
if (fieldDescriptor.isRepeated()) {
118+
return value().getRepeatedFieldCount(fieldDescriptor) > 0;
119+
}
120+
return value().hasField(fieldDescriptor);
121+
}
122+
123+
private static FieldDescriptor findFieldByNumber(Descriptor descriptor, SelectField field) {
124+
FieldDescriptor fieldDescriptor = descriptor.findFieldByNumber(field.fieldNumber());
125+
if (fieldDescriptor != null) {
126+
return fieldDescriptor;
127+
}
128+
129+
throw new IllegalArgumentException(
130+
String.format(
131+
"field '%s' (number %d) is not declared in message '%s'",
132+
field.fieldName(), field.fieldNumber(), descriptor.getFullName()));
133+
}
134+
93135
private FieldDescriptor findField(
94136
CelDescriptorPool celDescriptorPool, Descriptor descriptor, String fieldName) {
95137
if (enableJsonFieldNames()) {
@@ -114,4 +156,6 @@ private FieldDescriptor findField(
114156
"field '%s' is not declared in message '%s'",
115157
fieldName, descriptor.getFullName())));
116158
}
159+
160+
ProtoMessageValue() {}
117161
}

0 commit comments

Comments
 (0)