diff --git a/eval/compiler/BUILD b/eval/compiler/BUILD index a8fe79b56..746b6af8f 100644 --- a/eval/compiler/BUILD +++ b/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", diff --git a/eval/compiler/flat_expr_builder.cc b/eval/compiler/flat_expr_builder.cc index 8228764a4..8e1c5606d 100644 --- a/eval/compiler/flat_expr_builder.cc +++ b/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 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(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 field_type; if (options_.enable_typed_field_access) { std::optional 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> program_optimizers_; const absl::flat_hash_map& type_map_; - absl::flat_hash_map 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) { diff --git a/eval/compiler/flat_expr_builder_extensions.h b/eval/compiler/flat_expr_builder_extensions.h index 66b80d215..90916f5b4 100644 --- a/eval/compiler/flat_expr_builder_extensions.h +++ b/eval/compiler/flat_expr_builder_extensions.h @@ -25,6 +25,7 @@ #include #include #include +#include #include #include #include @@ -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 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& arena_; - const bool explicit_arena_; const std::shared_ptr message_factory_; + const bool explicit_arena_; + bool resolve_types_; + absl::flat_hash_map resolved_types_; }; // Interface for Ast Transforms. diff --git a/eval/compiler/flat_expr_builder_extensions_test.cc b/eval/compiler/flat_expr_builder_extensions_test.cc index 07a6cfbe6..7d9426205 100644 --- a/eval/compiler/flat_expr_builder_extensions_test.cc +++ b/eval/compiler/flat_expr_builder_extensions_test.cc @@ -14,6 +14,7 @@ #include "eval/compiler/flat_expr_builder_extensions.h" #include +#include #include #include "absl/base/nullability.h" @@ -21,6 +22,7 @@ #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 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_() {} diff --git a/eval/compiler/flat_expr_builder_test.cc b/eval/compiler/flat_expr_builder_test.cc index 5ce27a603..5d36dfb37 100644 --- a/eval/compiler/flat_expr_builder_test.cc +++ b/eval/compiler/flat_expr_builder_test.cc @@ -17,6 +17,7 @@ #include #include #include +#include #include #include #include @@ -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 env = + NewTestingRuntimeEnv(); + ASSERT_OK_AND_ASSIGN(std::unique_ptr 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 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 ast, + validation_result.ReleaseAst()); + + class TypeRecordingOptimizer : public ProgramOptimizer { + public: + explicit TypeRecordingOptimizer( + absl::flat_hash_map& pre_visit_types, + absl::flat_hash_map& 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 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 type = context.GetResolvedType(node.id()); + type.has_value()) { + post_visit_types_[node.id()] = *type; + } + return absl::OkStatus(); + } + + private: + absl::flat_hash_map& pre_visit_types_; + absl::flat_hash_map& 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 pre_visit_types; + absl::flat_hash_map post_visit_types; + builder.flat_expr_builder().AddProgramOptimizer( + [&](PlannerContext& context, const cel::Ast&) + -> absl::StatusOr> { + context.set_resolve_types(); + return std::make_unique(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