google · Nov 30, 2023 · Nov 30, 2023
diff --git a/base/values/struct_value.cc b/base/values/struct_value.cc
@@ -110,7 +110,7 @@ absl::StatusOr<Handle<Value>> StructValue::GetFieldByNumber(
                                             number);
 }
 
-absl::StatusOr<Handle<Value>> StructValue::Qualify(
+absl::StatusOr<QualifyResult> StructValue::Qualify(
     ValueFactory& value_factory, absl::Span<const SelectQualifier> qualifiers,
     bool presence_test) const {
   return CEL_INTERNAL_STRUCT_VALUE_DISPATCH(Qualify, value_factory, qualifiers,
@@ -345,7 +345,7 @@ absl::StatusOr<Handle<Value>> LegacyStructValue::GetFieldByNumber(
                                       true);
 }
 
-absl::StatusOr<Handle<Value>> LegacyStructValue::Qualify(
+absl::StatusOr<QualifyResult> LegacyStructValue::Qualify(
     ValueFactory& value_factory, absl::Span<const SelectQualifier> qualifiers,
     bool presence_test) const {
   return MessageValueQualify(msg_, type_info_, value_factory, qualifiers,

diff --git a/base/values/struct_value.h b/base/values/struct_value.h
@@ -66,6 +66,14 @@ struct FieldSpecifier {
 
 using SelectQualifier = absl::variant<FieldSpecifier, AttributeQualifier>;
 
+struct QualifyResult {
+  // The possibly intermediate result of qualification.
+  Handle<Value> value;
+  // The number of qualifiers applied.
+  // A value < 0 indicates all qualifiers were effectively applied.
+  int qualifier_count;
+};
+
 // StructValue represents an instance of cel::StructType.
 class StructValue : public Value,
                     public base_internal::EnableHandleFromThis<StructValue> {
@@ -117,7 +125,7 @@ class StructValue : public Value,
   // absl::StatusCode::kUnimplemented has special meaning: if returned the
   // evaluator will attempt to apply the operation using the standard Get/Has
   // operations.
-  absl::StatusOr<Handle<Value>> Qualify(
+  absl::StatusOr<QualifyResult> Qualify(
       ValueFactory& value_factory,
       absl::Span<const SelectQualifier> select_qualifiers,
       bool presence_test) const;
@@ -225,7 +233,7 @@ ABSL_ATTRIBUTE_WEAK absl::StatusOr<bool> MessageValueHasFieldByName(
 ABSL_ATTRIBUTE_WEAK absl::StatusOr<Handle<Value>> MessageValueGetFieldByNumber(
     uintptr_t msg, uintptr_t type_info, ValueFactory& value_factory,
     int64_t number, bool unbox_null_wrapper_types);
-ABSL_ATTRIBUTE_WEAK absl::StatusOr<Handle<Value>> MessageValueQualify(
+ABSL_ATTRIBUTE_WEAK absl::StatusOr<QualifyResult> MessageValueQualify(
     uintptr_t msg, uintptr_t type_info, ValueFactory& value_factory,
     absl::Span<const SelectQualifier> qualifiers, bool presence_test);
 ABSL_ATTRIBUTE_WEAK absl::StatusOr<Handle<Value>> MessageValueGetFieldByName(
@@ -263,7 +271,7 @@ class LegacyStructValue final : public StructValue, public InlineData {
   absl::StatusOr<Handle<Value>> GetFieldByNumber(ValueFactory& value_factory,
                                                  int64_t number) const;
 
-  absl::StatusOr<Handle<Value>> Qualify(
+  absl::StatusOr<QualifyResult> Qualify(
       ValueFactory& value_factory,
       absl::Span<const SelectQualifier> select_qualifiers,
       bool presence_test) const;
@@ -358,7 +366,7 @@ class AbstractStructValue : public StructValue,
   virtual absl::StatusOr<Handle<Value>> GetFieldByNumber(
       ValueFactory& value_factory, int64_t number) const = 0;
 
-  virtual absl::StatusOr<Handle<Value>> Qualify(
+  virtual absl::StatusOr<QualifyResult> Qualify(
       ValueFactory& value_factory,
       absl::Span<const SelectQualifier> select_qualifiers,
       bool presence_test) const {

diff --git a/eval/internal/interop.cc b/eval/internal/interop.cc
@@ -333,7 +333,7 @@ absl::StatusOr<Handle<Value>> LegacyStructGetFieldImpl(
   return FromLegacyValue(arena, legacy_value);
 }
 
-absl::StatusOr<Handle<Value>> LegacyStructQualifyImpl(
+absl::StatusOr<QualifyResult> LegacyStructQualifyImpl(
     const MessageWrapper& wrapper, absl::Span<const cel::SelectQualifier> path,
     bool presence_test, MemoryManagerRef memory_manager) {
   if (path.empty()) {
@@ -353,15 +353,22 @@ absl::StatusOr<Handle<Value>> LegacyStructQualifyImpl(
               return field.GetStringKey().value_or("<invalid field>");
             }},
         path.front());
-    return interop_internal::CreateErrorValueFromView(
-        interop_internal::CreateNoSuchFieldError(arena, field_name));
+    return QualifyResult{
+        interop_internal::CreateErrorValueFromView(
+            interop_internal::CreateNoSuchFieldError(arena, field_name)),
+        -1};
   }
 
   CEL_ASSIGN_OR_RETURN(
-      auto legacy_value,
+      LegacyTypeAccessApis::LegacyQualifyResult legacy_result,
       access_api->Qualify(path, wrapper, presence_test, memory_manager));
 
-  return FromLegacyValue(arena, legacy_value);
+  QualifyResult result;
+  result.qualifier_count = legacy_result.qualifier_count;
+
+  CEL_ASSIGN_OR_RETURN(result.value,
+                       FromLegacyValue(arena, legacy_result.value));
+  return result;
 }
 
 }  // namespace
@@ -1009,7 +1016,7 @@ absl::StatusOr<Handle<Value>> MessageValueGetFieldByName(
       wrapper, name, unbox_null_wrapper_types, value_factory.memory_manager());
 }
 
-absl::StatusOr<Handle<Value>> MessageValueQualify(
+absl::StatusOr<QualifyResult> MessageValueQualify(
     uintptr_t msg, uintptr_t type_info, ValueFactory& value_factory,
     absl::Span<const SelectQualifier> qualifiers, bool presence_test) {
   auto wrapper = MessageWrapperAccess::Make(msg, type_info);

diff --git a/eval/public/structs/BUILD b/eval/public/structs/BUILD
@@ -277,6 +277,7 @@ cc_test(
         "//eval/public/testing:matchers",
         "//eval/testutil:test_message_cc_proto",
         "//extensions/protobuf:memory_manager",
+        "//internal:proto_matchers",
         "//internal:testing",
         "//runtime:runtime_options",
         "//testutil:util",

diff --git a/eval/public/structs/legacy_type_adapter.h b/eval/public/structs/legacy_type_adapter.h
@@ -77,6 +77,13 @@ class LegacyTypeMutationApis {
 // dynamic dispatch to concrete implementations).
 class LegacyTypeAccessApis {
  public:
+  struct LegacyQualifyResult {
+    // The possibly intermediate result of the select operation.
+    CelValue value;
+    // Number of qualifiers applied.
+    int qualifier_count;
+  };
+
   virtual ~LegacyTypeAccessApis() = default;
 
   // Return whether an instance of the type has field set to a non-default
@@ -107,7 +114,7 @@ class LegacyTypeAccessApis {
   // - presence_test controls whether to treat the call as a 'has' call,
   // returning
   //   whether the leaf field is set to a non-default value.
-  virtual absl::StatusOr<CelValue> Qualify(
+  virtual absl::StatusOr<LegacyQualifyResult> Qualify(
       absl::Span<const cel::SelectQualifier>,
       const CelValue::MessageWrapper& instance, bool presence_test,
       cel::MemoryManagerRef memory_manager) const {

diff --git a/eval/public/structs/proto_message_type_adapter.cc b/eval/public/structs/proto_message_type_adapter.cc
@@ -181,10 +181,11 @@ absl::StatusOr<CelValue> GetFieldImpl(const google::protobuf::Message* message,
   return CreateCelValueFromField(message, field_desc, unboxing_option, arena);
 }
 
-absl::StatusOr<CelValue> QualifyImpl(
+absl::StatusOr<LegacyTypeAccessApis::LegacyQualifyResult> QualifyImpl(
     const google::protobuf::Message* message, const google::protobuf::Descriptor* descriptor,
     absl::Span<const cel::SelectQualifier> path, bool presence_test,
     cel::MemoryManagerRef memory_manager) {
+  using LegacyQualifyResult = LegacyTypeAccessApis::LegacyQualifyResult;
   google::protobuf::Arena* arena = ProtoMemoryManagerArena(memory_manager);
   const google::protobuf::Message* select_path_elem = message;
   const google::protobuf::FieldDescriptor* field_desc = nullptr;
@@ -195,7 +196,7 @@ absl::StatusOr<CelValue> QualifyImpl(
     const auto& qualifier = path[i];
     if (!absl::holds_alternative<cel::FieldSpecifier>(qualifier)) {
       return absl::UnimplementedError(
-          "Select optimization only supports message traversal.");
+          "Select optimization only supports typed message field traversal.");
     }
     const auto& field_specifier = absl::get<cel::FieldSpecifier>(qualifier);
 
@@ -215,7 +216,8 @@ absl::StatusOr<CelValue> QualifyImpl(
     }
 
     if (field_desc == nullptr) {
-      return CreateNoSuchFieldError(memory_manager, field_specifier.name);
+      return LegacyQualifyResult{
+          CreateNoSuchFieldError(memory_manager, field_specifier.name), -1};
     }
 
     if (i == path.size() - 1) {
@@ -225,8 +227,14 @@ absl::StatusOr<CelValue> QualifyImpl(
     }
 
     if (field_desc->is_repeated()) {
-      return absl::UnimplementedError(
-          "Select optimization only supports singular field traversal.");
+      // Return early if we can't traverse the next step (a map or list).
+      LegacyQualifyResult result;
+      result.qualifier_count = i + 1;
+      CEL_ASSIGN_OR_RETURN(
+          result.value,
+          CreateCelValueFromField(select_path_elem, field_desc,
+                                  ProtoWrapperTypeOptions::kUnsetNull, arena));
+      return result;
     }
 
     switch (field_desc->type()) {
@@ -240,16 +248,24 @@ absl::StatusOr<CelValue> QualifyImpl(
     }
   }
 
+  LegacyQualifyResult result;
+  result.qualifier_count = -1;
+
   if (presence_test) {
-    return CelValue::CreateBool(
+    result.value = CelValue::CreateBool(
         CelFieldIsPresent(select_path_elem, field_desc, reflection));
+
+    return result;
   }
 
   // For simplicity, qualify only supports the spec behavior for wrapper types.
-  // The optimization is opt-in, so should not affect clients depening on the
+  // The optimization is opt-in, so should not affect clients depending on the
   // legacy behavior.
-  return CreateCelValueFromField(select_path_elem, field_desc,
-                                 ProtoWrapperTypeOptions::kUnsetNull, arena);
+  CEL_ASSIGN_OR_RETURN(
+      result.value,
+      CreateCelValueFromField(select_path_elem, field_desc,
+                              ProtoWrapperTypeOptions::kUnsetNull, arena));
+  return result;
 }
 
 std::vector<absl::string_view> ListFieldsImpl(
@@ -293,7 +309,7 @@ class DucktypedMessageAdapter : public LegacyTypeAccessApis,
                         unboxing_option, memory_manager);
   }
 
-  absl::StatusOr<CelValue> Qualify(
+  absl::StatusOr<LegacyTypeAccessApis::LegacyQualifyResult> Qualify(
       absl::Span<const cel::SelectQualifier> qualifiers,
       const CelValue::MessageWrapper& instance, bool presence_test,
       cel::MemoryManagerRef memory_manager) const override {
@@ -510,7 +526,8 @@ absl::StatusOr<CelValue> ProtoMessageTypeAdapter::GetField(
                       memory_manager);
 }
 
-absl::StatusOr<CelValue> ProtoMessageTypeAdapter::Qualify(
+absl::StatusOr<LegacyTypeAccessApis::LegacyQualifyResult>
+ProtoMessageTypeAdapter::Qualify(
     absl::Span<const cel::SelectQualifier> qualifiers,
     const CelValue::MessageWrapper& instance, bool presence_test,
     cel::MemoryManagerRef memory_manager) const {

diff --git a/eval/public/structs/proto_message_type_adapter.h b/eval/public/structs/proto_message_type_adapter.h
@@ -91,7 +91,7 @@ class ProtoMessageTypeAdapter : public LegacyTypeInfoApis,
       absl::string_view field_name,
       const CelValue::MessageWrapper& value) const override;
 
-  absl::StatusOr<CelValue> Qualify(
+  absl::StatusOr<LegacyTypeAccessApis::LegacyQualifyResult> Qualify(
       absl::Span<const cel::SelectQualifier> qualifiers,
       const CelValue::MessageWrapper& instance, bool presence_test,
       cel::MemoryManagerRef memory_manager) const override;

diff --git a/eval/public/structs/proto_message_type_adapter_test.cc b/eval/public/structs/proto_message_type_adapter_test.cc
@@ -21,7 +21,6 @@
 #include "absl/status/status.h"
 #include "base/attribute.h"
 #include "base/values/struct_value.h"
-#include "eval/public/cel_options.h"
 #include "eval/public/cel_value.h"
 #include "eval/public/containers/container_backed_list_impl.h"
 #include "eval/public/containers/container_backed_map_impl.h"
@@ -31,26 +30,32 @@
 #include "eval/public/testing/matchers.h"
 #include "eval/testutil/test_message.pb.h"
 #include "extensions/protobuf/memory_manager.h"
+#include "internal/proto_matchers.h"
 #include "internal/testing.h"
 #include "runtime/runtime_options.h"
-#include "testutil/util.h"
 #include "google/protobuf/arena.h"
 #include "google/protobuf/descriptor.h"
 #include "google/protobuf/message.h"
 
 namespace google::api::expr::runtime {
 namespace {
 
+using ::cel::ProtoWrapperTypeOptions;
 using ::cel::extensions::ProtoMemoryManagerRef;
+using ::cel::internal::test::EqualsProto;
 using ::google::protobuf::Int64Value;
 using testing::_;
+using testing::AllOf;
 using testing::ElementsAre;
+using testing::Eq;
+using testing::Field;
 using testing::HasSubstr;
 using testing::Optional;
 using testing::Truly;
 using cel::internal::IsOkAndHolds;
 using cel::internal::StatusIs;
-using testutil::EqualsProto;
+
+using LegacyQualifyResult = LegacyTypeAccessApis::LegacyQualifyResult;
 
 class ProtoMessageTypeAccessorTest : public testing::TestWithParam<bool> {
  public:
@@ -766,12 +771,13 @@ TEST(ProtoMesssageTypeAdapter, Qualify) {
   std::vector<cel::SelectQualifier> qualfiers{
       cel::FieldSpecifier{12, "message_value"},
       cel::FieldSpecifier{2, "int64_value"}};
-  EXPECT_THAT(api->Qualify(qualfiers, wrapped,
-                           /*presence_test=*/false, manager),
-              IsOkAndHolds(test::IsCelInt64(42)));
+  EXPECT_THAT(
+      api->Qualify(qualfiers, wrapped,
+                   /*presence_test=*/false, manager),
+      IsOkAndHolds(Field(&LegacyQualifyResult::value, test::IsCelInt64(42))));
 }
 
-TEST(ProtoMesssageTypeAdapter, QualifyMapsNotYetSupported) {
+TEST(ProtoMesssageTypeAdapter, QualifyMapsTraversalPartialSupport) {
   google::protobuf::Arena arena;
   ProtoMessageTypeAdapter adapter(
       google::protobuf::DescriptorPool::generated_pool()->FindMessageTypeByName(
@@ -787,12 +793,43 @@ TEST(ProtoMesssageTypeAdapter, QualifyMapsNotYetSupported) {
   ASSERT_NE(api, nullptr);
 
   std::vector<cel::SelectQualifier> qualfiers{
+      cel::FieldSpecifier{210, "string_message_map"},
       cel::AttributeQualifier::OfString("@key"),
       cel::FieldSpecifier{2, "int64_value"}};
 
+  EXPECT_THAT(
+      api->Qualify(qualfiers, wrapped,
+                   /*presence_test=*/false, manager),
+      IsOkAndHolds(
+          AllOf(Field(&LegacyQualifyResult::value, Truly([](const CelValue& v) {
+                  return v.IsMap() && v.MapOrDie()->size() == 1;
+                })),
+                Field(&LegacyQualifyResult::qualifier_count, Eq(1)))));
+}
+
+TEST(ProtoMesssageTypeAdapter, UntypedQualifiersNotYetSupported) {
+  google::protobuf::Arena arena;
+  ProtoMessageTypeAdapter adapter(
+      google::protobuf::DescriptorPool::generated_pool()->FindMessageTypeByName(
+          "google.api.expr.runtime.TestMessage"),
+      google::protobuf::MessageFactory::generated_factory());
+  auto manager = ProtoMemoryManagerRef(&arena);
+
+  TestMessage message;
+  (*message.mutable_string_message_map())["@key"].set_int64_value(42);
+  CelValue::MessageWrapper wrapped(&message, &adapter);
+
+  const LegacyTypeAccessApis* api = adapter.GetAccessApis(MessageWrapper());
+  ASSERT_NE(api, nullptr);
+
+  std::vector<cel::SelectQualifier> qualfiers{
+      cel::AttributeQualifier::OfString("string_message_map"),
+      cel::AttributeQualifier::OfString("@key"),
+      cel::AttributeQualifier::OfString("int64_value")};
+
   EXPECT_THAT(api->Qualify(qualfiers, wrapped,
                            /*presence_test=*/false, manager),
-              StatusIs(absl::StatusCode::kUnimplemented));
+              StatusIs(absl::StatusCode::kUnimplemented, _));
 }
 
 TEST(ProtoMesssageTypeAdapter, QualifyRepeatedLeaf) {
@@ -817,10 +854,12 @@ TEST(ProtoMesssageTypeAdapter, QualifyRepeatedLeaf) {
       cel::FieldSpecifier{102, "int64_list"},
   };
 
-  EXPECT_THAT(api->Qualify(qualfiers, wrapped,
-                           /*presence_test=*/false, manager),
-              IsOkAndHolds(test::IsCelList(
-                  ElementsAre(test::IsCelInt64(1), test::IsCelInt64(2)))));
+  EXPECT_THAT(
+      api->Qualify(qualfiers, wrapped,
+                   /*presence_test=*/false, manager),
+      IsOkAndHolds(Field(&LegacyQualifyResult::value,
+                         test::IsCelList(ElementsAre(test::IsCelInt64(1),
+                                                     test::IsCelInt64(2))))));
 }
 
 }  // namespace

diff --git a/extensions/BUILD b/extensions/BUILD
@@ -152,12 +152,15 @@ cc_library(
         "//internal:overloaded",
         "//internal:status_macros",
         "//runtime:runtime_options",
+        "//runtime/internal:errors",
         "@com_google_absl//absl/algorithm:container",
+        "@com_google_absl//absl/base:core_headers",
         "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings",
         "@com_google_absl//absl/types:optional",
+        "@com_google_absl//absl/types:span",
         "@com_google_absl//absl/types:variant",
     ],
 )