Add accessors for extracting legacy type from cel::LegacyStructType.

PiperOrigin-RevId: 534499645

Author: jnthntatum

base/type_manager.h

@@ -35,7 +35,7 @@ namespace cel {
 class TypeManager final {
  public:
   TypeManager(TypeFactory& type_factory ABSL_ATTRIBUTE_LIFETIME_BOUND,
-              TypeProvider& type_provider ABSL_ATTRIBUTE_LIFETIME_BOUND)
+              const TypeProvider& type_provider ABSL_ATTRIBUTE_LIFETIME_BOUND)
       : type_factory_(type_factory), type_provider_(type_provider) {}
 
   MemoryManager& memory_manager() const {

eval/internal/BUILD

@@ -58,7 +58,10 @@ cc_test(
         "//eval/public/containers:container_backed_list_impl",
         "//eval/public/containers:container_backed_map_impl",
         "//eval/public/structs:cel_proto_wrapper",
+        "//eval/public/structs:legacy_type_info_apis",
+        "//eval/public/structs:proto_message_type_adapter",
         "//eval/public/structs:trivial_legacy_type_info",
+        "//eval/testutil:test_message_cc_proto",
         "//extensions/protobuf:memory_manager",
         "//extensions/protobuf:type",
         "//extensions/protobuf:value",

eval/internal/interop.cc

@@ -248,6 +248,11 @@ Handle<StructType> LegacyStructTypeAccess::Create(uintptr_t message) {
       base_internal::LegacyStructType>(message);
 }
 
+uintptr_t LegacyStructTypeAccess::GetStorage(
+    const base_internal::LegacyStructType& type) {
+  return type.msg_;
+}
+
 Handle<StructValue> LegacyStructValueAccess::Create(
     const MessageWrapper& wrapper) {
   return Create(MessageWrapperAccess::Message(wrapper),
@@ -526,6 +531,27 @@ absl::StatusOr<CelValue> ToLegacyValue(google::protobuf::Arena* arena,
       ValueKindToString(value->kind()), " is not yet implemented"));
 }
 
+Handle<StructType> CreateStructTypeFromLegacyTypeInfo(
+    const LegacyTypeInfoApis* type_info) {
+  return LegacyStructTypeAccess::Create(reinterpret_cast<uintptr_t>(type_info));
+}
+
+const google::api::expr::runtime::LegacyTypeInfoApis* LegacyTypeInfoFromType(
+    const Handle<Type>& type) {
+  if (!type->Is<base_internal::LegacyStructType>()) {
+    return nullptr;
+  }
+  uintptr_t representation = LegacyStructTypeAccess::GetStorage(
+      type->As<base_internal::LegacyStructType>());
+  if ((representation & base_internal::kMessageWrapperTagMask) !=
+      base_internal::kMessageWrapperTagTypeInfoValue) {
+    return nullptr;
+  }
+
+  return reinterpret_cast<google::api::expr::runtime::LegacyTypeInfoApis*>(
+      representation & base_internal::kMessageWrapperPtrMask);
+}
+
 Handle<NullValue> CreateNullValue() {
   return HandleFactory<NullValue>::Make<NullValue>();
 }

eval/internal/interop.h

@@ -30,6 +30,7 @@
 #include "base/values/type_value.h"
 #include "eval/public/cel_value.h"
 #include "eval/public/message_wrapper.h"
+#include "eval/public/structs/legacy_type_info_apis.h"
 
 namespace cel::interop_internal {
 
@@ -45,6 +46,7 @@ struct CelMapAccess final {
 
 struct LegacyStructTypeAccess final {
   static Handle<StructType> Create(uintptr_t message);
+  static uintptr_t GetStorage(const base_internal::LegacyStructType& type);
 };
 
 struct LegacyStructValueAccess final {
@@ -68,6 +70,13 @@ struct MessageWrapperAccess final {
       google::api::expr::runtime::MessageWrapper& wrapper);
 };
 
+Handle<StructType> CreateStructTypeFromLegacyTypeInfo(
+    const google::api::expr::runtime::LegacyTypeInfoApis* type_info);
+
+// Nullptr is returned if this is not a legacy message type.
+const google::api::expr::runtime::LegacyTypeInfoApis* LegacyTypeInfoFromType(
+    const Handle<Type>& type);
+
 // Unlike ValueFactory::CreateStringValue, this does not copy input and instead
 // wraps it. It should only be used for interop with the legacy CelValue.
 Handle<StringValue> CreateStringValueFromView(absl::string_view value);

eval/internal/interop_test.cc

@@ -39,20 +39,26 @@
 #include "eval/public/containers/container_backed_map_impl.h"
 #include "eval/public/message_wrapper.h"
 #include "eval/public/structs/cel_proto_wrapper.h"
+#include "eval/public/structs/legacy_type_info_apis.h"
+#include "eval/public/structs/proto_message_type_adapter.h"
 #include "eval/public/structs/trivial_legacy_type_info.h"
 #include "eval/public/unknown_set.h"
+#include "eval/testutil/test_message.pb.h"
 #include "extensions/protobuf/memory_manager.h"
 #include "extensions/protobuf/type_provider.h"
 #include "extensions/protobuf/value.h"
 #include "internal/testing.h"
+#include "google/protobuf/message.h"
 
 namespace cel::interop_internal {
 namespace {
 
 using ::google::api::expr::runtime::CelProtoWrapper;
 using ::google::api::expr::runtime::CelValue;
 using ::google::api::expr::runtime::ContainerBackedListImpl;
+using ::google::api::expr::runtime::LegacyTypeInfoApis;
 using ::google::api::expr::runtime::MessageWrapper;
+using ::google::api::expr::runtime::TestMessage;
 using ::google::api::expr::runtime::UnknownSet;
 using testing::Eq;
 using testing::HasSubstr;
@@ -799,6 +805,42 @@ TEST(ValueInterop, StructFromLegacyMessageLite) {
             value_wrapper.legacy_type_info());
 }
 
+TEST(ValueInterop, StructTypeFromLegacyTypeInfo) {
+  google::protobuf::Arena arena;
+  extensions::ProtoMemoryManager memory_manager(&arena);
+  TypeFactory type_factory(memory_manager);
+  TypeManager type_manager(type_factory, TypeProvider::Builtin());
+  ValueFactory value_factory(type_manager);
+  google::protobuf::LinkMessageReflection<google::api::expr::runtime::TestMessage>();
+
+  google::api::expr::runtime::ProtoMessageTypeAdapter adapter(
+      TestMessage::descriptor(), google::protobuf::MessageFactory::generated_factory());
+
+  Handle<Type> type = CreateStructTypeFromLegacyTypeInfo(
+      static_cast<const LegacyTypeInfoApis*>(&adapter));
+
+  EXPECT_EQ(type->name(), "google.api.expr.runtime.TestMessage");
+}
+
+TEST(ValueInterop, StructTypeLegacyTypeInfoRoundTrip) {
+  google::protobuf::Arena arena;
+  extensions::ProtoMemoryManager memory_manager(&arena);
+  TypeFactory type_factory(memory_manager);
+  TypeManager type_manager(type_factory, TypeProvider::Builtin());
+  ValueFactory value_factory(type_manager);
+  google::protobuf::LinkMessageReflection<google::api::expr::runtime::TestMessage>();
+
+  google::api::expr::runtime::ProtoMessageTypeAdapter adapter(
+      TestMessage::descriptor(), google::protobuf::MessageFactory::generated_factory());
+
+  Handle<Type> type = CreateStructTypeFromLegacyTypeInfo(
+      static_cast<const LegacyTypeInfoApis*>(&adapter));
+
+  EXPECT_EQ(LegacyTypeInfoFromType(type),
+            static_cast<const LegacyTypeInfoApis*>(&adapter));
+  EXPECT_EQ(LegacyTypeInfoFromType(type_factory.GetBoolType()), nullptr);
+}
+
 TEST(ValueInterop, LegacyStructRoundtrip) {
   google::protobuf::Arena arena;
   extensions::ProtoMemoryManager memory_manager(&arena);

eval/public/BUILD

@@ -940,16 +940,19 @@ cc_library(
     srcs = ["cel_type_registry.cc"],
     hdrs = ["cel_type_registry.h"],
     deps = [
+        "//base:data",
         "//base:handle",
         "//base:memory",
-        "//base:type",
-        "//base:value",
         "//eval/internal:interop",
+        "//eval/public/structs:legacy_type_info_apis",
         "//eval/public/structs:legacy_type_provider",
+        "//internal:casts",
+        "//runtime/internal:composed_type_provider",
         "@com_google_absl//absl/base:core_headers",
         "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/container:flat_hash_set",
         "@com_google_absl//absl/container:node_hash_set",
+        "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings",
         "@com_google_absl//absl/synchronization",
         "@com_google_absl//absl/types:optional",
@@ -962,8 +965,8 @@ cc_test(
     srcs = ["cel_type_registry_test.cc"],
     deps = [
         ":cel_type_registry",
-        "//base:type",
-        "//base:value",
+        "//base:data",
+        "//base:memory",
         "//eval/public/structs:legacy_type_provider",
         "//eval/testutil:test_message_cc_proto",
         "//internal:testing",
@@ -974,6 +977,22 @@ cc_test(
     ],
 )
 
+cc_test(
+    name = "cel_type_registry_protobuf_reflection_test",
+    srcs = ["cel_type_registry_protobuf_reflection_test.cc"],
+    deps = [
+        ":cel_type_registry",
+        "//base:data",
+        "//base:handle",
+        "//base:memory",
+        "//eval/public/structs:protobuf_descriptor_type_provider",
+        "//eval/testutil:test_message_cc_proto",
+        "//internal:testing",
+        "@com_google_absl//absl/types:optional",
+        "@com_google_protobuf//:protobuf",
+    ],
+)
+
 cc_test(
     name = "builtin_func_test",
     size = "small",

eval/public/cel_type_registry.cc

@@ -7,15 +7,22 @@
 
 #include "absl/container/flat_hash_set.h"
 #include "absl/container/node_hash_set.h"
+#include "absl/status/statusor.h"
 #include "absl/strings/string_view.h"
 #include "absl/synchronization/mutex.h"
 #include "absl/types/optional.h"
 #include "base/handle.h"
 #include "base/memory.h"
+#include "base/type.h"
 #include "base/type_factory.h"
+#include "base/type_provider.h"
 #include "base/types/enum_type.h"
+#include "base/types/struct_type.h"
 #include "base/value.h"
 #include "eval/internal/interop.h"
+#include "eval/public/structs/legacy_type_info_apis.h"
+#include "eval/public/structs/legacy_type_provider.h"
+#include "internal/casts.h"
 #include "google/protobuf/descriptor.h"
 
 namespace google::api::expr::runtime {
@@ -24,6 +31,7 @@ namespace {
 
 using cel::Handle;
 using cel::MemoryManager;
+using cel::Type;
 using cel::TypeFactory;
 using cel::UniqueRef;
 using cel::Value;
@@ -54,6 +62,51 @@ cel::TypeFactory& GetDefaultTypeFactory() {
   return *factory;
 }
 
+class LegacyToModernTypeProviderAdapter : public cel::TypeProvider {
+ public:
+  explicit LegacyToModernTypeProviderAdapter(const LegacyTypeProvider& provider)
+      : provider_(provider) {}
+
+  absl::StatusOr<absl::optional<Handle<Type>>> ProvideType(
+      TypeFactory& factory, absl::string_view name) const override {
+    absl::optional<const LegacyTypeInfoApis*> type_info =
+        provider_.ProvideLegacyTypeInfo(name);
+
+    if (!type_info.has_value() || *type_info == nullptr) {
+      return absl::nullopt;
+    }
+
+    return cel::interop_internal::CreateStructTypeFromLegacyTypeInfo(
+        *type_info);
+  }
+
+ private:
+  const LegacyTypeProvider& provider_;
+};
+
+// A trivial type provider for registered Enums.
+//
+// Clients manually register the expected enums available to reference in the
+// input expressions. These may mask other defined enumerators (e.g. reference
+// by a protobuf message).
+class EnumTypeProvider : public cel::TypeProvider {
+ public:
+  explicit EnumTypeProvider(const EnumMap& enum_map) : enum_map_(enum_map) {}
+
+  absl::StatusOr<absl::optional<Handle<Type>>> ProvideType(
+      TypeFactory&, absl::string_view name) const override {
+    auto iter = enum_map_.find(name);
+    if (iter != enum_map_.end()) {
+      return iter->second;
+    }
+
+    return absl::nullopt;
+  }
+
+ private:
+  const EnumMap& enum_map_;
+};
+
 // EnumType implementation for generic enums that are defined at runtime that
 // can be resolved in expressions.
 //
@@ -176,6 +229,8 @@ ResolveableEnumType::FindConstantByNumber(int64_t number) const {
 
 CelTypeRegistry::CelTypeRegistry() : types_(GetCoreTypes()) {
   RegisterEnum("google.protobuf.NullValue", {{"NULL_VALUE", 0}});
+  type_provider_impl_.AddTypeProvider(
+      std::make_unique<EnumTypeProvider>(resolveable_enums_));
 }
 
 void CelTypeRegistry::Register(std::string fully_qualified_type_name) {
@@ -198,18 +253,27 @@ void CelTypeRegistry::RegisterEnum(absl::string_view enum_name,
   resolveable_enums_[enum_name] = std::move(result_or).value();
 }
 
+void CelTypeRegistry::RegisterTypeProvider(
+    std::unique_ptr<LegacyTypeProvider> provider) {
+  legacy_type_providers_.push_back(
+      std::shared_ptr<const LegacyTypeProvider>(std::move(provider)));
+  type_provider_impl_.AddTypeProvider(
+      std::make_unique<LegacyToModernTypeProviderAdapter>(
+          *legacy_type_providers_.back()));
+}
+
 std::shared_ptr<const LegacyTypeProvider>
 CelTypeRegistry::GetFirstTypeProvider() const {
-  if (type_providers_.empty()) {
+  if (legacy_type_providers_.empty()) {
     return nullptr;
   }
-  return type_providers_[0];
+  return legacy_type_providers_[0];
 }
 
 // Find a type's CelValue instance by its fully qualified name.
 absl::optional<LegacyTypeAdapter> CelTypeRegistry::FindTypeAdapter(
     absl::string_view fully_qualified_type_name) const {
-  for (const auto& provider : type_providers_) {
+  for (const auto& provider : legacy_type_providers_) {
     auto maybe_adapter = provider->ProvideLegacyType(fully_qualified_type_name);
     if (maybe_adapter.has_value()) {
       return maybe_adapter;