Update the default type provider (protobuf reflection impl) to return a specific

type info instance instead of the generic one.

PiperOrigin-RevId: 534552333

Author: jnthntatum

eval/public/structs/BUILD

@@ -259,6 +259,7 @@ cc_test(
         ":legacy_type_adapter",
         ":legacy_type_info_apis",
         ":proto_message_type_adapter",
+        "//eval/public:cel_options",
         "//eval/public:cel_value",
         "//eval/public:message_wrapper",
         "//eval/public/containers:container_backed_list_impl",

eval/public/structs/legacy_type_info_apis.h

@@ -31,6 +31,9 @@ class LegacyTypeMutationApis;
 // Provides ability to obtain field access apis, type info, and debug
 // representation of a message.
 //
+// The message parameter may wrap a nullptr to request generic accessors /
+// mutators for the TypeInfo instance if it is available.
+//
 // This is implemented as a separate class from LegacyTypeAccessApis to resolve
 // cyclic dependency between CelValue (which needs to access these apis to
 // provide DebugString and ObtainCelTypename) and LegacyTypeAccessApis (which

eval/public/structs/proto_message_type_adapter.cc

@@ -301,6 +301,34 @@ CelValue MessageCelValueFactory(const google::protobuf::Message* message) {
 
 }  // namespace
 
+std::string ProtoMessageTypeAdapter::DebugString(
+    const MessageWrapper& wrapped_message) const {
+  if (!wrapped_message.HasFullProto() ||
+      wrapped_message.message_ptr() == nullptr) {
+    return UnsupportedTypeName();
+  }
+  auto* message = cel::internal::down_cast<const google::protobuf::Message*>(
+      wrapped_message.message_ptr());
+  return message->ShortDebugString();
+}
+
+const std::string& ProtoMessageTypeAdapter::GetTypename(
+    const MessageWrapper& wrapped_message) const {
+  return descriptor_->full_name();
+}
+
+const LegacyTypeMutationApis* ProtoMessageTypeAdapter::GetMutationApis(
+    const MessageWrapper& wrapped_message) const {
+  // Defer checks for misuse on wrong message kind in the accessor calls.
+  return this;
+}
+
+const LegacyTypeAccessApis* ProtoMessageTypeAdapter::GetAccessApis(
+    const MessageWrapper& wrapped_message) const {
+  // Defer checks for misuse on wrong message kind in the builder calls.
+  return this;
+}
+
 absl::Status ProtoMessageTypeAdapter::ValidateSetFieldOp(
     bool assertion, absl::string_view field, absl::string_view detail) const {
   if (!assertion) {

eval/public/structs/proto_message_type_adapter.h

@@ -15,6 +15,7 @@
 #ifndef THIRD_PARTY_CEL_CPP_EVAL_PUBLIC_STRUCTS_PROTO_MESSAGE_TYPE_ADAPTER_H_
 #define THIRD_PARTY_CEL_CPP_EVAL_PUBLIC_STRUCTS_PROTO_MESSAGE_TYPE_ADAPTER_H_
 
+#include <string>
 #include <vector>
 
 #include "google/protobuf/descriptor.h"
@@ -29,7 +30,13 @@
 
 namespace google::api::expr::runtime {
 
-class ProtoMessageTypeAdapter : public LegacyTypeAccessApis,
+// Implementation for legacy struct (message) type apis using reflection.
+//
+// Note: The type info API implementation attached to message values is
+// generally the duck-typed instance to support the default behavior of
+// deferring to the protobuf reflection apis on the message instance.
+class ProtoMessageTypeAdapter : public LegacyTypeInfoApis,
+                                public LegacyTypeAccessApis,
                                 public LegacyTypeMutationApis {
  public:
   ProtoMessageTypeAdapter(const google::protobuf::Descriptor* descriptor,
@@ -38,6 +45,19 @@ class ProtoMessageTypeAdapter : public LegacyTypeAccessApis,
 
   ~ProtoMessageTypeAdapter() override = default;
 
+  // Implement LegacyTypeInfoApis
+  std::string DebugString(const MessageWrapper& wrapped_message) const override;
+
+  const std::string& GetTypename(
+      const MessageWrapper& wrapped_message) const override;
+
+  const LegacyTypeAccessApis* GetAccessApis(
+      const MessageWrapper& wrapped_message) const override;
+
+  const LegacyTypeMutationApis* GetMutationApis(
+      const MessageWrapper& wrapped_message) const override;
+
+  // Implement LegacyTypeMutation APIs.
   absl::StatusOr<CelValue::MessageWrapper::Builder> NewInstance(
       cel::MemoryManager& memory_manager) const override;
 
@@ -52,6 +72,7 @@ class ProtoMessageTypeAdapter : public LegacyTypeAccessApis,
       cel::MemoryManager& memory_manager,
       CelValue::MessageWrapper::Builder instance) const override;
 
+  // Implement LegacyTypeAccessAPIs.
   absl::StatusOr<CelValue> GetField(
       absl::string_view field_name, const CelValue::MessageWrapper& instance,
       ProtoWrapperTypeOptions unboxing_option,

eval/public/structs/proto_message_type_adapter_test.cc

@@ -20,6 +20,7 @@
 #include "google/protobuf/message.h"
 #include "google/protobuf/message_lite.h"
 #include "absl/status/status.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"
@@ -671,5 +672,70 @@ TEST(ProtoMessageTypeAdapter, AdaptFromWellKnownTypeNotAMessageError) {
               StatusIs(absl::StatusCode::kInternal));
 }
 
+TEST(ProtoMesssageTypeAdapter, TypeInfoDebug) {
+  google::protobuf::Arena arena;
+  ProtoMessageTypeAdapter adapter(
+      google::protobuf::DescriptorPool::generated_pool()->FindMessageTypeByName(
+          "google.api.expr.runtime.TestMessage"),
+      google::protobuf::MessageFactory::generated_factory());
+  ProtoMemoryManager manager(&arena);
+
+  TestMessage message;
+  message.set_int64_value(42);
+  EXPECT_THAT(adapter.DebugString(MessageWrapper(&message, &adapter)),
+              HasSubstr(message.ShortDebugString()));
+
+  EXPECT_THAT(adapter.DebugString(MessageWrapper()),
+              HasSubstr("<unknown message>"));
+}
+
+TEST(ProtoMesssageTypeAdapter, TypeInfoName) {
+  google::protobuf::Arena arena;
+  ProtoMessageTypeAdapter adapter(
+      google::protobuf::DescriptorPool::generated_pool()->FindMessageTypeByName(
+          "google.api.expr.runtime.TestMessage"),
+      google::protobuf::MessageFactory::generated_factory());
+  ProtoMemoryManager manager(&arena);
+
+  EXPECT_EQ(adapter.GetTypename(MessageWrapper()),
+            "google.api.expr.runtime.TestMessage");
+}
+
+TEST(ProtoMesssageTypeAdapter, TypeInfoMutator) {
+  google::protobuf::Arena arena;
+  ProtoMessageTypeAdapter adapter(
+      google::protobuf::DescriptorPool::generated_pool()->FindMessageTypeByName(
+          "google.api.expr.runtime.TestMessage"),
+      google::protobuf::MessageFactory::generated_factory());
+  ProtoMemoryManager manager(&arena);
+
+  const LegacyTypeMutationApis* api = adapter.GetMutationApis(MessageWrapper());
+  ASSERT_NE(api, nullptr);
+
+  ASSERT_OK_AND_ASSIGN(MessageWrapper::Builder builder,
+                       api->NewInstance(manager));
+  EXPECT_NE(dynamic_cast<TestMessage*>(builder.message_ptr()), nullptr);
+}
+
+TEST(ProtoMesssageTypeAdapter, TypeInfoAccesor) {
+  google::protobuf::Arena arena;
+  ProtoMessageTypeAdapter adapter(
+      google::protobuf::DescriptorPool::generated_pool()->FindMessageTypeByName(
+          "google.api.expr.runtime.TestMessage"),
+      google::protobuf::MessageFactory::generated_factory());
+  ProtoMemoryManager manager(&arena);
+
+  TestMessage message;
+  message.set_int64_value(42);
+  CelValue::MessageWrapper wrapped(&message, &adapter);
+
+  const LegacyTypeAccessApis* api = adapter.GetAccessApis(MessageWrapper());
+  ASSERT_NE(api, nullptr);
+
+  EXPECT_THAT(api->GetField("int64_value", wrapped,
+                            ProtoWrapperTypeOptions::kUnsetNull, manager),
+              IsOkAndHolds(test::IsCelInt64(42)));
+}
+
 }  // namespace
 }  // namespace google::api::expr::runtime

eval/public/structs/protobuf_descriptor_type_provider.cc

@@ -25,18 +25,7 @@ namespace google::api::expr::runtime {
 
 absl::optional<LegacyTypeAdapter> ProtobufDescriptorProvider::ProvideLegacyType(
     absl::string_view name) const {
-  const ProtoMessageTypeAdapter* result = nullptr;
-  {
-    absl::MutexLock lock(&mu_);
-    auto it = type_cache_.find(name);
-    if (it != type_cache_.end()) {
-      result = it->second.get();
-    } else {
-      auto type_provider = GetType(name);
-      result = type_provider.get();
-      type_cache_[name] = std::move(type_provider);
-    }
-  }
+  const ProtoMessageTypeAdapter* result = GetTypeAdapter(name);
   if (result == nullptr) {
     return absl::nullopt;
   }
@@ -47,11 +36,11 @@ absl::optional<LegacyTypeAdapter> ProtobufDescriptorProvider::ProvideLegacyType(
 absl::optional<const LegacyTypeInfoApis*>
 ProtobufDescriptorProvider::ProvideLegacyTypeInfo(
     absl::string_view name) const {
-  return &GetGenericProtoTypeInfoInstance();
+  return GetTypeAdapter(name);
 }
 
-std::unique_ptr<ProtoMessageTypeAdapter> ProtobufDescriptorProvider::GetType(
-    absl::string_view name) const {
+std::unique_ptr<ProtoMessageTypeAdapter>
+ProtobufDescriptorProvider::CreateTypeAdapter(absl::string_view name) const {
   const google::protobuf::Descriptor* descriptor =
       descriptor_pool_->FindMessageTypeByName(name);
   if (descriptor == nullptr) {
@@ -61,4 +50,17 @@ std::unique_ptr<ProtoMessageTypeAdapter> ProtobufDescriptorProvider::GetType(
   return std::make_unique<ProtoMessageTypeAdapter>(descriptor,
                                                    message_factory_);
 }
+
+const ProtoMessageTypeAdapter* ProtobufDescriptorProvider::GetTypeAdapter(
+    absl::string_view name) const {
+  absl::MutexLock lock(&mu_);
+  auto it = type_cache_.find(name);
+  if (it != type_cache_.end()) {
+    return it->second.get();
+  }
+  auto type_provider = CreateTypeAdapter(name);
+  const ProtoMessageTypeAdapter* result = type_provider.get();
+  type_cache_[name] = std::move(type_provider);
+  return result;
+}
 }  // namespace google::api::expr::runtime

eval/public/structs/protobuf_descriptor_type_provider.h

@@ -45,11 +45,13 @@ class ProtobufDescriptorProvider : public LegacyTypeProvider {
       absl::string_view name) const override;
 
  private:
-  // Run a lookup if the type adapter hasn't already been built.
-  // returns nullptr if not found.
-  std::unique_ptr<ProtoMessageTypeAdapter> GetType(
+  // Create a new type instance if found in the registered descriptor pool.
+  // Otherwise, returns nullptr.
+  std::unique_ptr<ProtoMessageTypeAdapter> CreateTypeAdapter(
       absl::string_view name) const;
 
+  const ProtoMessageTypeAdapter* GetTypeAdapter(absl::string_view name) const;
+
   const google::protobuf::DescriptorPool* descriptor_pool_;
   google::protobuf::MessageFactory* message_factory_;
   mutable absl::flat_hash_map<std::string,