Prototype: Refactor scope handling to allow wiring in the runtime representation of annotations.

PiperOrigin-RevId: 723628841

Author: jnthntatum

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

@@ -137,13 +137,15 @@ cc_library(
         "//common:source",
         "//common:type",
         "//common:type_kind",
+        "//internal:annotations",
         "//internal:status_macros",
         "//parser:macro",
         "@com_google_absl//absl/algorithm:container",
         "@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/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;
 

checker/internal/type_checker_impl.cc

@@ -51,6 +51,7 @@
 #include "common/source.h"
 #include "common/type.h"
 #include "common/type_kind.h"
+#include "internal/annotations.h"
 #include "internal/status_macros.h"
 #include "google/protobuf/arena.h"
 
@@ -68,6 +69,10 @@ std::string FormatCandidate(absl::Span<const std::string> qualifiers) {
   return absl::StrJoin(qualifiers, ".");
 }
 
+static const TraversalOptions kTraversalOptions = {
+    /*.use_comprehension_callbacks=*/true,
+};
+
 SourceLocation ComputeSourceLocation(const AstImpl& ast, int64_t expr_id) {
   const auto& source_info = ast.source_info();
   auto iter = source_info.positions().find(expr_id);
@@ -235,6 +240,8 @@ absl::StatusOr<AstType> FlattenType(const Type& type) {
   }
 }
 
+using AnnotationMap = internal::AnnotationMap;
+
 class ResolveVisitor : public AstVisitorBase {
  public:
   struct FunctionResolution {
@@ -247,13 +254,17 @@ class ResolveVisitor : public AstVisitorBase {
                  const TypeCheckEnv& env, const AstImpl& ast,
                  TypeInferenceContext& inference_context,
                  std::vector<TypeCheckIssue>& issues,
+                 AnnotationMap& annotations,
+                 cel::CheckerAnnotationSupport annotation_support,
                  absl::Nonnull<google::protobuf::Arena*> arena)
       : container_(container),
+        annotation_support_(annotation_support),
         namespace_generator_(std::move(namespace_generator)),
         env_(&env),
         inference_context_(&inference_context),
         issues_(&issues),
         ast_(&ast),
+        annotations_(&annotations),
         root_scope_(env.MakeVariableScope()),
         arena_(arena),
         current_scope_(&root_scope_) {}
@@ -265,6 +276,43 @@ class ResolveVisitor : public AstVisitorBase {
       return;
     }
     expr_stack_.pop_back();
+    if (!status_.ok()) {
+      return;
+    }
+    if (annotation_support_ != CheckerAnnotationSupport::kCheck) {
+      return;
+    }
+    auto annotations = annotations_->find(expr.id());
+    if (annotations == annotations_->end()) {
+      return;
+    }
+    if (annotation_context_.has_value()) {
+      issues_->push_back(
+          TypeCheckIssue::CreateError(ComputeSourceLocation(*ast_, expr.id()),
+                                      "Nested annotations are not supported."));
+      return;
+    }
+    auto annotation_scope = current_scope_->MakeNestedScope();
+    VariableScope* annotation_scope_ptr = annotation_scope.get();
+    // bit of a misuse, but annotation scope is largely the same as for
+    // comprehensions.
+    comprehension_vars_.push_back(std::move(annotation_scope));
+
+    annotation_context_ = {current_scope_};
+    current_scope_ = annotation_scope_ptr;
+    Type annotated_expr_type = GetDeducedType(&expr);
+    annotation_scope_ptr->InsertVariableIfAbsent(
+        MakeVariableDecl("cel.annotated_value", annotated_expr_type));
+
+    // Note: this does not need to happen now during the main traversal, but
+    // it's a easier to reason about for me. It's equally valid to just record
+    // the relevant annotations and do a separate check pass later.
+    for (const auto& annotation : annotations->second) {
+      CheckAnnotation(annotation, expr, annotated_expr_type);
+    }
+
+    current_scope_ = annotation_context_->parent;
+    annotation_context_.reset();
   }
 
   void PostVisitConst(const Expr& expr, const Constant& constant) override;
@@ -341,6 +389,10 @@ class ResolveVisitor : public AstVisitorBase {
     const FunctionDecl* decl;
   };
 
+  struct AnnotationContext {
+    const VariableScope* parent;
+  };
+
   void ResolveSimpleIdentifier(const Expr& expr, absl::string_view name);
 
   void ResolveQualifiedIdentifier(const Expr& expr,
@@ -459,12 +511,17 @@ class ResolveVisitor : public AstVisitorBase {
     return DynType();
   }
 
+  void CheckAnnotation(const internal::AnnotationRep& annotation_expr,
+                       const Expr& annotated_expr, const Type& annotated_type);
+
   absl::string_view container_;
+  CheckerAnnotationSupport annotation_support_;
   NamespaceGenerator namespace_generator_;
   absl::Nonnull<const TypeCheckEnv*> env_;
   absl::Nonnull<TypeInferenceContext*> inference_context_;
   absl::Nonnull<std::vector<TypeCheckIssue>*> issues_;
   absl::Nonnull<const ast_internal::AstImpl*> ast_;
+  absl::Nonnull<AnnotationMap*> annotations_;
   VariableScope root_scope_;
   absl::Nonnull<google::protobuf::Arena*> arena_;
 
@@ -479,6 +536,7 @@ class ResolveVisitor : public AstVisitorBase {
   absl::flat_hash_set<const Expr*> deferred_select_operations_;
   std::vector<std::unique_ptr<VariableScope>> comprehension_vars_;
   std::vector<ComprehensionScope> comprehension_scopes_;
+  absl::optional<AnnotationContext> annotation_context_;
   absl::Status status_;
   int error_count_ = 0;
 
@@ -535,6 +593,64 @@ void ResolveVisitor::PostVisitIdent(const Expr& expr, const IdentExpr& ident) {
   }
 }
 
+void ResolveVisitor::CheckAnnotation(const internal::AnnotationRep& annotation,
+                                     const Expr& annotated_expr,
+                                     const Type& annotated_type) {
+  const auto* annotation_decl = env_->LookupAnnotation(annotation.name);
+  if (annotation_decl == nullptr) {
+    ReportIssue(TypeCheckIssue::CreateError(
+        ComputeSourceLocation(*ast_, annotated_expr.id()),
+        absl::StrCat("undefined annotation '", annotation.name, "'")));
+    return;
+  }
+
+  // Checking if assignable to Dyn may influence the type inference so skip
+  // here.
+  if (!annotation_decl->applicable_type().IsDyn()) {
+    if (!inference_context_->IsAssignable(annotated_type,
+                                          annotation_decl->applicable_type())) {
+      ReportIssue(TypeCheckIssue::CreateError(
+          ComputeSourceLocation(*ast_, annotated_expr.id()),
+          absl::StrCat(
+              "annotation '", annotation.name, "' is not applicable to type '",
+              inference_context_->FinalizeType(annotated_type).DebugString(),
+              "'")));
+      return;
+    }
+  }
+
+  if (annotation.inspect_only) {
+    // Nothing to do -- the value expression is not intended to be evaluated.
+    // Examples are for things like a pointer to another file if the
+    // subexpression is inlined from somewhere else.
+    return;
+  }
+
+  // TODO - re-entrant traversal bypasses the complexity limits.
+  AstTraverse(*annotation.value_expr, *this, kTraversalOptions);
+
+  if (!status_.ok()) {
+    return;
+  }
+
+  Type value_expression_type = GetDeducedType(annotation.value_expr);
+
+  if (!annotation_decl->expected_type().IsDyn()) {
+    if (!inference_context_->IsAssignable(value_expression_type,
+                                          annotation_decl->expected_type())) {
+      ReportIssue(TypeCheckIssue::CreateError(
+          ComputeSourceLocation(*ast_, annotated_expr.id()),
+          absl::StrCat("annotation '", annotation.name,
+                       "' value expression type '",
+                       inference_context_->FinalizeType(value_expression_type)
+                           .DebugString(),
+                       "' is not assignable to '",
+                       annotation_decl->expected_type().DebugString(), "'")));
+      return;
+    }
+  }
+}
+
 void ResolveVisitor::PostVisitConst(const Expr& expr,
                                     const Constant& constant) {
   switch (constant.kind().index()) {
@@ -1276,13 +1392,28 @@ absl::StatusOr<ValidationResult> TypeCheckerImpl::Check(
 
   TypeInferenceContext type_inference_context(
       &type_arena, options_.enable_legacy_null_assignment);
+
+  internal::AnnotationMap annotation_exprs;
+  Expr* root = &ast_impl.root_expr();
+  if (ast_impl.root_expr().has_call_expr() &&
+      ast_impl.root_expr().call_expr().function() == "cel.@annotated" &&
+      ast_impl.root_expr().call_expr().args().size() == 2) {
+    if (options_.annotation_support == CheckerAnnotationSupport::kStrip) {
+      ast_impl.root_expr() =
+          std::move(ast_impl.root_expr().mutable_call_expr().mutable_args()[0]);
+      root = &ast_impl.root_expr();
+    } else {
+      annotation_exprs = internal::BuildAnnotationMap(ast_impl);
+      root = &ast_impl.root_expr().mutable_call_expr().mutable_args()[0];
+    }
+  }
+
   ResolveVisitor visitor(env_.container(), std::move(generator), env_, ast_impl,
-                         type_inference_context, issues, &type_arena);
+                         type_inference_context, issues, annotation_exprs,
+                         options_.annotation_support, &type_arena);
 
-  TraversalOptions opts;
-  opts.use_comprehension_callbacks = true;
   bool error_limit_reached = false;
-  auto traversal = AstTraversal::Create(ast_impl.root_expr(), opts);
+  auto traversal = AstTraversal::Create(*root, kTraversalOptions);
 
   for (int step = 0; step < options_.max_expression_node_count * 2; ++step) {
     bool has_next = traversal.Step(visitor);
@@ -1300,7 +1431,7 @@ absl::StatusOr<ValidationResult> TypeCheckerImpl::Check(
 
   if (!traversal.IsDone() && !error_limit_reached) {
     return absl::InvalidArgumentError(
-        absl::StrCat("Maximum expression node count exceeded: ",
+        absl::StrCat("maximum expression node count exceeded: ",
                      options_.max_expression_node_count));
   }
 
@@ -1309,7 +1440,7 @@ absl::StatusOr<ValidationResult> TypeCheckerImpl::Check(
         {}, absl::StrCat("maximum number of ERROR issues exceeded: ",
                          options_.max_error_issues)));
   } else if (env_.expected_type().has_value()) {
-    visitor.AssertExpectedType(ast_impl.root_expr(), *env_.expected_type());
+    visitor.AssertExpectedType(*root, *env_.expected_type());
   }
 
   // If any issues are errors, return without an AST.
@@ -1324,7 +1455,14 @@ absl::StatusOr<ValidationResult> TypeCheckerImpl::Check(
   // been invalidated by other updates.
   ResolveRewriter rewriter(visitor, type_inference_context, options_,
                            ast_impl.reference_map(), ast_impl.type_map());
-  AstRewrite(ast_impl.root_expr(), rewriter);
+  AstRewrite(*root, rewriter);
+  if (options_.annotation_support == CheckerAnnotationSupport::kCheck) {
+    for (auto& annotations : annotation_exprs) {
+      for (auto& annotation : annotations.second) {
+        AstRewrite(*annotation.value_expr, rewriter);
+      }
+    }
+  }
 
   CEL_RETURN_IF_ERROR(rewriter.status());