Memoize enum lookup table in FlatExprBuilder.

PiperOrigin-RevId: 739351773

Author: jnthntatum

eval/compiler/BUILD

@@ -143,10 +143,12 @@ cc_library(
         "@com_google_absl//absl/container:node_hash_map",
         "@com_google_absl//absl/functional:any_invocable",
         "@com_google_absl//absl/log:absl_check",
+        "@com_google_absl//absl/log:absl_log",
         "@com_google_absl//absl/log:check",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/synchronization",
         "@com_google_absl//absl/types:optional",
         "@com_google_absl//absl/types:span",
         "@com_google_absl//absl/types:variant",
@@ -368,7 +370,9 @@ cc_test(
         "//runtime/internal:issue_collector",
         "//runtime/internal:runtime_env",
         "//runtime/internal:runtime_env_testing",
+        "@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/status",
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings",
@@ -417,11 +421,12 @@ cc_library(
         "//internal:status_macros",
         "//runtime:function_overload_reference",
         "//runtime:function_registry",
-        "//runtime:type_registry",
+        "@com_google_absl//absl/base:no_destructor",
         "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings",
         "@com_google_absl//absl/types:optional",
+        "@com_google_absl//absl/types:span",
     ],
 )
 
@@ -436,19 +441,22 @@ cc_test(
         "//base:ast",
         "//base:builtins",
         "//common:expr",
+        "//common:value",
         "//common/ast:ast_impl",
         "//common/ast:expr",
         "//common/ast:expr_proto",
         "//eval/public:builtin_func_registrar",
         "//eval/public:cel_function",
         "//eval/public:cel_function_registry",
+        "//eval/public:cel_value",
         "//extensions/protobuf:ast_converters",
         "//internal:casts",
         "//internal:proto_matchers",
         "//internal:testing",
         "//runtime:runtime_issue",
         "//runtime:type_registry",
         "//runtime/internal:issue_collector",
+        "@com_google_absl//absl/base:no_destructor",
         "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/log:absl_check",
         "@com_google_absl//absl/memory",
@@ -496,7 +504,6 @@ cc_test(
         "//eval/testutil:test_message_cc_proto",
         "//internal:testing",
         "@com_google_absl//absl/status",
-        "@com_google_absl//absl/types:optional",
         "@com_google_absl//absl/types:span",
         "@com_google_protobuf//:protobuf",
     ],

eval/compiler/cel_expression_builder_flat_impl.h

@@ -74,6 +74,7 @@ class CelExpressionBuilderFlatImpl : public CelExpressionBuilder {
   FlatExprBuilder& flat_expr_builder() { return flat_expr_builder_; }
 
   void set_container(std::string container) override {
+    flat_expr_builder_.InvalidateResolverIndex();
     flat_expr_builder_.set_container(std::move(container));
   }
 
@@ -87,6 +88,7 @@ class CelExpressionBuilderFlatImpl : public CelExpressionBuilder {
   // CelValue instances, and to extend the set of types and enums known to
   // expressions by registering them ahead of time.
   CelTypeRegistry* GetTypeRegistry() const override {
+    flat_expr_builder_.InvalidateResolverIndex();
     return &env_->legacy_type_registry;
   }
 

eval/compiler/constant_folding_test.cc

@@ -15,10 +15,14 @@
 #include "eval/compiler/constant_folding.h"
 
 #include <memory>
+#include <string>
 #include <utility>
+#include <vector>
 
 #include "cel/expr/syntax.pb.h"
+#include "absl/base/no_destructor.h"
 #include "absl/base/nullability.h"
+#include "absl/container/flat_hash_map.h"
 #include "absl/status/status.h"
 #include "absl/status/statusor.h"
 #include "absl/strings/string_view.h"
@@ -68,16 +72,28 @@ using ::google::api::expr::runtime::ProgramOptimizerFactory;
 using ::google::api::expr::runtime::Resolver;
 using ::testing::SizeIs;
 
+const std::vector<std::string>& EmptyNamespacePrefixes() {
+  static const absl::NoDestructor<std::vector<std::string>> kEmptyPrefixes(
+      {""});
+  return *kEmptyPrefixes;
+}
+
+const absl::flat_hash_map<std::string, cel::Value>& EmptyEnumValueMap() {
+  static const absl::NoDestructor<absl::flat_hash_map<std::string, cel::Value>>
+      kEmptyEnumValueMap({});
+  return *kEmptyEnumValueMap;
+}
+
 class UpdatedConstantFoldingTest : public testing::Test {
  public:
   UpdatedConstantFoldingTest()
       : env_(NewTestingRuntimeEnv()),
         function_registry_(env_->function_registry),
         type_registry_(env_->type_registry),
         issue_collector_(RuntimeIssue::Severity::kError),
-        resolver_("", function_registry_, type_registry_,
-                  type_registry_.GetComposedTypeProvider(),
-                  type_registry_.resolveable_enums()) {}
+        resolver_(EmptyNamespacePrefixes(), EmptyEnumValueMap(),
+                  function_registry_,
+                  type_registry_.GetComposedTypeProvider()) {}
 
  protected:
   absl::Nonnull<std::shared_ptr<RuntimeEnv>> env_;

eval/compiler/flat_expr_builder.cc

@@ -41,6 +41,7 @@
 #include "absl/strings/match.h"
 #include "absl/strings/numbers.h"
 #include "absl/strings/str_cat.h"
+#include "absl/strings/str_split.h"
 #include "absl/strings/string_view.h"
 #include "absl/strings/strip.h"
 #include "absl/types/optional.h"
@@ -2459,14 +2460,50 @@ std::vector<ExecutionPathView> FlattenExpressionTable(
 
 }  // namespace
 
+std::shared_ptr<FlatExprBuilder::ResolverIndex>
+FlatExprBuilder::BuildResolverIndex() const {
+  auto result = std::make_shared<ResolverIndex>();
+  auto& namespace_prefixes = result->namespace_prefixes;
+  auto& enum_value_map = result->enum_value_map;
+  const auto& resolveable_enums = type_registry_.resolveable_enums();
+
+  std::string prefix = "";
+  namespace_prefixes.push_back(prefix);
+  auto container_elements = absl::StrSplit(container_, '.');
+  for (const auto& elem : container_elements) {
+    // Tolerate trailing / leading '.'.
+    if (elem.empty()) {
+      continue;
+    }
+    absl::StrAppend(&prefix, elem, ".");
+    // longest prefix first.
+    namespace_prefixes.insert(namespace_prefixes.begin(), prefix);
+  }
+
+  for (auto iter = resolveable_enums.begin(); iter != resolveable_enums.end();
+       ++iter) {
+    absl::string_view enum_name = iter->first;
+    const auto& enum_type = iter->second;
+
+    for (const auto& enumerator : enum_type.enumerators) {
+      auto key = absl::StrCat(enum_name, ".", enumerator.name);
+      enum_value_map[key] = cel::IntValue(enumerator.number);
+    }
+  }
+
+  return result;
+}
+
 absl::StatusOr<FlatExpression> FlatExprBuilder::CreateExpressionImpl(
     std::unique_ptr<Ast> ast, std::vector<RuntimeIssue>* issues) const {
   RuntimeIssue::Severity max_severity = options_.fail_on_warnings
                                             ? RuntimeIssue::Severity::kWarning
                                             : RuntimeIssue::Severity::kError;
   IssueCollector issue_collector(max_severity);
-  Resolver resolver(container_, function_registry_, type_registry_,
-                    GetTypeProvider(), type_registry_.resolveable_enums(),
+  auto resolver_index = GetResolverIndex();
+  Resolver resolver(resolver_index->namespace_prefixes,
+                    resolver_index->enum_value_map, function_registry_,
+                    GetTypeProvider(),
                     options_.enable_qualified_type_identifiers);
 
   std::shared_ptr<google::protobuf::Arena> arena;

eval/compiler/flat_expr_builder.h

@@ -23,10 +23,15 @@
 #include <vector>
 
 #include "absl/base/nullability.h"
+#include "absl/base/thread_annotations.h"
+#include "absl/container/flat_hash_map.h"
+#include "absl/log/absl_log.h"
 #include "absl/status/statusor.h"
 #include "absl/strings/string_view.h"
+#include "absl/synchronization/mutex.h"
 #include "base/ast.h"
 #include "base/type_provider.h"
+#include "common/value.h"
 #include "eval/compiler/flat_expr_builder_extensions.h"
 #include "eval/eval/evaluator_core.h"
 #include "runtime/function_registry.h"
@@ -93,11 +98,55 @@ class FlatExprBuilder {
   // `optional_type` handling is needed.
   void enable_optional_types() { enable_optional_types_ = true; }
 
+  // Note: this is a temporary solution to support the legacy expression
+  // builder which remains mutable after building expressions. This is not
+  // correct for thread safety (e.g. storing a reference to the type
+  // registry), however the builder is otherwise thread hostile if used in this
+  // way so only mitigates some cases.
+  void InvalidateResolverIndex() const {
+    absl::MutexLock lock(&resolver_index_mutex_);
+    if (resolver_index_ != nullptr) {
+      ABSL_LOG(WARNING)
+          << "attempted to update CEL expression builder after use";
+    }
+    resolver_index_.reset();
+  }
+
  private:
   const cel::TypeProvider& GetTypeProvider() const;
 
   const absl::Nonnull<std::shared_ptr<const cel::runtime_internal::RuntimeEnv>>
       env_;
+
+  struct ResolverIndex {
+    std::vector<std::string> namespace_prefixes;
+    absl::flat_hash_map<std::string, cel::Value> enum_value_map;
+  };
+
+  std::shared_ptr<ResolverIndex> BuildResolverIndex() const;
+
+  std::shared_ptr<ResolverIndex> GetResolverIndex() const {
+    std::shared_ptr<ResolverIndex> result;
+    {
+      absl::ReaderMutexLock lock(&resolver_index_mutex_);
+      result = resolver_index_;
+    }
+
+    if (result != nullptr) {
+      return result;
+    }
+    // Slow path: build the resolver index.
+    absl::MutexLock lock(&resolver_index_mutex_);
+    result = resolver_index_;
+    if (result != nullptr) {
+      return result;
+    }
+
+    result = BuildResolverIndex();
+    resolver_index_ = result;
+    return result;
+  }
+
   cel::RuntimeOptions options_;
   std::string container_;
   bool enable_optional_types_ = false;
@@ -108,6 +157,11 @@ class FlatExprBuilder {
   bool use_legacy_type_provider_;
   std::vector<std::unique_ptr<AstTransform>> ast_transforms_;
   std::vector<ProgramOptimizerFactory> program_optimizers_;
+
+  // See note on accessor.
+  mutable std::shared_ptr<ResolverIndex> resolver_index_
+      ABSL_GUARDED_BY(resolver_index_mutex_);
+  mutable absl::Mutex resolver_index_mutex_;
 };
 
 }  // namespace google::api::expr::runtime

eval/compiler/flat_expr_builder_extensions_test.cc

@@ -55,15 +55,26 @@ using ::testing::Optional;
 
 using Subexpression = ProgramBuilder::Subexpression;
 
+const std::vector<std::string>& EmptyNamespacePrefixes() {
+  static const absl::NoDestructor<std::vector<std::string>> kEmptyPrefixes(
+      {""});
+  return *kEmptyPrefixes;
+}
+
+const absl::flat_hash_map<std::string, cel::Value>& EmptyEnumValueMap() {
+  static const absl::NoDestructor<absl::flat_hash_map<std::string, cel::Value>>
+      kEmptyEnumValueMap({});
+  return *kEmptyEnumValueMap;
+}
+
 class PlannerContextTest : public testing::Test {
  public:
   PlannerContextTest()
       : env_(NewTestingRuntimeEnv()),
         type_registry_(env_->type_registry),
         function_registry_(env_->function_registry),
-        resolver_("", function_registry_, type_registry_,
-                  type_registry_.GetComposedTypeProvider(),
-                  type_registry_.resolveable_enums()),
+        resolver_(EmptyNamespacePrefixes(), EmptyEnumValueMap(),
+                  function_registry_, type_registry_.GetComposedTypeProvider()),
         issue_collector_(RuntimeIssue::Severity::kError) {}
 
  protected: