Compile policy tooling with field selection optimization applied

PiperOrigin-RevId: 975271045

Author: l46kok

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

@@ -317,6 +317,7 @@ java_library(
     srcs = [
         "ProtoLiteCelValueConverter.java",
         "ProtoMessageLiteValue.java",
+        "RawProtoMessageLiteValue.java",
     ],
     tags = [
     ],
@@ -325,6 +326,7 @@ java_library(
         ":values",
         "//:auto_value",
         "//common/annotations",
+        "//common/exceptions:attribute_not_found",
         "//common/internal:cel_lite_descriptor_pool",
         "//common/internal:well_known_proto",
         "//common/types",
@@ -333,6 +335,7 @@ java_library(
         "//protobuf:cel_lite_descriptor",
         "@maven//:com_google_errorprone_error_prone_annotations",
         "@maven//:com_google_guava_guava",
+        "@maven//:org_jspecify_jspecify",
         "@maven_android//:com_google_protobuf_protobuf_javalite",
     ],
 )
@@ -342,6 +345,7 @@ cel_android_library(
     srcs = [
         "ProtoLiteCelValueConverter.java",
         "ProtoMessageLiteValue.java",
+        "RawProtoMessageLiteValue.java",
     ],
     tags = [
     ],
@@ -350,6 +354,7 @@ cel_android_library(
         ":values_android",
         "//:auto_value",
         "//common/annotations",
+        "//common/exceptions:attribute_not_found",
         "//common/internal:cel_lite_descriptor_pool_android",
         "//common/internal:well_known_proto_android",
         "//common/types:type_providers_android",
@@ -358,6 +363,7 @@ cel_android_library(
         "//protobuf:cel_lite_descriptor",
         "@maven//:com_google_errorprone_error_prone_annotations",
         "@maven//:com_google_guava_guava",
+        "@maven//:org_jspecify_jspecify",
         "@maven_android//:com_google_guava_guava",
         "@maven_android//:com_google_protobuf_protobuf_javalite",
     ],

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

@@ -20,6 +20,7 @@
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Defaults;
 import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableListMultimap;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.Multimap;
 import com.google.common.collect.Multimaps;
@@ -45,6 +46,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.NoSuchElementException;
+import java.util.Optional;
 import java.util.TreeMap;
 
 /**
@@ -160,6 +162,19 @@ Object getDefaultCelValue(String protoTypeName, String fieldName) {
     return toRuntimeValue(defaultValue);
   }
 
+  @Internal
+  public Optional<FieldLiteDescriptor> findFieldDescriptor(String protoTypeName, String fieldName) {
+    return descriptorPool
+        .findDescriptor(protoTypeName)
+        .flatMap(desc -> desc.findByFieldName(fieldName));
+  }
+
+  @Internal
+  public Optional<Object> findDefaultCelValue(String protoTypeName, String fieldName) {
+    return findFieldDescriptor(protoTypeName, fieldName)
+        .map(fieldDescriptor -> toRuntimeValue(getDefaultValue(fieldDescriptor)));
+  }
+
   @Override
   @SuppressWarnings("LiteProtoToString") // No alternative identifier to use. Debug only info is OK.
   public Object toRuntimeValue(Object value) {
@@ -193,7 +208,10 @@ protected Object fromWellKnownProto(MessageLiteOrBuilder msg, WellKnownProto wel
           descriptorPool
               .findDescriptor(message)
               .orElseThrow(
-                  () -> new NoSuchElementException("Could not find a descriptor for: " + message));
+                  () ->
+                      new NoSuchElementException(
+                          "Could not find a descriptor for message of type: "
+                              + message.getClass().getName()));
       return ProtoMessageLiteValue.create(message, descriptor.getProtoTypeName(), this);
     }
 
@@ -369,11 +387,14 @@ MessageFields readAllFields(byte[] bytes, String protoTypeName) throws IOExcepti
 
   ImmutableMap<String, Object> readAllFields(MessageLite msg, String protoTypeName)
       throws IOException {
-    return readAllFields(msg.toByteArray(), protoTypeName).values();
+    return readMessageFields(msg, protoTypeName).values();
   }
 
-  private static Object readUnknownField(int tagWireType, CodedInputStream inputStream)
-      throws IOException {
+  MessageFields readMessageFields(MessageLite msg, String protoTypeName) throws IOException {
+    return readAllFields(msg.toByteArray(), protoTypeName);
+  }
+
+  static Object readUnknownField(int tagWireType, CodedInputStream inputStream) throws IOException {
     switch (tagWireType) {
       case WireFormat.WIRETYPE_VARINT:
         return inputStream.readInt64();
@@ -393,16 +414,19 @@ private static Object readUnknownField(int tagWireType, CodedInputStream inputSt
   }
 
   @AutoValue
-  @SuppressWarnings("AutoValueImmutableFields") // Unknowns are inaccessible to users.
+  @AutoValue.CopyAnnotations
+  @Immutable
+  @SuppressWarnings("Immutable") // Safe immutable fields
   abstract static class MessageFields {
 
     abstract ImmutableMap<String, Object> values();
 
-    abstract Multimap<Integer, Object> unknowns();
+    abstract ImmutableListMultimap<Integer, Object> unknowns();
 
     static MessageFields create(
         ImmutableMap<String, Object> fieldValues, Multimap<Integer, Object> unknownFields) {
-      return new AutoValue_ProtoLiteCelValueConverter_MessageFields(fieldValues, unknownFields);
+      return new AutoValue_ProtoLiteCelValueConverter_MessageFields(
+          fieldValues, ImmutableListMultimap.copyOf(unknownFields));
     }
   }
 

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

@@ -17,11 +17,14 @@
 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 java.io.IOException;
 import java.util.Optional;
 
@@ -43,17 +46,27 @@ public abstract class ProtoMessageLiteValue extends StructValue<String, MessageL
   @Override
   public abstract CelType celType();
 
-  abstract ProtoLiteCelValueConverter protoLiteCelValueConverter();
+  @Internal
+  public abstract ProtoLiteCelValueConverter protoLiteCelValueConverter();
 
   @Memoized
-  ImmutableMap<String, Object> fieldValues() {
+  MessageFields messageFields() {
     try {
-      return protoLiteCelValueConverter().readAllFields(value(), celType().name());
+      return protoLiteCelValueConverter().readMessageFields(value(), celType().name());
     } catch (IOException e) {
       throw new IllegalStateException("Unable to read message fields for " + celType().name(), e);
     }
   }
 
+  @Internal
+  public ImmutableMap<String, Object> fieldValues() {
+    return messageFields().values();
+  }
+
+  public ImmutableListMultimap<Integer, Object> unknownFields() {
+    return messageFields().unknowns();
+  }
+
   @Override
   public boolean isZeroValue() {
     return value().getDefaultInstanceForType().equals(value());