Mandate passing proto2::DescriptorPool to the type checker

PiperOrigin-RevId: 689498942

Author: jcking

checker/BUILD

@@ -91,14 +91,18 @@ cc_library(
         "//common:decl",
         "//common:type",
         "//internal:status_macros",
+        "//internal:well_known_types",
         "//parser:macro",
         "@com_google_absl//absl/base:no_destructor",
+        "@com_google_absl//absl/base:nullability",
         "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/container:flat_hash_set",
         "@com_google_absl//absl/functional:any_invocable",
+        "@com_google_absl//absl/log:absl_check",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings",
+        "@com_google_protobuf//:protobuf",
     ],
 )
 
@@ -112,6 +116,7 @@ cc_test(
         "//common:decl",
         "//common:type",
         "//internal:testing",
+        "//internal:testing_descriptor_pool",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:status_matchers",
     ],
@@ -150,6 +155,7 @@ cc_test(
         "//common:decl",
         "//common:type",
         "//internal:testing",
+        "//internal:testing_descriptor_pool",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:status_matchers",
         "@com_google_protobuf//:protobuf",
@@ -185,11 +191,9 @@ cc_test(
         "//base/ast_internal:ast_impl",
         "//base/ast_internal:expr",
         "//checker/internal:test_ast_helpers",
-        "//extensions/protobuf:value",
         "//internal:testing",
+        "//internal:testing_descriptor_pool",
         "@com_google_absl//absl/status:status_matchers",
         "@com_google_absl//absl/strings",
-        "@com_google_cel_spec//proto/test/v1/proto3:test_all_types_cc_proto",
-        "@com_google_protobuf//:protobuf",
     ],
 )

checker/internal/BUILD

@@ -167,6 +167,7 @@ cc_test(
         "//extensions/protobuf:value",
         "//internal:status_macros",
         "//internal:testing",
+        "//internal:testing_descriptor_pool",
         "@com_google_absl//absl/base:no_destructor",
         "@com_google_absl//absl/base:nullability",
         "@com_google_absl//absl/container:flat_hash_set",

checker/internal/type_check_env.cc

@@ -30,6 +30,7 @@
 #include "common/type_introspector.h"
 #include "internal/status_macros.h"
 #include "google/protobuf/arena.h"
+#include "google/protobuf/descriptor.h"
 
 namespace cel::checker_internal {
 
@@ -59,8 +60,21 @@ absl::Nullable<const FunctionDecl*> TypeCheckEnv::LookupFunction(
 
 absl::StatusOr<absl::optional<Type>> TypeCheckEnv::LookupTypeName(
     TypeFactory& type_factory, absl::string_view name) const {
+  {
+    // Check the descriptor pool first, then fallback to custom type providers.
+    absl::Nullable<const google::protobuf::Descriptor*> descriptor =
+        descriptor_pool_->FindMessageTypeByName(name);
+    if (descriptor != nullptr) {
+      return Type::Message(descriptor);
+    }
+    absl::Nullable<const google::protobuf::EnumDescriptor*> enum_descriptor =
+        descriptor_pool_->FindEnumTypeByName(name);
+    if (enum_descriptor != nullptr) {
+      return Type::Enum(enum_descriptor);
+    }
+  }
   const TypeCheckEnv* scope = this;
-  while (scope != nullptr) {
+  do {
     for (auto iter = type_providers_.rbegin(); iter != type_providers_.rend();
          ++iter) {
       auto type = (*iter)->FindType(type_factory, name);
@@ -69,15 +83,34 @@ absl::StatusOr<absl::optional<Type>> TypeCheckEnv::LookupTypeName(
       }
     }
     scope = scope->parent_;
-  }
+  } while ((scope != nullptr));
   return absl::nullopt;
 }
 
 absl::StatusOr<absl::optional<VariableDecl>> TypeCheckEnv::LookupEnumConstant(
     TypeFactory& type_factory, absl::string_view type,
     absl::string_view value) const {
+  {
+    // Check the descriptor pool first, then fallback to custom type providers.
+    absl::Nullable<const google::protobuf::EnumDescriptor*> enum_descriptor =
+        descriptor_pool_->FindEnumTypeByName(type);
+    if (enum_descriptor != nullptr) {
+      absl::Nullable<const google::protobuf::EnumValueDescriptor*> enum_value_descriptor =
+          enum_descriptor->FindValueByName(value);
+      if (enum_value_descriptor == nullptr) {
+        return absl::nullopt;
+      }
+      auto decl =
+          MakeVariableDecl(absl::StrCat(enum_descriptor->full_name(), ".",
+                                        enum_value_descriptor->name()),
+                           Type::Enum(enum_descriptor));
+      decl.set_value(
+          Constant(static_cast<int64_t>(enum_value_descriptor->number())));
+      return decl;
+    }
+  }
   const TypeCheckEnv* scope = this;
-  while (scope != nullptr) {
+  do {
     for (auto iter = type_providers_.rbegin(); iter != type_providers_.rend();
          ++iter) {
       auto enum_constant = (*iter)->FindEnumConstant(type_factory, type, value);
@@ -95,7 +128,7 @@ absl::StatusOr<absl::optional<VariableDecl>> TypeCheckEnv::LookupEnumConstant(
       }
     }
     scope = scope->parent_;
-  }
+  } while (scope != nullptr);
   return absl::nullopt;
 }
 
@@ -122,8 +155,25 @@ absl::StatusOr<absl::optional<VariableDecl>> TypeCheckEnv::LookupTypeConstant(
 absl::StatusOr<absl::optional<StructTypeField>> TypeCheckEnv::LookupStructField(
     TypeFactory& type_factory, absl::string_view type_name,
     absl::string_view field_name) const {
+  {
+    // Check the descriptor pool first, then fallback to custom type providers.
+    absl::Nullable<const google::protobuf::Descriptor*> descriptor =
+        descriptor_pool_->FindMessageTypeByName(type_name);
+    if (descriptor != nullptr) {
+      absl::Nullable<const google::protobuf::FieldDescriptor*> field_descriptor =
+          descriptor->FindFieldByName(field_name);
+      if (field_descriptor == nullptr) {
+        field_descriptor = descriptor_pool_->FindExtensionByPrintableName(
+            descriptor, field_name);
+        if (field_descriptor == nullptr) {
+          return absl::nullopt;
+        }
+      }
+      return cel::MessageTypeField(field_descriptor);
+    }
+  }
   const TypeCheckEnv* scope = this;
-  while (scope != nullptr) {
+  do {
     // Check the type providers in reverse registration order.
     // Note: this doesn't allow for shadowing a type with a subset type of the
     // same name -- the parent type provider will still be considered when
@@ -137,7 +187,7 @@ absl::StatusOr<absl::optional<StructTypeField>> TypeCheckEnv::LookupStructField(
       }
     }
     scope = scope->parent_;
-  }
+  } while (scope != nullptr);
   return absl::nullopt;
 }
 

checker/internal/type_check_env.h

@@ -34,6 +34,7 @@
 #include "common/type_factory.h"
 #include "common/type_introspector.h"
 #include "google/protobuf/arena.h"
+#include "google/protobuf/descriptor.h"
 
 namespace cel::checker_internal {
 
@@ -89,11 +90,12 @@ class TypeCheckEnv {
   using FunctionDeclPtr = absl::Nonnull<const FunctionDecl*>;
 
  public:
-  TypeCheckEnv() : container_(""), parent_(nullptr) {};
-
-  explicit TypeCheckEnv(const TypeCheckEnv* parent)
-      : container_(parent != nullptr ? parent->container() : ""),
-        parent_(parent) {}
+  explicit TypeCheckEnv(
+      absl::Nonnull<std::shared_ptr<const google::protobuf::DescriptorPool>>
+          descriptor_pool)
+      : descriptor_pool_(std::move(descriptor_pool)),
+        container_(""),
+        parent_(nullptr) {};
 
   // Move-only.
   TypeCheckEnv(TypeCheckEnv&&) = default;
@@ -165,14 +167,28 @@ class TypeCheckEnv {
       TypeFactory& type_factory, absl::Nonnull<google::protobuf::Arena*> arena,
       absl::string_view type_name) const;
 
-  TypeCheckEnv MakeExtendedEnvironment() const { return TypeCheckEnv(this); }
-  VariableScope MakeVariableScope() const { return VariableScope(*this); }
+  TypeCheckEnv MakeExtendedEnvironment() const ABSL_ATTRIBUTE_LIFETIME_BOUND {
+    return TypeCheckEnv(this);
+  }
+  VariableScope MakeVariableScope() const ABSL_ATTRIBUTE_LIFETIME_BOUND {
+    return VariableScope(*this);
+  }
+
+  absl::Nonnull<const google::protobuf::DescriptorPool*> descriptor_pool() const {
+    return descriptor_pool_.get();
+  }
 
  private:
+  explicit TypeCheckEnv(absl::Nonnull<const TypeCheckEnv*> parent)
+      : descriptor_pool_(parent->descriptor_pool_),
+        container_(parent != nullptr ? parent->container() : ""),
+        parent_(parent) {}
+
   absl::StatusOr<absl::optional<VariableDecl>> LookupEnumConstant(
       TypeFactory& type_factory, absl::string_view type,
       absl::string_view value) const;
 
+  absl::Nonnull<std::shared_ptr<const google::protobuf::DescriptorPool>> descriptor_pool_;
   std::string container_;
   absl::Nullable<const TypeCheckEnv*> parent_;