Add support for managing modern type providers in the CEL type registry.

PiperOrigin-RevId: 535726361

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/interop.cc

@@ -526,6 +526,11 @@ 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));
+}
+
 Handle<NullValue> CreateNullValue() {
   return HandleFactory<NullValue>::Make<NullValue>();
 }

eval/internal/interop.h

@@ -68,6 +68,9 @@ struct MessageWrapperAccess final {
       google::api::expr::runtime::MessageWrapper& wrapper);
 };
 
+Handle<StructType> CreateStructTypeFromLegacyTypeInfo(
+    const google::api::expr::runtime::LegacyTypeInfoApis* type_info);
+
 // 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/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;

eval/public/cel_type_registry.h

@@ -1,6 +1,7 @@
 #ifndef THIRD_PARTY_CEL_CPP_EVAL_PUBLIC_CEL_TYPE_REGISTRY_H_
 #define THIRD_PARTY_CEL_CPP_EVAL_PUBLIC_CEL_TYPE_REGISTRY_H_
 
+#include <algorithm>
 #include <memory>
 #include <string>
 #include <utility>
@@ -16,6 +17,7 @@
 #include "base/types/enum_type.h"
 #include "base/value.h"
 #include "eval/public/structs/legacy_type_provider.h"
+#include "runtime/internal/composed_type_provider.h"
 
 namespace google::api::expr::runtime {
 
@@ -67,13 +69,31 @@ class CelTypeRegistry {
   // Register a new type provider.
   //
   // Type providers are consulted in the order they are added.
-  void RegisterTypeProvider(std::unique_ptr<LegacyTypeProvider> provider) {
-    type_providers_.push_back(std::move(provider));
-  }
+  void RegisterTypeProvider(std::unique_ptr<LegacyTypeProvider> provider);
 
   // Get the first registered type provider.
   std::shared_ptr<const LegacyTypeProvider> GetFirstTypeProvider() const;
 
+  // Returns the effective type provider that has been configured with the
+  // registry.
+  //
+  // This is a composited type provider that should check in order:
+  // - builtins (via TypeManager)
+  // - custom enumerations
+  // - registered extension type providers in the order registered.
+  const cel::TypeProvider& GetTypeProvider() const {
+    return type_provider_impl_;
+  }
+
+  // Register an additional type provider with the registry.
+  //
+  // A pointer to the registered provider is returned to support testing,
+  // but users should prefer to use the composed type provider from
+  // GetTypeProvider()
+  void RegisterModernTypeProvider(std::unique_ptr<cel::TypeProvider> provider) {
+    return type_provider_impl_.AddTypeProvider(std::move(provider));
+  }
+
   // Find a type adapter given a fully qualified type name.
   // Adapter provides a generic interface for the reflection operations the
   // interpreter needs to provide.
@@ -117,7 +137,11 @@ class CelTypeRegistry {
   // Internal representation for enums.
   absl::flat_hash_map<std::string, cel::Handle<cel::EnumType>>
       resolveable_enums_;
-  std::vector<std::shared_ptr<LegacyTypeProvider>> type_providers_;
+  cel::runtime_internal::ComposedTypeProvider type_provider_impl_;
+  // TODO(uncreated-issue/44): This is needed to inspect the registered legacy type
+  // providers for client tests. This can be removed when they are migrated to
+  // use the modern APIs.
+  std::vector<std::shared_ptr<const LegacyTypeProvider>> legacy_type_providers_;
 };
 
 }  // namespace google::api::expr::runtime