Add ToBuilder() functions for parser/checker/compiler

Adds support for extending a configured compiler similar to cel-java.

PiperOrigin-RevId: 905220235

Author: jnthntatum

checker/BUILD

@@ -123,6 +123,7 @@ cc_test(
     srcs = ["type_checker_builder_factory_test.cc"],
     deps = [
         ":checker_options",
+        ":optional",
         ":standard_library",
         ":type_checker",
         ":type_checker_builder",

checker/internal/BUILD

@@ -74,6 +74,7 @@ cc_library(
         "@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/log:absl_check",
         "@com_google_absl//absl/memory",
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings",
@@ -175,6 +176,7 @@ cc_test(
         ":type_checker_impl",
         "//checker:checker_options",
         "//checker:type_check_issue",
+        "//checker:type_checker_builder",
         "//checker:validation_result",
         "//common:ast",
         "//common:container",

checker/internal/type_check_env.h

@@ -23,6 +23,7 @@
 #include "absl/base/attributes.h"
 #include "absl/base/nullability.h"
 #include "absl/container/flat_hash_map.h"
+#include "absl/log/absl_check.h"
 #include "absl/memory/memory.h"
 #include "absl/status/statusor.h"
 #include "absl/strings/string_view.h"
@@ -100,7 +101,8 @@ class TypeCheckEnv {
     type_providers_.push_back(proto_type_introspector_);
   }
 
-  // Move-only.
+  TypeCheckEnv(const TypeCheckEnv&) = default;
+  TypeCheckEnv& operator=(const TypeCheckEnv&) = default;
   TypeCheckEnv(TypeCheckEnv&&) = default;
   TypeCheckEnv& operator=(TypeCheckEnv&&) = default;
 
@@ -193,19 +195,26 @@ class TypeCheckEnv {
 
   // Used to keep an arena alive if one was needed to allocate types.
   //
-  // The TypeCheckEnv does not otherwise use it.
-  void set_arena(std::shared_ptr<const google::protobuf::Arena> arena) {
+  // Expected to be called exactly once if at all.
+  void set_arena(std::shared_ptr<google::protobuf::Arena> arena) {
+    ABSL_DCHECK(arena_ == nullptr || arena == arena_);
     arena_ = std::move(arena);
   }
 
+  // Returns the arena if one was set, nullptr otherwise.
+  std::shared_ptr<google::protobuf::Arena> arena() const { return arena_; }
+
  private:
   absl::StatusOr<absl::optional<VariableDecl>> LookupEnumConstant(
       absl::string_view type, absl::string_view value) const;
 
   absl_nonnull std::shared_ptr<const google::protobuf::DescriptorPool> descriptor_pool_;
 
   // If set, an arena was needed to allocate types in the environment.
-  absl_nullable std::shared_ptr<const google::protobuf::Arena> arena_;
+  //
+  // The TypeCheckEnv does not otherwise use the arena, though it may be used by
+  // derived TypeCheckerBuilders.
+  absl_nullable std::shared_ptr<google::protobuf::Arena> arena_;
   ExpressionContainer container_;
 
   // Used to resolve fields on message types.

checker/internal/type_checker_builder_impl.cc

@@ -34,6 +34,7 @@
 #include "checker/internal/type_checker_impl.h"
 #include "checker/type_checker.h"
 #include "checker/type_checker_builder.h"
+#include "common/container.h"
 #include "common/decl.h"
 #include "common/type.h"
 #include "common/type_introspector.h"
@@ -342,8 +343,16 @@ absl::Status TypeCheckerBuilderImpl::ApplyConfig(
 }
 
 absl::StatusOr<std::unique_ptr<TypeChecker>> TypeCheckerBuilderImpl::Build() {
-  TypeCheckEnv env(descriptor_pool_);
-  env.set_container(expression_container_);
+  TypeCheckEnv env(template_env_);
+  CEL_RETURN_IF_ERROR(ConfigureTypeCheckEnv(env));
+  return std::make_unique<checker_internal::TypeCheckerImpl>(std::move(env),
+                                                             options_);
+}
+
+absl::Status TypeCheckerBuilderImpl::ConfigureTypeCheckEnv(TypeCheckEnv& env) {
+  if (expression_container_.has_value()) {
+    env.set_container(*expression_container_);
+  }
   if (expected_type_.has_value()) {
     env.set_expected_type(*expected_type_);
   }
@@ -377,12 +386,10 @@ absl::StatusOr<std::unique_ptr<TypeChecker>> TypeCheckerBuilderImpl::Build() {
                                   /*subset=*/nullptr, env));
 
   CEL_RETURN_IF_ERROR(ApplyConfig(default_config_, /*subset=*/nullptr, env));
-  // A library may have been the first to initialize the arena, so we need to
-  // set it as the last step.
-  env.set_arena(arena_);
-  auto checker = std::make_unique<checker_internal::TypeCheckerImpl>(
-      std::move(env), options_);
-  return checker;
+  if (type_arena_ != nullptr) {
+    env.set_arena(type_arena_);
+  }
+  return absl::OkStatus();
 }
 
 absl::Status TypeCheckerBuilderImpl::AddLibrary(CheckerLibrary library) {
@@ -432,7 +439,7 @@ absl::Status TypeCheckerBuilderImpl::AddOrReplaceVariable(
 absl::Status TypeCheckerBuilderImpl::AddContextDeclaration(
     absl::string_view type) {
   const google::protobuf::Descriptor* desc =
-      descriptor_pool_->FindMessageTypeByName(type);
+      template_env_.descriptor_pool()->FindMessageTypeByName(type);
   if (desc == nullptr) {
     return absl::NotFoundError(
         absl::StrCat("context declaration '", type, "' not found"));
@@ -479,7 +486,10 @@ void TypeCheckerBuilderImpl::AddTypeProvider(
 }
 
 void TypeCheckerBuilderImpl::set_container(absl::string_view container) {
-  expression_container_.SetContainer(container).IgnoreError();
+  if (!expression_container_.has_value()) {
+    expression_container_.emplace();
+  }
+  expression_container_->SetContainer(container).IgnoreError();
 }
 
 void TypeCheckerBuilderImpl::SetExpressionContainer(

checker/internal/type_checker_builder_impl.h

@@ -40,8 +40,6 @@
 
 namespace cel::checker_internal {
 
-class TypeCheckerBuilderImpl;
-
 // Builder for TypeChecker instances.
 class TypeCheckerBuilderImpl : public TypeCheckerBuilder {
  public:
@@ -51,7 +49,18 @@ class TypeCheckerBuilderImpl : public TypeCheckerBuilder {
       const CheckerOptions& options)
       : options_(options),
         target_config_(&default_config_),
-        descriptor_pool_(std::move(descriptor_pool)) {}
+        template_env_(std::move(descriptor_pool)) {}
+
+  // Constructor for building an extended TypeChecker.
+  explicit TypeCheckerBuilderImpl(const CheckerOptions& options,
+                                  const TypeCheckEnv& template_env)
+      : options_(options),
+        target_config_(&default_config_),
+        template_env_(template_env) {
+    if (auto arena = template_env_.arena(); arena != nullptr) {
+      type_arena_ = std::move(arena);
+    }
+  }
 
   // Move only.
   TypeCheckerBuilderImpl(const TypeCheckerBuilderImpl&) = delete;
@@ -83,14 +92,14 @@ class TypeCheckerBuilderImpl : public TypeCheckerBuilder {
   const CheckerOptions& options() const override { return options_; }
 
   google::protobuf::Arena* absl_nonnull arena() override {
-    if (arena_ == nullptr) {
-      arena_ = std::make_shared<google::protobuf::Arena>();
+    if (type_arena_ == nullptr) {
+      type_arena_ = std::make_shared<google::protobuf::Arena>();
     }
-    return arena_.get();
+    return type_arena_.get();
   }
 
   const google::protobuf::DescriptorPool* absl_nonnull descriptor_pool() const override {
-    return descriptor_pool_.get();
+    return template_env_.descriptor_pool();
   }
 
  private:
@@ -129,19 +138,21 @@ class TypeCheckerBuilderImpl : public TypeCheckerBuilder {
   absl::Status ApplyConfig(ConfigRecord config, const TypeCheckerSubset* subset,
                            TypeCheckEnv& env);
 
+  absl::Status ConfigureTypeCheckEnv(TypeCheckEnv& env);
+
   CheckerOptions options_;
   // Default target for configuration changes. Used for direct calls to
   // AddVariable, AddFunction, etc.
   ConfigRecord default_config_;
   // Active target for configuration changes.
   // This is used to track which library the change is made on behalf of.
   ConfigRecord* absl_nonnull target_config_;
-  std::shared_ptr<const google::protobuf::DescriptorPool> descriptor_pool_;
-  std::shared_ptr<google::protobuf::Arena> arena_;
+  TypeCheckEnv template_env_;
+  std::shared_ptr<google::protobuf::Arena> type_arena_;
   std::vector<CheckerLibrary> libraries_;
   absl::flat_hash_map<std::string, TypeCheckerSubset> subsets_;
   absl::flat_hash_set<std::string> library_ids_;
-  ExpressionContainer expression_container_;
+  absl::optional<ExpressionContainer> expression_container_;
   absl::optional<Type> expected_type_;
 };
 

checker/internal/type_checker_impl.cc

@@ -37,8 +37,10 @@
 #include "checker/internal/format_type_name.h"
 #include "checker/internal/namespace_generator.h"
 #include "checker/internal/type_check_env.h"
+#include "checker/internal/type_checker_builder_impl.h"
 #include "checker/internal/type_inference_context.h"
 #include "checker/type_check_issue.h"
+#include "checker/type_checker_builder.h"
 #include "checker/validation_result.h"
 #include "common/ast.h"
 #include "common/ast_rewrite.h"
@@ -1326,4 +1328,8 @@ absl::StatusOr<ValidationResult> TypeCheckerImpl::Check(
   return ValidationResult(std::move(ast), std::move(issues));
 }
 
+std::unique_ptr<TypeCheckerBuilder> TypeCheckerImpl::ToBuilder() const {
+  return std::make_unique<TypeCheckerBuilderImpl>(options_, env_);
+}
+
 }  // namespace cel::checker_internal

checker/internal/type_checker_impl.h

@@ -22,6 +22,7 @@
 #include "checker/checker_options.h"
 #include "checker/internal/type_check_env.h"
 #include "checker/type_checker.h"
+#include "checker/type_checker_builder.h"
 #include "checker/validation_result.h"
 #include "common/ast.h"
 #include "google/protobuf/arena.h"
@@ -44,6 +45,8 @@ class TypeCheckerImpl : public TypeChecker {
   absl::StatusOr<ValidationResult> Check(
       std::unique_ptr<Ast> ast) const override;
 
+  std::unique_ptr<TypeCheckerBuilder> ToBuilder() const override;
+
  private:
   TypeCheckEnv env_;
   google::protobuf::Arena type_arena_;

checker/internal/type_checker_impl_test.cc

@@ -33,6 +33,7 @@
 #include "checker/internal/test_ast_helpers.h"
 #include "checker/internal/type_check_env.h"
 #include "checker/type_check_issue.h"
+#include "checker/type_checker_builder.h"
 #include "checker/validation_result.h"
 #include "common/ast.h"
 #include "common/container.h"
@@ -1413,6 +1414,44 @@ TEST(TypeCheckerImplTest, ExpectedTypeDoesntMatch) {
           "expected type 'map(string, string)' but found 'map(string, int)'")));
 }
 
+TEST(TypeCheckerImplTest, ToBuilder) {
+  TypeCheckEnv env(GetSharedTestingDescriptorPool());
+  TypeCheckerImpl impl(std::move(env));
+  auto builder = impl.ToBuilder();
+  ASSERT_THAT(builder->AddVariable(MakeVariableDecl("x", IntType())), IsOk());
+  ASSERT_OK_AND_ASSIGN(auto new_checker, builder->Build());
+
+  ASSERT_OK_AND_ASSIGN(auto ast, MakeTestParsedAst("x"));
+  ASSERT_OK_AND_ASSIGN(ValidationResult result,
+                       new_checker->Check(std::move(ast)));
+  EXPECT_TRUE(result.IsValid());
+}
+
+TEST(TypeCheckerImplTest, ToBuilderPropagatesArena) {
+  auto arena = std::make_shared<google::protobuf::Arena>();
+
+  TypeCheckEnv env(GetSharedTestingDescriptorPool());
+  env.set_arena(arena);
+
+  Type list_type = ListType(arena.get(), IntType());
+  ASSERT_TRUE(
+      env.InsertVariableIfAbsent(MakeVariableDecl("my_list", list_type)));
+
+  auto base_checker = std::make_unique<TypeCheckerImpl>(std::move(env));
+
+  std::unique_ptr<TypeCheckerBuilder> builder = base_checker->ToBuilder();
+
+  base_checker.reset();
+  arena.reset();
+
+  ASSERT_OK_AND_ASSIGN(auto derived_checker, builder->Build());
+
+  ASSERT_OK_AND_ASSIGN(auto ast, MakeTestParsedAst("my_list"));
+  ASSERT_OK_AND_ASSIGN(ValidationResult result,
+                       derived_checker->Check(std::move(ast)));
+  EXPECT_TRUE(result.IsValid());
+}
+
 TEST(TypeCheckerImplTest, BadSourcePosition) {
   TypeCheckEnv env(GetSharedTestingDescriptorPool());
 

checker/type_checker.h

@@ -23,6 +23,8 @@
 
 namespace cel {
 
+class TypeCheckerBuilder;
+
 // TypeChecker interface.
 //
 // Checks references and type agreement for a parsed CEL expression.
@@ -43,6 +45,9 @@ class TypeChecker {
   virtual absl::StatusOr<ValidationResult> Check(
       std::unique_ptr<Ast> ast) const = 0;
 
+  // Returns a builder initialized with the configuration of this type checker.
+  virtual std::unique_ptr<TypeCheckerBuilder> ToBuilder() const = 0;
+
   // TODO(uncreated-issue/73): add overload for cref AST.
 };
 

checker/type_checker_builder.h

@@ -35,7 +35,6 @@
 namespace cel {
 
 class TypeCheckerBuilder;
-class TypeCheckerBuilderImpl;
 
 // Functional implementation to apply the library features to a
 // TypeCheckerBuilder.