Add an optional protobuf Arena tied to TypeCheckerBuilder instances.

Intended to be used to make it easier to guarantee lifetime of types needed by the checker.

PiperOrigin-RevId: 733896891

Author: jnthntatum

checker/BUILD

@@ -88,10 +88,12 @@ cc_library(
         ":type_checker",
         "//common:decl",
         "//common:type",
+        "@com_google_absl//absl/base:nullability",
         "@com_google_absl//absl/functional:any_invocable",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings",
+        "@com_google_protobuf//:protobuf",
     ],
 )
 

checker/internal/type_check_env.h

@@ -179,6 +179,17 @@ class TypeCheckEnv {
     return descriptor_pool_.get();
   }
 
+  // Return an arena that can be used to allocate memory for types that will be
+  // used by the TypeChecker being built.
+  //
+  // This is only intended to be used for configuration.
+  absl::Nonnull<google::protobuf::Arena*> arena() {
+    if (arena_ == nullptr) {
+      arena_ = std::make_unique<google::protobuf::Arena>();
+    }
+    return arena_.get();
+  }
+
  private:
   explicit TypeCheckEnv(absl::Nonnull<const TypeCheckEnv*> parent)
       : descriptor_pool_(parent->descriptor_pool_),
@@ -189,6 +200,7 @@ class TypeCheckEnv {
       absl::string_view type, absl::string_view value) const;
 
   absl::Nonnull<std::shared_ptr<const google::protobuf::DescriptorPool>> descriptor_pool_;
+  absl::Nullable<std::unique_ptr<google::protobuf::Arena>> arena_;
   std::string container_;
   absl::Nullable<const TypeCheckEnv*> parent_;
 

checker/internal/type_checker_builder_impl.h

@@ -32,6 +32,7 @@
 #include "common/decl.h"
 #include "common/type.h"
 #include "common/type_introspector.h"
+#include "google/protobuf/arena.h"
 #include "google/protobuf/descriptor.h"
 
 namespace cel::checker_internal {
@@ -71,6 +72,13 @@ class TypeCheckerBuilderImpl : public TypeCheckerBuilder {
 
   const CheckerOptions& options() const override { return options_; }
 
+  absl::Nonnull<google::protobuf::Arena*> arena() override { return env_.arena(); }
+
+  absl::Nonnull<const google::protobuf::DescriptorPool*> descriptor_pool()
+      const override {
+    return env_.descriptor_pool();
+  }
+
  private:
   absl::Status AddContextDeclarationVariables(
       absl::Nonnull<const google::protobuf::Descriptor*> descriptor);

checker/type_checker.h

@@ -27,8 +27,8 @@ namespace cel {
 //
 // Checks references and type agreement for a parsed CEL expression.
 //
-// TODO: see Compiler for bundled parse and type check from a
-// source expression string.
+// See Compiler for bundled parse and type check from a source expression
+// string.
 class TypeChecker {
  public:
   virtual ~TypeChecker() = default;

checker/type_checker_builder.h

@@ -18,6 +18,7 @@
 #include <memory>
 #include <string>
 
+#include "absl/base/nullability.h"
 #include "absl/functional/any_invocable.h"
 #include "absl/status/status.h"
 #include "absl/status/statusor.h"
@@ -27,6 +28,8 @@
 #include "common/decl.h"
 #include "common/type.h"
 #include "common/type_introspector.h"
+#include "google/protobuf/arena.h"
+#include "google/protobuf/descriptor.h"
 
 namespace cel {
 
@@ -106,6 +109,16 @@ class TypeCheckerBuilder {
   // This operation is destructive: the builder instance should not be used
   // after this method is called.
   virtual absl::StatusOr<std::unique_ptr<TypeChecker>> Build() && = 0;
+
+  // Returns a pointer to an arena that can be used to allocate memory for types
+  // that will be used by the TypeChecker being built.
+  //
+  // On Build(), the arena is transferred to the TypeChecker being built.
+  virtual absl::Nonnull<google::protobuf::Arena*> arena() = 0;
+
+  // The configured descriptor pool.
+  virtual absl::Nonnull<const google::protobuf::DescriptorPool*> descriptor_pool()
+      const = 0;
 };
 
 }  // namespace cel

checker/type_checker_builder_factory.cc

@@ -1,4 +1,3 @@
-
 // Copyright 2024 Google LLC
 //
 // Licensed under the Apache License, Version 2.0 (the "License");

checker/type_checker_builder_factory.h

@@ -26,6 +26,9 @@ namespace cel {
 
 // Creates a new `TypeCheckerBuilder`.
 //
+// The builder implementation is thread-hostile and should only be used from a
+// single thread, but the resulting `TypeChecker` instance is thread-safe.
+//
 // When passing a raw pointer to a descriptor pool, the descriptor pool must
 // outlive the type checker builder and the type checker builder it creates.
 //

checker/type_checker_builder_factory_test.cc

@@ -49,6 +49,22 @@ TEST(TypeCheckerBuilderTest, AddVariable) {
   EXPECT_TRUE(result.IsValid());
 }
 
+TEST(TypeCheckerBuilderTest, AddComplexType) {
+  ASSERT_OK_AND_ASSIGN(
+      std::unique_ptr<TypeCheckerBuilder> builder,
+      CreateTypeCheckerBuilder(GetSharedTestingDescriptorPool()));
+
+  MapType map_type(builder->arena(), StringType(), IntType());
+
+  ASSERT_THAT(builder->AddVariable(MakeVariableDecl("m", map_type)), IsOk());
+
+  ASSERT_OK_AND_ASSIGN(auto checker, std::move(*builder).Build());
+  builder.reset();
+  ASSERT_OK_AND_ASSIGN(auto ast, MakeTestParsedAst("m.foo"));
+  ASSERT_OK_AND_ASSIGN(ValidationResult result, checker->Check(std::move(ast)));
+  EXPECT_TRUE(result.IsValid());
+}
+
 TEST(TypeCheckerBuilderTest, AddVariableRedeclaredError) {
   ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<TypeCheckerBuilder> builder,