Add support for looking up type constants and enums.

PiperOrigin-RevId: 678867445

checker/BUILD

@@ -84,6 +84,7 @@ cc_library(
         "//checker/internal:type_check_env",
         "//checker/internal:type_checker_impl",
         "//common:decl",
+        "//common:type",
         "@com_google_absl//absl/container:flat_hash_set",
         "@com_google_absl//absl/functional:any_invocable",
         "@com_google_absl//absl/status",

checker/internal/BUILD

@@ -61,12 +61,20 @@ cc_library(
     srcs = ["type_check_env.cc"],
     hdrs = ["type_check_env.h"],
     deps = [
+        "//common:constant",
         "//common:decl",
+        "//common:type",
+        "//internal:status_macros",
         "@com_google_absl//absl/base:core_headers",
         "@com_google_absl//absl/base:nullability",
         "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/memory",
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings",
         "@com_google_absl//absl/strings:string_view",
+        "@com_google_absl//absl/types:optional",
+        "@com_google_absl//absl/types:span",
+        "@com_google_protobuf//:protobuf",
     ],
 )
 
@@ -118,15 +126,16 @@ cc_library(
         "//common:constant",
         "//common:decl",
         "//common:expr",
+        "//common:memory",
         "//common:source",
         "//common:type",
         "//common:type_kind",
+        "//extensions/protobuf:memory_manager",
         "//internal:status_macros",
         "@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/log:absl_check",
-        "@com_google_absl//absl/log:absl_log",
+        "@com_google_absl//absl/container:flat_hash_set",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings",
@@ -152,6 +161,7 @@ cc_test(
         "//common:ast",
         "//common:decl",
         "//common:type",
+        "//extensions/protobuf:value",
         "//internal:status_macros",
         "//internal:testing",
         "@com_google_absl//absl/base:no_destructor",
@@ -161,6 +171,7 @@ cc_test(
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:status_matchers",
         "@com_google_absl//absl/strings",
+        "@com_google_cel_spec//proto/test/v1/proto2:test_all_types_cc_proto",
         "@com_google_cel_spec//proto/test/v1/proto3:test_all_types_cc_proto",
         "@com_google_protobuf//:protobuf",
     ],

checker/internal/type_check_env.cc

@@ -14,9 +14,22 @@
 
 #include "checker/internal/type_check_env.h"
 
+#include <cstddef>
+#include <cstdint>
+#include <string>
+
 #include "absl/base/nullability.h"
+#include "absl/status/statusor.h"
+#include "absl/strings/str_cat.h"
 #include "absl/strings/string_view.h"
+#include "absl/types/optional.h"
+#include "common/constant.h"
 #include "common/decl.h"
+#include "common/type.h"
+#include "common/type_factory.h"
+#include "common/type_introspector.h"
+#include "internal/status_macros.h"
+#include "google/protobuf/arena.h"
 
 namespace cel::checker_internal {
 
@@ -44,6 +57,90 @@ absl::Nullable<const FunctionDecl*> TypeCheckEnv::LookupFunction(
   return nullptr;
 }
 
+absl::StatusOr<absl::optional<Type>> TypeCheckEnv::LookupTypeName(
+    TypeFactory& type_factory, absl::string_view name) const {
+  const TypeCheckEnv* scope = this;
+  while (scope != nullptr) {
+    for (auto iter = type_providers_.rbegin(); iter != type_providers_.rend();
+         ++iter) {
+      auto type = (*iter)->FindType(type_factory, name);
+      if (!type.ok() || type->has_value()) {
+        return type;
+      }
+    }
+    scope = scope->parent_;
+  }
+  return absl::nullopt;
+}
+
+absl::StatusOr<absl::optional<VariableDecl>> TypeCheckEnv::LookupEnumConstant(
+    TypeFactory& type_factory, absl::string_view type,
+    absl::string_view value) const {
+  const TypeCheckEnv* scope = this;
+  while (scope != nullptr) {
+    for (auto iter = type_providers_.rbegin(); iter != type_providers_.rend();
+         ++iter) {
+      auto enum_constant = (*iter)->FindEnumConstant(type_factory, type, value);
+      if (!enum_constant.ok()) {
+        return enum_constant.status();
+      }
+      if (enum_constant->has_value()) {
+        auto decl =
+            MakeVariableDecl(absl::StrCat((**enum_constant).type_full_name, ".",
+                                          (**enum_constant).value_name),
+                             (**enum_constant).type);
+        decl.set_value(
+            Constant(static_cast<int64_t>((**enum_constant).number)));
+        return decl;
+      }
+    }
+    scope = scope->parent_;
+  }
+  return absl::nullopt;
+}
+
+absl::StatusOr<absl::optional<VariableDecl>> TypeCheckEnv::LookupTypeConstant(
+    TypeFactory& type_factory, absl::Nonnull<google::protobuf::Arena*> arena,
+    absl::string_view name) const {
+  CEL_ASSIGN_OR_RETURN(absl::optional<Type> type,
+                       LookupTypeName(type_factory, name));
+  if (type.has_value()) {
+    return MakeVariableDecl(std::string(type->name()), TypeType(arena, *type));
+  }
+
+  if (name.find('.') != name.npos) {
+    size_t last_dot = name.rfind('.');
+    absl::string_view enum_name_candidate = name.substr(0, last_dot);
+    absl::string_view value_name_candidate = name.substr(last_dot + 1);
+    return LookupEnumConstant(type_factory, enum_name_candidate,
+                              value_name_candidate);
+  }
+
+  return absl::nullopt;
+}
+
+absl::StatusOr<absl::optional<StructTypeField>> TypeCheckEnv::LookupStructField(
+    TypeFactory& type_factory, absl::string_view type_name,
+    absl::string_view field_name) const {
+  const TypeCheckEnv* scope = this;
+  while (scope != nullptr) {
+    // 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
+    // checking field accesses.
+    for (auto iter = type_providers_.rbegin(); iter != type_providers_.rend();
+         ++iter) {
+      auto field_info = (*iter)->FindStructTypeFieldByName(
+          type_factory, type_name, field_name);
+      if (!field_info.ok() || field_info->has_value()) {
+        return field_info;
+      }
+    }
+    scope = scope->parent_;
+  }
+  return absl::nullopt;
+}
+
 absl::Nullable<const VariableDecl*> VariableScope::LookupVariable(
     absl::string_view name) const {
   const VariableScope* scope = this;

checker/internal/type_check_env.h

@@ -18,13 +18,22 @@
 #include <memory>
 #include <string>
 #include <utility>
+#include <vector>
 
 #include "absl/base/attributes.h"
 #include "absl/base/nullability.h"
 #include "absl/container/flat_hash_map.h"
 #include "absl/memory/memory.h"
+#include "absl/status/statusor.h"
 #include "absl/strings/string_view.h"
+#include "absl/types/optional.h"
+#include "absl/types/span.h"
+#include "common/constant.h"
 #include "common/decl.h"
+#include "common/type.h"
+#include "common/type_factory.h"
+#include "common/type_introspector.h"
+#include "google/protobuf/arena.h"
 
 namespace cel::checker_internal {
 
@@ -68,9 +77,10 @@ class VariableScope {
   absl::flat_hash_map<std::string, VariableDecl> variables_;
 };
 
-// Class managing the type check environment.
+// Class managing the state of the type check environment.
 //
-// Maintains lookup maps for variables and functions.
+// Maintains lookup maps for variables and functions and the set of type
+// providers.
 //
 // This class is thread-compatible.
 class TypeCheckEnv {
@@ -95,6 +105,14 @@ class TypeCheckEnv {
     container_ = std::move(container);
   }
 
+  absl::Span<const std::unique_ptr<TypeIntrospector>> type_providers() const {
+    return type_providers_;
+  }
+
+  void AddTypeProvider(std::unique_ptr<TypeIntrospector> provider) {
+    type_providers_.push_back(std::move(provider));
+  }
+
   const absl::flat_hash_map<std::string, VariableDecl>& variables() const {
     return variables_;
   }
@@ -132,16 +150,34 @@ class TypeCheckEnv {
   absl::Nullable<const FunctionDecl*> LookupFunction(
       absl::string_view name) const;
 
+  absl::StatusOr<absl::optional<Type>> LookupTypeName(
+      TypeFactory& type_factory, absl::string_view name) const;
+
+  absl::StatusOr<absl::optional<StructTypeField>> LookupStructField(
+      TypeFactory& type_factory, absl::string_view type_name,
+      absl::string_view field_name) const;
+
+  absl::StatusOr<absl::optional<VariableDecl>> LookupTypeConstant(
+      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); }
 
  private:
+  absl::StatusOr<absl::optional<VariableDecl>> LookupEnumConstant(
+      TypeFactory& type_factory, absl::string_view type,
+      absl::string_view value) const;
+
   std::string container_;
   absl::Nullable<const TypeCheckEnv*> parent_;
 
   // Maps fully qualified names to declarations.
   absl::flat_hash_map<std::string, VariableDecl> variables_;
   absl::flat_hash_map<std::string, FunctionDecl> functions_;
+
+  // Type providers for custom types.
+  std::vector<std::unique_ptr<TypeIntrospector>> type_providers_;
 };
 
 }  // namespace cel::checker_internal