Walk select chains through in-memory proto messages without wrapping each one

PiperOrigin-RevId: 999700409

Author: l46kok

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

@@ -280,6 +280,7 @@ java_library(
     ],
     deps = [
         ":base_proto_cel_value_converter",
+        ":optimized_select_traversal",
         ":optimized_selectable",
         ":preadapted_list",
         ":select_field",
@@ -328,6 +329,7 @@ java_library(
     ],
     deps = [
         ":base_proto_cel_value_converter",
+        ":optimized_select_traversal",
         ":optimized_selectable",
         ":preadapted_list",
         ":select_field",
@@ -360,6 +362,7 @@ cel_android_library(
     ],
     deps = [
         ":base_proto_cel_value_converter_android",
+        ":optimized_select_traversal_android",
         ":optimized_selectable_android",
         ":preadapted_list_android",
         ":select_field_android",
@@ -483,6 +486,7 @@ java_library(
         ":select_field",
         "//common/annotations",
         "@maven//:com_google_errorprone_error_prone_annotations",
+        "@maven//:com_google_guava_guava",
     ],
 )
 
@@ -495,6 +499,7 @@ cel_android_library(
         ":select_field_android",
         "//common/annotations",
         "@maven//:com_google_errorprone_error_prone_annotations",
+        "@maven_android//:com_google_guava_guava",
     ],
 )
 

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

@@ -35,8 +35,17 @@ public final class OptimizedSelectTraversal {
    * Qualifies {@code target} through every field in {@code fields} and returns the terminal value.
    */
   public static Object qualify(Object target, ImmutableList<SelectField> fields) {
+    return qualify(target, fields, 0);
+  }
+
+  /** Qualifies {@code target} through the fields of {@code fields} from index {@code from} on. */
+  static Object qualify(Object target, ImmutableList<SelectField> fields, int from) {
     Object current = target;
-    for (int i = 0; i < fields.size(); i++) {
+    for (int i = from; i < fields.size(); i++) {
+      // With only the last field left, there's nothing to walk.
+      if (i < fields.size() - 1 && current instanceof OptimizedSelectable) {
+        return ((OptimizedSelectable) current).selectPath(fields, i);
+      }
       current = qualifyField(current, fields.get(i));
     }
     return current;

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

@@ -14,6 +14,7 @@
 
 package dev.cel.common.values;
 
+import com.google.common.collect.ImmutableList;
 import com.google.errorprone.annotations.Immutable;
 import dev.cel.common.annotations.Internal;
 import java.util.Optional;
@@ -23,7 +24,8 @@
  *
  * <p>Implementations resolve individual field selections against themselves by protobuf field
  * number. Walking the chain across multiple fields and heterogeneous values belongs to {@link
- * OptimizedSelectTraversal}.
+ * OptimizedSelectTraversal}, though an implementation may walk the part of it that it can read
+ * itself through {@link #selectPath}.
  *
  * <p>CEL Library Internals. Do Not Use.
  */
@@ -42,4 +44,11 @@ public interface OptimizedSelectable {
    * step of a presence test, or empty if absent.
    */
   Optional<Object> findByFieldNumber(SelectField field);
+
+  /**
+   * Selects {@code fields} from index {@code from} on, starting at this value, and returns the last
+   * one's value. Implementations may walk the nested values they can read without materializing
+   * each one, and hand the rest back to {@link OptimizedSelectTraversal}.
+   */
+  Object selectPath(ImmutableList<SelectField> fields, int from);
 }

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

@@ -506,6 +506,46 @@ boolean hasFieldByNumber(
         "Unexpected encoding type: " + fieldDescriptor.getEncodingType());
   }
 
+  /**
+   * Like {@link OptimizedSelectable#selectPath}: walks {@code message}'s in-memory submessages
+   * through their getters without wrapping them, then selects the rest from the last one reached.
+   */
+  Object selectPath(ProtoMessageLiteValue message, ImmutableList<SelectField> fields, int from) {
+    MessageLite current = message.rawValue();
+    MessageLiteDescriptor descriptor = message.descriptor();
+    if (current == null || !descriptor.hasGetters()) {
+      // Only in-memory messages read through generated getters can be walked.
+      return OptimizedSelectTraversal.qualify(
+          message.selectByFieldNumber(fields.get(from)), fields, from + 1);
+    }
+    int i = from;
+    // Walks submessages up to the one holding the last field, stopping at one that's unset or has
+    // no descriptor with getters, such as well-known types.
+    while (i < fields.size() - 1) {
+      Object submessage = descriptor.getFieldValue(current, fields.get(i).fieldNumber());
+      if (!(submessage instanceof MessageLite)) {
+        break;
+      }
+      Optional<MessageLiteDescriptor> submessageDescriptor =
+          descriptorPool
+              .findDescriptor((MessageLite) submessage)
+              .filter(MessageLiteDescriptor::hasGetters);
+      if (!submessageDescriptor.isPresent()) {
+        break;
+      }
+      current = (MessageLite) submessage;
+      descriptor = submessageDescriptor.get();
+      i++;
+    }
+    // Selects the next field from the last message reached, or from this one if none was.
+    ProtoMessageLiteValue reached = message;
+    if (i > from) {
+      reached = ProtoMessageLiteValue.create(current, descriptor, this);
+    }
+    return OptimizedSelectTraversal.qualify(
+        reached.selectByFieldNumber(fields.get(i)), fields, i + 1);
+  }
+
   /** Like {@link #hasFieldByNumber}, through the field's generated getter. */
   boolean hasByGetter(ProtoMessageLiteValue message, SelectField field) {
     return readGetter(message, field.fieldNumber()) != null;

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

@@ -18,6 +18,7 @@
 
 import com.google.auto.value.AutoValue;
 import com.google.auto.value.extension.memoized.Memoized;
+import com.google.common.collect.ImmutableList;
 import com.google.errorprone.annotations.Immutable;
 import com.google.protobuf.ByteString;
 import com.google.protobuf.MessageLite;
@@ -147,6 +148,11 @@ public Optional<Object> find(String field) {
     return Optional.ofNullable(descriptor().findByFieldName(field)).map(this::readField);
   }
 
+  @Override
+  public Object selectPath(ImmutableList<SelectField> fields, int from) {
+    return protoLiteCelValueConverter().selectPath(this, fields, from);
+  }
+
   @Override
   public Object selectByFieldNumber(SelectField field) {
     FieldLiteDescriptor fd = descriptor().findByFieldNumber(field.fieldNumber());

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

@@ -17,11 +17,14 @@
 import static com.google.common.base.Preconditions.checkNotNull;
 
 import com.google.auto.value.AutoValue;
+import com.google.common.collect.ImmutableList;
 import com.google.errorprone.annotations.Immutable;
 import com.google.protobuf.Descriptors.Descriptor;
 import com.google.protobuf.Descriptors.FieldDescriptor;
+import com.google.protobuf.DynamicMessage;
 import com.google.protobuf.Message;
 import dev.cel.common.internal.CelDescriptorPool;
+import dev.cel.common.internal.WellKnownProto;
 import dev.cel.common.types.CelType;
 import dev.cel.common.types.StructTypeReference;
 import java.util.Optional;
@@ -76,9 +79,7 @@ public Optional<Object> find(String field) {
 
   @Override
   public Object selectByFieldNumber(SelectField field) {
-    FieldDescriptor fieldDescriptor = findFieldByNumber(value().getDescriptorForType(), field);
-
-    return protoCelValueConverter().fromProtoMessageFieldToCelValue(value(), fieldDescriptor);
+    return selectField(value(), field);
   }
 
   @Override
@@ -95,6 +96,42 @@ public Optional<Object> findByFieldNumber(SelectField field) {
     return findFieldValue(fieldDescriptor);
   }
 
+  @Override
+  public Object selectPath(ImmutableList<SelectField> fields, int from) {
+    Message message = value();
+    if (message instanceof DynamicMessage) {
+      // Each hop adapts a dynamic message into its generated class, if it has one.
+      return OptimizedSelectTraversal.qualify(
+          selectByFieldNumber(fields.get(from)), fields, from + 1);
+    }
+    int i = from;
+    // Walks nested messages without wrapping them, up to the one holding the last field.
+    while (i < fields.size() - 1) {
+      FieldDescriptor fieldDescriptor =
+          findFieldByNumber(message.getDescriptorForType(), fields.get(i));
+      if (!isWalkable(fieldDescriptor)) {
+        break;
+      }
+      message = (Message) message.getField(fieldDescriptor);
+      i++;
+    }
+    return OptimizedSelectTraversal.qualify(selectField(message, fields.get(i)), fields, i + 1);
+  }
+
+  // A singular message field that isn't of a well-known type, which the converter would adapt.
+  private static boolean isWalkable(FieldDescriptor fieldDescriptor) {
+    return fieldDescriptor.getJavaType() == FieldDescriptor.JavaType.MESSAGE
+        && !fieldDescriptor.isRepeated()
+        && !WellKnownProto.getByTypeName(fieldDescriptor.getMessageType().getFullName())
+            .isPresent();
+  }
+
+  private Object selectField(Message message, SelectField field) {
+    FieldDescriptor fieldDescriptor = findFieldByNumber(message.getDescriptorForType(), field);
+
+    return protoCelValueConverter().fromProtoMessageFieldToCelValue(message, fieldDescriptor);
+  }
+
   public static ProtoMessageValue create(
       Message value,
       CelDescriptorPool celDescriptorPool,

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

@@ -17,6 +17,7 @@
 import static com.google.common.base.Preconditions.checkNotNull;
 
 import com.google.auto.value.AutoValue;
+import com.google.common.collect.ImmutableList;
 import com.google.errorprone.annotations.Immutable;
 import com.google.protobuf.ByteString;
 import com.google.protobuf.CodedInputStream;
@@ -136,6 +137,19 @@ public Optional<Object> findByFieldNumber(SelectField field) {
     return protoLiteCelValueConverter().findByFieldNumber(toByteString(), protoTypeName(), field);
   }
 
+  @Override
+  public Object selectPath(ImmutableList<SelectField> fields, int from) {
+    // Wire bytes have no submessages in memory to walk. Consecutive hops through wire bytes are
+    // selected here, rather than each going back through the traversal.
+    Object current = this;
+    int i = from;
+    while (i < fields.size() - 1 && current instanceof RawProtoMessageLiteValue) {
+      current = ((RawProtoMessageLiteValue) current).selectByFieldNumber(fields.get(i));
+      i++;
+    }
+    return OptimizedSelectTraversal.qualify(current, fields, i);
+  }
+
   static RawProtoMessageLiteValue create(
       ByteString rawWireBytes, ProtoLiteCelValueConverter protoLiteCelValueConverter) {
     return create(rawWireBytes, "", protoLiteCelValueConverter);

common/src/test/java/dev/cel/common/values/OptimizedSelectTraversalTest.java

@@ -237,6 +237,27 @@ public void hasField_intermediateUnsupportedTarget_throwsException() {
     assertThat(thrown).hasMessageThat().contains("child_key");
   }
 
+  @Test
+  public void qualify_valueWalkingRestOfChain_returnsItsResult() {
+    FakeOptimizedSelectable walking =
+        new FakeOptimizedSelectable(ImmutableMap.of()) {
+          @Override
+          public Object selectPath(ImmutableList<SelectField> fields, int from) {
+            return "walked from " + fields.get(from).fieldName();
+          }
+        };
+    FakeOptimizedSelectable root = new FakeOptimizedSelectable(ImmutableMap.of("outer", walking));
+    ImmutableList<SelectField> fields =
+        ImmutableList.of(
+            SelectField.create(1L, "outer", 11),
+            SelectField.create(2L, "middle", 11),
+            SelectField.create(3L, "leaf", 9, ""));
+
+    Object result = OptimizedSelectTraversal.qualify(root, fields);
+
+    assertThat(result).isEqualTo("walked from middle");
+  }
+
   @Test
   public void hasField_unsupportedTarget_throwsException() {
     ImmutableList<SelectField> fields = ImmutableList.of(SelectField.create(1L, "invalid_field"));
@@ -272,7 +293,7 @@ public void hasField_errorValue_returnsFalse() {
   }
 
   @SuppressWarnings("Immutable")
-  private static final class FakeOptimizedSelectable implements OptimizedSelectable {
+  private static class FakeOptimizedSelectable implements OptimizedSelectable {
     private final ImmutableMap<String, Object> values;
 
     @Override
@@ -297,6 +318,12 @@ public Optional<Object> findByFieldNumber(SelectField field) {
       return Optional.ofNullable(values.get(field.fieldName()));
     }
 
+    @Override
+    public Object selectPath(ImmutableList<SelectField> fields, int from) {
+      return OptimizedSelectTraversal.qualify(
+          selectByFieldNumber(fields.get(from)), fields, from + 1);
+    }
+
     FakeOptimizedSelectable(Map<String, Object> values) {
       this.values = ImmutableMap.copyOf(values);
     }