Expose resolved types through PlannerContext in FlatExprBuilder.

eval/compiler/BUILD

@@ -32,6 +32,7 @@ cc_library(
         "//base:ast",
         "//base:data",
         "//common:expr",
+        "//common:type",
         "//common:value",
         "//eval/eval:evaluator_core",
         "//eval/eval:expression_step_logic",
@@ -59,6 +60,7 @@ cc_test(
         ":flat_expr_builder_extensions",
         ":resolver",
         "//common:expr",
+        "//common:type",
         "//eval/eval:evaluator_core",
         "//eval/eval:expression_step_logic",
         "//internal:testing",
@@ -146,9 +148,19 @@ cc_test(
         ":flat_expr_builder",
         ":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/compiler:flat_expr_builder_extensions",
+        "//eval/eval:evaluator_core",
         "//eval/public:activation",
         "//eval/public:ast_rewrite",
         "//eval/public:builtin_func_registrar",

eval/compiler/flat_expr_builder.cc

@@ -611,21 +611,6 @@ class FlatExprVisitor : public cel::AstVisitor {
     }
   }
 
-  void SetResolvedType(const cel::Expr& expr, cel::Type type) {
-    resolved_types_[&expr] = std::move(type);
-  }
-
-  std::optional<cel::Type> GetResolvedType(const cel::Expr* expr) const {
-    if (expr == nullptr) {
-      return std::nullopt;
-    }
-    auto it = resolved_types_.find(expr);
-    if (it != resolved_types_.end()) {
-      return it->second;
-    }
-    return std::nullopt;
-  }
-
   void PreVisitExpr(const cel::Expr& expr) override {
     ValidateOrError(!absl::holds_alternative<cel::UnspecifiedExpr>(expr.kind()),
                     "Invalid empty expression");
@@ -637,7 +622,7 @@ class FlatExprVisitor : public cel::AstVisitor {
       resume_from_suppressed_branch_ = &expr;
     }
 
-    if (options_.enable_typed_field_access) {
+    if (extension_context_.resolve_types()) {
       MaybeResolveType(expr);
     }
 
@@ -965,7 +950,7 @@ class FlatExprVisitor : public cel::AstVisitor {
     std::optional<cel::StructTypeField> field_type;
     if (options_.enable_typed_field_access) {
       std::optional<cel::Type> operand_type =
-          GetResolvedType(&select_expr.operand());
+          extension_context_.GetResolvedType(select_expr.operand().id());
       if (operand_type.has_value() && operand_type->IsStruct()) {
         struct_type = operand_type->GetStruct();
         if (struct_type.has_value()) {
@@ -1720,7 +1705,6 @@ class FlatExprVisitor : public cel::AstVisitor {
   const cel::Expr* resume_from_suppressed_branch_ = nullptr;
   std::vector<std::unique_ptr<ProgramOptimizer>> program_optimizers_;
   const absl::flat_hash_map<int64_t, cel::TypeSpec>& type_map_;
-  absl::flat_hash_map<const cel::Expr*, cel::Type> resolved_types_;
   IssueCollector& issue_collector_;
 
   ProgramBuilder& program_builder_;
@@ -1869,7 +1853,7 @@ void FlatExprVisitor::MaybeResolveType(const cel::Expr& expr) {
   if (!type.ok()) {
     return;
   }
-  SetResolvedType(expr, *type);
+  extension_context_.SetResolvedType(expr.id(), *type);
 }
 
 void LogicalCondVisitor::PreVisit(const cel::Expr* expr) {

eval/compiler/flat_expr_builder_extensions.h

@@ -25,6 +25,7 @@
 #include <cstddef>
 #include <cstdint>
 #include <memory>
+#include <optional>
 #include <utility>
 #include <variant>
 #include <vector>
@@ -40,6 +41,7 @@
 #include "base/ast.h"
 #include "base/type_provider.h"
 #include "common/expr.h"
+#include "common/type.h"
 #include "common/type_reflector.h"
 #include "eval/compiler/resolver.h"
 #include "eval/eval/evaluator_core.h"
@@ -297,8 +299,9 @@ class PlannerContext {
         issue_collector_(issue_collector),
         program_builder_(program_builder),
         arena_(arena),
+        message_factory_(std::move(message_factory)),
         explicit_arena_(arena_ != nullptr),
-        message_factory_(std::move(message_factory)) {}
+        resolve_types_(options_.enable_typed_field_access) {}
 
   ProgramBuilder& program_builder() { return program_builder_; }
 
@@ -358,6 +361,30 @@ class PlannerContext {
     return arena_.get();
   }
 
+  // Whether the planner should resolve types in the AST to concrete cel::Type
+  // values.
+  bool resolve_types() const { return resolve_types_; }
+
+  void set_resolve_types() { resolve_types_ = true; }
+
+  // Returns the resolved type if available.
+  //
+  // This can only be used as a hint since we can't guarantee that the AST was
+  // type checked or can be fully trusted (e.g. post-processing broke IDs).
+  //
+  // This is populated by FlatExprBuilder on the PreVisit pass for the Expr.
+  std::optional<cel::Type> GetResolvedType(int64_t expr_id) const {
+    auto it = resolved_types_.find(expr_id);
+    if (it != resolved_types_.end()) {
+      return it->second;
+    }
+    return std::nullopt;
+  }
+
+  void SetResolvedType(int64_t expr_id, cel::Type type) {
+    resolved_types_[expr_id] = std::move(type);
+  }
+
   // Returns `true` if a message factory was explicitly provided during
   // planning.
   bool HasExplicitMessageFactory() const { return message_factory_ != nullptr; }
@@ -375,8 +402,10 @@ class PlannerContext {
   cel::runtime_internal::IssueCollector& issue_collector_;
   ProgramBuilder& program_builder_;
   std::shared_ptr<google::protobuf::Arena>& arena_;
-  const bool explicit_arena_;
   const std::shared_ptr<google::protobuf::MessageFactory> message_factory_;
+  const bool explicit_arena_;
+  bool resolve_types_;
+  absl::flat_hash_map<int64_t, cel::Type> resolved_types_;
 };
 
 // Interface for Ast Transforms.

eval/compiler/flat_expr_builder_extensions_test.cc

@@ -14,13 +14,15 @@
 #include "eval/compiler/flat_expr_builder_extensions.h"
 
 #include <memory>
+#include <optional>
 #include <utility>
 
 #include "absl/base/nullability.h"
 #include "absl/status/status.h"
 #include "absl/status/status_matchers.h"
 #include "absl/status/statusor.h"
 #include "common/expr.h"
+#include "common/type.h"
 #include "eval/compiler/resolver.h"
 #include "eval/eval/evaluator_core.h"
 #include "eval/eval/expression_step_logic.h"
@@ -347,6 +349,23 @@ TEST_F(PlannerContextTest, AddSubplanStepFailsOnUnknownNode) {
               StatusIs(absl::StatusCode::kInternal));
 }
 
+TEST_F(PlannerContextTest, ResolvedTypes) {
+  ProgramBuilder program_builder;
+  std::shared_ptr<google::protobuf::Arena> arena;
+  options_.enable_typed_field_access = false;
+  PlannerContext context(env_, resolver_, options_,
+                         type_registry_.GetComposedTypeProvider(),
+                         issue_collector_, program_builder, arena);
+
+  EXPECT_FALSE(context.resolve_types());
+  context.set_resolve_types();
+  EXPECT_TRUE(context.resolve_types());
+
+  EXPECT_EQ(context.GetResolvedType(1), std::nullopt);
+  context.SetResolvedType(1, cel::IntType());
+  EXPECT_THAT(context.GetResolvedType(1), testing::Optional(cel::IntType()));
+}
+
 class ProgramBuilderTest : public testing::Test {
  public:
   ProgramBuilderTest() : type_registry_(), function_registry_() {}

eval/compiler/flat_expr_builder_test.cc

@@ -17,6 +17,7 @@
 #include <cstdint>
 #include <functional>
 #include <memory>
+#include <optional>
 #include <string>
 #include <utility>
 #include <vector>
@@ -35,12 +36,22 @@
 #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/public/activation.h"
 #include "eval/public/ast_rewrite.h"
 #include "eval/public/builtin_func_registrar.h"
@@ -3148,6 +3159,83 @@ TEST(FlatExprBuilderTest, ListAppendComprehensionWithTargetFailsPlanning) {
       StatusIs(absl::StatusCode::kInvalidArgument));
 }
 
+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());
+
+  class TypeRecordingOptimizer : public ProgramOptimizer {
+   public:
+    explicit TypeRecordingOptimizer(
+        absl::flat_hash_map<int64_t, cel::Type>& pre_visit_types,
+        absl::flat_hash_map<int64_t, cel::Type>& post_visit_types)
+        : pre_visit_types_(pre_visit_types),
+          post_visit_types_(post_visit_types) {}
+
+    absl::Status OnPreVisit(PlannerContext& context,
+                            const cel::Expr& node) override {
+      if (std::optional<cel::Type> type = context.GetResolvedType(node.id());
+          type.has_value()) {
+        pre_visit_types_[node.id()] = *type;
+      }
+      return absl::OkStatus();
+    }
+
+    absl::Status OnPostVisit(PlannerContext& context,
+                             const cel::Expr& node) override {
+      if (std::optional<cel::Type> type = context.GetResolvedType(node.id());
+          type.has_value()) {
+        post_visit_types_[node.id()] = *type;
+      }
+      return absl::OkStatus();
+    }
+
+   private:
+    absl::flat_hash_map<int64_t, cel::Type>& pre_visit_types_;
+    absl::flat_hash_map<int64_t, cel::Type>& post_visit_types_;
+  };
+
+  cel::RuntimeOptions options;
+  options.enable_typed_field_access = false;
+  CelExpressionBuilderFlatImpl builder(env, options);
+  ASSERT_THAT(RegisterBuiltinFunctions(builder.GetRegistry()), IsOk());
+
+  absl::flat_hash_map<int64_t, cel::Type> pre_visit_types;
+  absl::flat_hash_map<int64_t, cel::Type> post_visit_types;
+  builder.flat_expr_builder().AddProgramOptimizer(
+      [&](PlannerContext& context, const cel::Ast&)
+          -> absl::StatusOr<std::unique_ptr<ProgramOptimizer>> {
+        context.set_resolve_types();
+        return std::make_unique<TypeRecordingOptimizer>(pre_visit_types,
+                                                        post_visit_types);
+      });
+
+  ASSERT_OK_AND_ASSIGN(FlatExpression flat_expr,
+                       builder.flat_expr_builder().CreateExpressionImpl(
+                           std::move(ast), /*issues=*/nullptr));
+  EXPECT_THAT(pre_visit_types,
+              testing::UnorderedElementsAre(testing::Pair(1, cel::IntType()),
+                                            testing::Pair(2, cel::BoolType()),
+                                            testing::Pair(3, cel::IntType())));
+  EXPECT_THAT(post_visit_types,
+              testing::UnorderedElementsAre(testing::Pair(1, cel::IntType()),
+                                            testing::Pair(2, cel::BoolType()),
+                                            testing::Pair(3, cel::IntType())));
+}
+
 }  // namespace
 
 }  // namespace google::api::expr::runtime