Prototype: add support for type-checking annotations.

checker/checker_options.h

@@ -17,6 +17,8 @@
 
 namespace cel {
 
+enum class CheckerAnnotationSupport { kStrip, kRetain, kCheck };
+
 // Options for enabling core type checker features.
 struct CheckerOptions {
   // Enable overloads for numeric comparisons across types.
@@ -55,6 +57,12 @@ struct CheckerOptions {
   // If exceeded, the checker will stop processing the ast and return
   // the current set of issues.
   int max_error_issues = 20;
+
+  // Annotation support level.
+  //
+  // Default behavior is to strip annotations.
+  CheckerAnnotationSupport annotation_support =
+      CheckerAnnotationSupport::kStrip;
 };
 
 }  // namespace cel

checker/internal/BUILD

@@ -119,6 +119,7 @@ cc_library(
         ":namespace_generator",
         ":type_check_env",
         ":type_inference_context",
+        "//base:googleinit",
         "//base/ast_internal:ast_impl",
         "//base/ast_internal:expr",
         "//checker:checker_options",
@@ -144,6 +145,7 @@ cc_library(
         "@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/log:absl_log",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings",

checker/internal/type_check_env.cc

@@ -45,6 +45,19 @@ absl::Nullable<const VariableDecl*> TypeCheckEnv::LookupVariable(
   return nullptr;
 }
 
+absl::Nullable<const AnnotationDecl*> TypeCheckEnv::LookupAnnotation(
+    absl::string_view name) const {
+  const TypeCheckEnv* scope = this;
+  while (scope != nullptr) {
+    if (auto it = scope->annotations_.find(name);
+        it != scope->annotations_.end()) {
+      return &it->second;
+    }
+    scope = scope->parent_;
+  }
+  return nullptr;
+}
+
 absl::Nullable<const FunctionDecl*> TypeCheckEnv::LookupFunction(
     absl::string_view name) const {
   const TypeCheckEnv* scope = this;

checker/internal/type_check_env.h

@@ -84,10 +84,6 @@ class VariableScope {
 //
 // This class is thread-compatible.
 class TypeCheckEnv {
- private:
-  using VariableDeclPtr = absl::Nonnull<const VariableDecl*>;
-  using FunctionDeclPtr = absl::Nonnull<const FunctionDecl*>;
-
  public:
   explicit TypeCheckEnv(
       absl::Nonnull<std::shared_ptr<const google::protobuf::DescriptorPool>>
@@ -130,6 +126,10 @@ class TypeCheckEnv {
     return variables_.insert({decl.name(), std::move(decl)}).second;
   }
 
+  bool InsertAnnotationIfAbsent(AnnotationDecl decl) {
+    return annotations_.insert({decl.name(), std::move(decl)}).second;
+  }
+
   const absl::flat_hash_map<std::string, FunctionDecl>& functions() const {
     return functions_;
   }
@@ -158,6 +158,8 @@ class TypeCheckEnv {
       absl::string_view name) const;
   absl::Nullable<const FunctionDecl*> LookupFunction(
       absl::string_view name) const;
+  absl::Nullable<const AnnotationDecl*> LookupAnnotation(
+      absl::string_view name) const;
 
   absl::StatusOr<absl::optional<Type>> LookupTypeName(
       absl::string_view name) const;
@@ -195,6 +197,7 @@ class TypeCheckEnv {
   // Maps fully qualified names to declarations.
   absl::flat_hash_map<std::string, VariableDecl> variables_;
   absl::flat_hash_map<std::string, FunctionDecl> functions_;
+  absl::flat_hash_map<std::string, AnnotationDecl> annotations_;
 
   // Type providers for custom types.
   std::vector<std::unique_ptr<TypeIntrospector>> type_providers_;

checker/internal/type_checker_builder_impl.cc

@@ -205,6 +205,17 @@ absl::Status TypeCheckerBuilderImpl::MergeFunction(const FunctionDecl& decl) {
   return absl::OkStatus();
 }
 
+absl::Status TypeCheckerBuilderImpl::AddAnnotation(const AnnotationDecl& decl) {
+  if (decl.name().empty()) {
+    return absl::InvalidArgumentError("annotation name must not be empty");
+  }
+  if (!env_.InsertAnnotationIfAbsent(decl)) {
+    return absl::AlreadyExistsError(
+        absl::StrCat("annotation '", decl.name(), "' already exists"));
+  }
+  return absl::OkStatus();
+}
+
 void TypeCheckerBuilderImpl::AddTypeProvider(
     std::unique_ptr<TypeIntrospector> provider) {
   env_.AddTypeProvider(std::move(provider));

checker/internal/type_checker_builder_impl.h

@@ -60,6 +60,7 @@ class TypeCheckerBuilderImpl : public TypeCheckerBuilder {
   absl::Status AddVariable(const VariableDecl& decl) override;
   absl::Status AddContextDeclaration(absl::string_view type) override;
   absl::Status AddFunction(const FunctionDecl& decl) override;
+  absl::Status AddAnnotation(const AnnotationDecl& decl) override;
 
   void SetExpectedType(const Type& type) override;