Implement OptimizedSelectable for ProtoMessageLiteValue

PiperOrigin-RevId: 988430966

Author: l46kok

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

@@ -14,19 +14,21 @@
 
 package dev.cel.common.values;
 
+import static com.google.common.base.Preconditions.checkNotNull;
+
 import com.google.auto.value.AutoValue;
 import com.google.auto.value.extension.memoized.Memoized;
-import com.google.common.base.Preconditions;
 import com.google.common.collect.ImmutableListMultimap;
 import com.google.common.collect.ImmutableMap;
 import com.google.errorprone.annotations.Immutable;
 import com.google.protobuf.MessageLite;
-import dev.cel.common.annotations.Internal;
 import dev.cel.common.types.CelType;
 import dev.cel.common.types.StructTypeReference;
 import dev.cel.common.values.ProtoLiteCelValueConverter.MessageFields;
+import dev.cel.protobuf.CelLiteDescriptor.FieldLiteDescriptor;
 import java.io.IOException;
 import java.util.Optional;
+import org.jspecify.annotations.Nullable;
 
 /**
  * ProtoMessageLiteValue is a struct value with protobuf support for {@link MessageLite}.
@@ -35,10 +37,23 @@
  *
  * <p>If the codebase has access to full protobuf messages with descriptors, use {@code
  * ProtoMessageValue} instead.
+ *
+ * <p>Implements {@link OptimizedSelectable} so that select chains can address fields by number:
+ *
+ * <ul>
+ *   <li><b>Field renames:</b> If a protobuf field is renamed in schema after an AST was compiled,
+ *       resolving by {@link SelectField#fieldNumber()} maps the number to the runtime descriptor's
+ *       current field name, preventing {@code CelAttributeNotFoundException}.
+ *   <li><b>Version skew / unknown fields:</b> When evaluating payloads serialized by a newer binary
+ *       containing fields absent from the local {@code CelLiteDescriptor}, the unknown wire bytes
+ *       are preserved in {@link #unknownFields()} and decoded on demand using the compile-time wire
+ *       type and default metadata in {@link SelectField}.
+ * </ul>
  */
 @AutoValue
 @Immutable
-public abstract class ProtoMessageLiteValue extends StructValue<String, MessageLite> {
+public abstract class ProtoMessageLiteValue extends StructValue<String, MessageLite>
+    implements OptimizedSelectable {
 
   @Override
   public abstract MessageLite value();
@@ -57,12 +72,11 @@ MessageFields messageFields() {
     }
   }
 
-  @Internal
-  public ImmutableMap<String, Object> fieldValues() {
+  ImmutableMap<String, Object> fieldValues() {
     return messageFields().values();
   }
 
-  public ImmutableListMultimap<Integer, Object> unknownFields() {
+  ImmutableListMultimap<Integer, Object> unknownFields() {
     return messageFields().unknowns();
   }
 
@@ -84,11 +98,59 @@ public Optional<Object> find(String field) {
         .map(value -> protoLiteCelValueConverter().toRuntimeValue(fieldValue));
   }
 
+  @Override
+  public Object selectByFieldNumber(SelectField field) {
+    FieldLiteDescriptor fd = findFieldDescriptor(field);
+    Object known = findKnownFieldValue(fd);
+    if (known != null) {
+      return protoLiteCelValueConverter().toRuntimeValue(known);
+    }
+    return RawProtoMessageLiteValue.selectWireOrDefault(
+        field, fd, unknownFields().get(field.fieldNumber()), protoLiteCelValueConverter());
+  }
+
+  @Override
+  public boolean hasFieldByNumber(SelectField field) {
+    FieldLiteDescriptor fd = findFieldDescriptor(field);
+    if (findKnownFieldValue(fd) != null) {
+      return true;
+    }
+    return RawProtoMessageLiteValue.isPresentInWire(
+        field, fd, unknownFields().get(field.fieldNumber()));
+  }
+
+  @Override
+  public Optional<Object> findByFieldNumber(SelectField field) {
+    FieldLiteDescriptor fd = findFieldDescriptor(field);
+    Object known = findKnownFieldValue(fd);
+    if (known != null) {
+      return Optional.of(protoLiteCelValueConverter().toRuntimeValue(known));
+    }
+    return RawProtoMessageLiteValue.navigateWire(
+        field, fd, unknownFields().get(field.fieldNumber()), protoLiteCelValueConverter());
+  }
+
+  private @Nullable FieldLiteDescriptor findFieldDescriptor(SelectField field) {
+    return protoLiteCelValueConverter()
+        .findFieldDescriptor(celType().name(), field.fieldNumber())
+        .orElse(null);
+  }
+
+  private @Nullable Object findKnownFieldValue(@Nullable FieldLiteDescriptor fieldDescriptor) {
+    if (fieldDescriptor == null) {
+      return null;
+    }
+    return fieldValues().get(fieldDescriptor.getFieldName());
+  }
+
   public static ProtoMessageLiteValue create(
       MessageLite value, String typeName, ProtoLiteCelValueConverter protoLiteCelValueConverter) {
-    Preconditions.checkNotNull(value);
-    Preconditions.checkNotNull(typeName);
+    checkNotNull(value);
+    checkNotNull(typeName);
+    checkNotNull(protoLiteCelValueConverter);
     return new AutoValue_ProtoMessageLiteValue(
         value, StructTypeReference.create(typeName), protoLiteCelValueConverter);
   }
+
+  ProtoMessageLiteValue() {}
 }