Default opt in select optimization to use eagerly resolve field descriptor lookups.

PiperOrigin-RevId: 1002720724

Author: jnthntatum

eval/compiler/BUILD

@@ -31,8 +31,10 @@ cc_library(
         ":resolver",
         "//base:ast",
         "//base:data",
+        "//common:ast",
         "//common:expr",
         "//common:type",
+        "//common:type_spec_resolver",
         "//common:value",
         "//eval/eval:evaluator_core",
         "//eval/eval:expression_step_logic",
@@ -147,21 +149,11 @@ cc_test(
         ":cel_expression_builder_flat_impl",
         ":constant_folding",
         ":flat_expr_builder",
-        ":flat_expr_builder_extensions",
         ":qualified_reference_resolver",
         "//base:builtins",
-        "//checker:validation_result",
-        "//common:ast",
-        "//common:decl",
-        "//common:expr",
         "//common:function_descriptor",
         "//common:kind",
-        "//common:type",
         "//common:value",
-        "//compiler",
-        "//compiler:compiler_factory",
-        "//compiler:standard_library",
-        "//eval/eval:evaluator_core",
         "//eval/eval:expression_step_logic",
         "//eval/public:activation",
         "//eval/public:ast_rewrite",

eval/compiler/flat_expr_builder.cc

@@ -1860,13 +1860,7 @@ void FlatExprVisitor::MaybeResolveType(const cel::Expr& expr) {
   if (it == type_map_.end()) {
     return;
   }
-  absl::StatusOr<cel::Type> type = cel::ConvertTypeSpecToType(
-      it->second, extension_context_.type_reflector(),
-      extension_context_.MutableArena());
-  if (!type.ok()) {
-    return;
-  }
-  extension_context_.SetResolvedType(expr.id(), *type);
+  extension_context_.TryResolveType(expr.id(), it->second);
 }
 
 void LogicalCondVisitor::PreVisit(const cel::Expr* expr) {

eval/compiler/flat_expr_builder_extensions.cc

@@ -14,8 +14,10 @@
 #include "eval/compiler/flat_expr_builder_extensions.h"
 
 #include <cstddef>
+#include <cstdint>
 #include <iterator>
 #include <memory>
+#include <optional>
 #include <utility>
 #include <vector>
 
@@ -27,7 +29,10 @@
 #include "absl/status/status.h"
 #include "absl/status/statusor.h"
 #include "absl/types/variant.h"
+#include "common/ast.h"
 #include "common/expr.h"
+#include "common/type.h"
+#include "common/type_spec_resolver.h"
 #include "eval/eval/evaluator_core.h"
 
 namespace google::api::expr::runtime {
@@ -350,6 +355,21 @@ absl::Status PlannerContext::ReplaceSubplan(const cel::Expr& node,
   return absl::OkStatus();
 }
 
+std::optional<cel::Type> PlannerContext::TryResolveType(
+    int64_t expr_id, const cel::TypeSpec& ast_type) {
+  auto it = resolved_types_.find(expr_id);
+  if (it != resolved_types_.end()) {
+    return it->second;
+  }
+  absl::StatusOr<cel::Type> type =
+      cel::ConvertTypeSpecToType(ast_type, type_reflector_, MutableArena());
+  if (!type.ok()) {
+    return std::nullopt;
+  }
+  resolved_types_.insert(it, {expr_id, *type});
+  return *type;
+}
+
 void ProgramBuilder::Reset() {
   root_ = nullptr;
   current_ = nullptr;

eval/compiler/flat_expr_builder_extensions.h

@@ -381,6 +381,9 @@ class PlannerContext {
     return std::nullopt;
   }
 
+  std::optional<cel::Type> TryResolveType(int64_t expr_id,
+                                          const cel::TypeSpec& ast_type);
+
   void SetResolvedType(int64_t expr_id, cel::Type type) {
     resolved_types_[expr_id] = std::move(type);
   }

eval/compiler/flat_expr_builder_test.cc

@@ -18,7 +18,6 @@
 #include <functional>
 #include <limits>
 #include <memory>
-#include <optional>
 #include <string>
 #include <utility>
 #include <vector>
@@ -36,22 +35,12 @@
 #include "absl/time/time.h"
 #include "absl/types/span.h"
 #include "base/builtins.h"
-#include "checker/validation_result.h"
-#include "common/ast.h"
-#include "common/decl.h"
-#include "common/expr.h"
 #include "common/function_descriptor.h"
 #include "common/kind.h"
-#include "common/type.h"
 #include "common/value.h"
-#include "compiler/compiler.h"
-#include "compiler/compiler_factory.h"
-#include "compiler/standard_library.h"
 #include "eval/compiler/cel_expression_builder_flat_impl.h"
 #include "eval/compiler/constant_folding.h"
-#include "eval/compiler/flat_expr_builder_extensions.h"
 #include "eval/compiler/qualified_reference_resolver.h"
-#include "eval/eval/evaluator_core.h"
 #include "eval/eval/expression_step_logic.h"
 #include "eval/public/activation.h"
 #include "eval/public/ast_rewrite.h"
@@ -3417,22 +3406,33 @@ TEST(FlatExprBuilderTest, ListAppendComprehensionWithTargetFailsPlanning) {
 }
 
 TEST(FlatExprBuilderTest, ProgramOptimizerResolvesTypesWhenRequested) {
-  std::shared_ptr<cel::runtime_internal::RuntimeEnv> env =
-      NewTestingRuntimeEnv();
-  ASSERT_OK_AND_ASSIGN(std::unique_ptr<cel::CompilerBuilder> compiler_builder,
-                       cel::NewCompilerBuilder(env->descriptor_pool));
-  ASSERT_THAT(compiler_builder->AddLibrary(cel::StandardCompilerLibrary()),
-              IsOk());
-  ASSERT_THAT(compiler_builder->GetCheckerBuilder().AddVariable(
-                  cel::MakeVariableDecl("x", cel::IntType())),
-              IsOk());
-  ASSERT_OK_AND_ASSIGN(std::unique_ptr<cel::Compiler> compiler,
-                       compiler_builder->Build());
-  ASSERT_OK_AND_ASSIGN(cel::ValidationResult validation_result,
-                       compiler->Compile("x == 1"));
-  ASSERT_TRUE(validation_result.IsValid());
-  ASSERT_OK_AND_ASSIGN(std::unique_ptr<cel::Ast> ast,
-                       validation_result.ReleaseAst());
+  CheckedExpr expr;
+  ASSERT_TRUE(google::protobuf::TextFormat::ParseFromString(
+      R"pb(
+        type_map {
+          key: 1
+          value { primitive: BOOL }
+        }
+        type_map {
+          key: 2
+          value { primitive: INT64 }
+        }
+        expr {
+          id: 1
+          call_expr {
+            function: "_==_"
+            args {
+              id: 2
+              ident_expr { name: "x" }
+            }
+            args {
+              id: 3
+              const_expr { int64_value: 1 }
+            }
+          }
+        }
+      )pb",
+      &expr));
 
   class TypeRecordingOptimizer : public ProgramOptimizer {
    public:
@@ -3467,7 +3467,7 @@ TEST(FlatExprBuilderTest, ProgramOptimizerResolvesTypesWhenRequested) {
 
   cel::RuntimeOptions options;
   options.enable_typed_field_access = false;
-  CelExpressionBuilderFlatImpl builder(env, options);
+  CelExpressionBuilderFlatImpl builder(NewTestingRuntimeEnv(), options);
   ASSERT_THAT(RegisterBuiltinFunctions(builder.GetRegistry()), IsOk());
 
   absl::flat_hash_map<int64_t, cel::Type> pre_visit_types;
@@ -3480,17 +3480,14 @@ TEST(FlatExprBuilderTest, ProgramOptimizerResolvesTypesWhenRequested) {
                                                         post_visit_types);
       });
 
-  ASSERT_OK_AND_ASSIGN(FlatExpression flat_expr,
-                       builder.flat_expr_builder().CreateExpressionImpl(
-                           std::move(ast), /*issues=*/nullptr));
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<CelExpression> cel_expr,
+                       builder.CreateExpression(&expr));
   EXPECT_THAT(pre_visit_types,
-              testing::UnorderedElementsAre(testing::Pair(1, cel::IntType()),
-                                            testing::Pair(2, cel::BoolType()),
-                                            testing::Pair(3, cel::IntType())));
+              testing::UnorderedElementsAre(testing::Pair(1, cel::BoolType()),
+                                            testing::Pair(2, cel::IntType())));
   EXPECT_THAT(post_visit_types,
-              testing::UnorderedElementsAre(testing::Pair(1, cel::IntType()),
-                                            testing::Pair(2, cel::BoolType()),
-                                            testing::Pair(3, cel::IntType())));
+              testing::UnorderedElementsAre(testing::Pair(1, cel::BoolType()),
+                                            testing::Pair(2, cel::IntType())));
 }
 
 }  // namespace

extensions/BUILD

@@ -329,6 +329,7 @@ cc_library(
         "//common:expr",
         "//common:function_descriptor",
         "//common:kind",
+        "//common:legacy_value",
         "//common:native_type",
         "//common:type",
         "//common:value",
@@ -337,6 +338,7 @@ cc_library(
         "//eval/eval:attribute_trail",
         "//eval/eval:evaluator_core",
         "//eval/eval:expression_step_base",
+        "//extensions/protobuf/internal:map_reflection",
         "//internal:casts",
         "//internal:number",
         "//internal:status_macros",
@@ -346,6 +348,7 @@ cc_library(
         "//runtime/internal:runtime_friend_access",
         "//runtime/internal:runtime_impl",
         "@com_google_absl//absl/algorithm:container",
+        "@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/functional:overload",
@@ -394,15 +397,19 @@ cc_test(
         "//internal:status_macros",
         "//internal:testing",
         "//parser",
+        "//runtime",
         "//runtime:activation",
         "//runtime:function_adapter",
         "//runtime:function_registry",
+        "//runtime:runtime_builder",
         "//runtime:runtime_issue",
         "//runtime:runtime_options",
+        "//runtime:standard_runtime_builder_factory",
         "//runtime:type_registry",
         "//runtime/internal:issue_collector",
         "//runtime/internal:runtime_env",
         "//runtime/internal:runtime_env_testing",
+        "//testutil:test_external_extensions_descriptor_set",
         "@com_google_absl//absl/base:nullability",
         "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/log:absl_check",