Skip to content

Commit bc1e3e6

Browse files
jnthntatumcopybara-github
jnthntatum
authored and
copybara-github
committed
Add support for resolving types at plan time.
Add an option to prefetch field descriptors when type is known. PiperOrigin-RevId: 953496164
1 parent 6cb7d12 commit bc1e3e6

17 files changed

Lines changed: 554 additions & 123 deletions

‎conformance/BUILD‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -64,6 +64,7 @@ cc_library(
6464
"//runtime:constant_folding",
6565
"//runtime:optional_types",
6666
"//runtime:reference_resolver",
67+
"//runtime:regex_precompilation",
6768
"//runtime:runtime_options",
6869
"//runtime:standard_runtime_builder_factory",
6970
"//testutil:test_macros",

‎conformance/service.cc‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -76,6 +76,7 @@
7676
#include "runtime/constant_folding.h"
7777
#include "runtime/optional_types.h"
7878
#include "runtime/reference_resolver.h"
79+
#include "runtime/regex_precompilation.h"
7980
#include "runtime/runtime.h"
8081
#include "runtime/runtime_options.h"
8182
#include "runtime/standard_runtime_builder_factory.h"
@@ -283,6 +284,7 @@ class LegacyConformanceServiceImpl : public ConformanceServiceInterface {
283284
std::cerr << "Enabling optimizations" << std::endl;
284285
options.constant_folding = true;
285286
options.constant_arena = constant_arena;
287+
options.enable_typed_field_access = true;
286288
}
287289

288290
if (select_optimization) {
@@ -486,13 +488,17 @@ class ModernConformanceServiceImpl : public ConformanceServiceInterface {
486488
absl::string_view container) {
487489
RuntimeOptions options(options_);
488490
options.container = std::string(container);
491+
if (enable_optimizations_) {
492+
options.enable_typed_field_access = true;
493+
}
489494
CEL_ASSIGN_OR_RETURN(
490495
auto builder, CreateStandardRuntimeBuilder(
491496
google::protobuf::DescriptorPool::generated_pool(), options));
492497

493498
if (enable_optimizations_) {
494499
CEL_RETURN_IF_ERROR(cel::extensions::EnableConstantFolding(
495500
builder, google::protobuf::MessageFactory::generated_factory()));
501+
CEL_RETURN_IF_ERROR(cel::extensions::EnableRegexPrecompilation(builder));
496502
}
497503
CEL_RETURN_IF_ERROR(cel::EnableReferenceResolver(
498504
builder, cel::ReferenceResolverEnabled::kAlways));

‎eval/compiler/BUILD‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -109,6 +109,7 @@ cc_library(
109109
"//common:expr",
110110
"//common:kind",
111111
"//common:type",
112+
"//common:type_spec_resolver",
112113
"//common:value",
113114
"//eval/eval:comprehension_step",
114115
"//eval/eval:const_value_step",

‎eval/compiler/flat_expr_builder.cc‎

Lines changed: 71 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -23,15 +23,14 @@
2323
#include <iterator>
2424
#include <limits>
2525
#include <memory>
26+
#include <optional>
2627
#include <stack>
2728
#include <string>
2829
#include <type_traits>
2930
#include <utility>
3031
#include <vector>
3132

3233
#include "absl/algorithm/container.h"
33-
#include "absl/base/attributes.h"
34-
#include "absl/base/optimization.h"
3534
#include "absl/container/flat_hash_map.h"
3635
#include "absl/container/flat_hash_set.h"
3736
#include "absl/container/node_hash_map.h"
@@ -45,7 +44,6 @@
4544
#include "absl/strings/str_cat.h"
4645
#include "absl/strings/string_view.h"
4746
#include "absl/strings/strip.h"
48-
#include "absl/types/optional.h"
4947
#include "absl/types/span.h"
5048
#include "absl/types/variant.h"
5149
#include "base/ast.h"
@@ -59,6 +57,7 @@
5957
#include "common/expr.h"
6058
#include "common/kind.h"
6159
#include "common/type.h"
60+
#include "common/type_spec_resolver.h"
6261
#include "common/value.h"
6362
#include "eval/compiler/check_ast_extensions.h"
6463
#include "eval/compiler/flat_expr_builder_extensions.h"
@@ -528,7 +527,7 @@ class FlatExprVisitor : public cel::AstVisitor {
528527
FlatExprVisitor(
529528
const Resolver& resolver, const cel::RuntimeOptions& options,
530529
std::vector<std::unique_ptr<ProgramOptimizer>> program_optimizers,
531-
const absl::flat_hash_map<int64_t, cel::Reference>& reference_map,
530+
const absl::flat_hash_map<int64_t, cel::TypeSpec>& type_map,
532531
const cel::TypeProvider& type_provider, IssueCollector& issue_collector,
533532
ProgramBuilder& program_builder, PlannerContext& extension_context,
534533
bool enable_optional_types)
@@ -538,6 +537,7 @@ class FlatExprVisitor : public cel::AstVisitor {
538537
resolved_select_expr_(nullptr),
539538
options_(options),
540539
program_optimizers_(std::move(program_optimizers)),
540+
type_map_(type_map),
541541
issue_collector_(issue_collector),
542542
program_builder_(program_builder),
543543
extension_context_(extension_context),
@@ -606,6 +606,21 @@ class FlatExprVisitor : public cel::AstVisitor {
606606

607607
bool PlanRecursiveProgram() const { return max_recursion_depth_ > 0; }
608608

609+
void SetResolvedType(const cel::Expr& expr, cel::Type type) {
610+
resolved_types_[&expr] = std::move(type);
611+
}
612+
613+
std::optional<cel::Type> GetResolvedType(const cel::Expr* expr) const {
614+
if (expr == nullptr) {
615+
return std::nullopt;
616+
}
617+
auto it = resolved_types_.find(expr);
618+
if (it != resolved_types_.end()) {
619+
return it->second;
620+
}
621+
return std::nullopt;
622+
}
623+
609624
void PreVisitExpr(const cel::Expr& expr) override {
610625
ValidateOrError(!absl::holds_alternative<cel::UnspecifiedExpr>(expr.kind()),
611626
"Invalid empty expression");
@@ -617,6 +632,10 @@ class FlatExprVisitor : public cel::AstVisitor {
617632
resume_from_suppressed_branch_ = &expr;
618633
}
619634

635+
if (options_.enable_typed_field_access) {
636+
MaybeResolveType(expr);
637+
}
638+
620639
if (block_.has_value()) {
621640
BlockInfo& block = *block_;
622641
if (block.in && block.bindings_set.contains(&expr)) {
@@ -977,6 +996,24 @@ class FlatExprVisitor : public cel::AstVisitor {
977996
}
978997

979998
StringValue field = cel::StringValue(select_expr.field());
999+
std::optional<cel::StructType> struct_type;
1000+
std::optional<cel::StructTypeField> field_type;
1001+
if (options_.enable_typed_field_access) {
1002+
std::optional<cel::Type> operand_type =
1003+
GetResolvedType(&select_expr.operand());
1004+
if (operand_type.has_value() && operand_type->IsStruct()) {
1005+
struct_type = operand_type->GetStruct();
1006+
if (struct_type.has_value()) {
1007+
auto field_lookup =
1008+
extension_context_.type_reflector().FindStructTypeFieldByName(
1009+
*struct_type, select_expr.field());
1010+
// Swallow error to fallback to duck typing behavior.
1011+
if (field_lookup.ok() && field_lookup->has_value()) {
1012+
field_type = *std::move(field_lookup);
1013+
}
1014+
}
1015+
}
1016+
}
9801017
if (auto depth = RecursionEligible(); depth.has_value()) {
9811018
auto deps = ExtractRecursiveDependencies();
9821019
if (deps.size() != 1) {
@@ -994,6 +1031,13 @@ class FlatExprVisitor : public cel::AstVisitor {
9941031
return;
9951032
}
9961033

1034+
if (field_type.has_value()) {
1035+
AddStep(CreateTypedSelectStep(
1036+
std::move(field), *struct_type, *std::move(field_type),
1037+
select_expr.test_only(), expr.id(),
1038+
options_.enable_empty_wrapper_null_unboxing, enable_optional_types_));
1039+
return;
1040+
}
9971041
AddStep(CreateSelectStep(
9981042
std::move(field), select_expr.test_only(), expr.id(),
9991043
options_.enable_empty_wrapper_null_unboxing, enable_optional_types_));
@@ -1921,6 +1965,8 @@ class FlatExprVisitor : public cel::AstVisitor {
19211965
CallHandlerResult HandleHeterogeneousEqualityIn(const cel::Expr& expr,
19221966
const cel::CallExpr& call);
19231967

1968+
void MaybeResolveType(const cel::Expr& expr);
1969+
19241970
const Resolver& resolver_;
19251971
const cel::TypeProvider& type_provider_;
19261972
absl::Status progress_status_;
@@ -1942,6 +1988,8 @@ class FlatExprVisitor : public cel::AstVisitor {
19421988
absl::flat_hash_set<const cel::Expr*> suppressed_branches_;
19431989
const cel::Expr* resume_from_suppressed_branch_ = nullptr;
19441990
std::vector<std::unique_ptr<ProgramOptimizer>> program_optimizers_;
1991+
const absl::flat_hash_map<int64_t, cel::TypeSpec>& type_map_;
1992+
absl::flat_hash_map<const cel::Expr*, cel::Type> resolved_types_;
19451993
IssueCollector& issue_collector_;
19461994

19471995
ProgramBuilder& program_builder_;
@@ -2161,6 +2209,23 @@ FlatExprVisitor::HandleHeterogeneousEqualityIn(const cel::Expr& expr,
21612209
return CallHandlerResult::kIntercepted;
21622210
}
21632211

2212+
void FlatExprVisitor::MaybeResolveType(const cel::Expr& expr) {
2213+
// Try to resolve the type from the type map, but don't fail if it's not
2214+
// there. This permits cases where the runtime type is compatible but not
2215+
// the same as the type checked type.
2216+
auto it = type_map_.find(expr.id());
2217+
if (it == type_map_.end()) {
2218+
return;
2219+
}
2220+
absl::StatusOr<cel::Type> type = cel::ConvertTypeSpecToType(
2221+
it->second, extension_context_.type_reflector(),
2222+
extension_context_.MutableArena());
2223+
if (!type.ok()) {
2224+
return;
2225+
}
2226+
SetResolvedType(expr, *type);
2227+
}
2228+
21642229
void LogicalCondVisitor::PreVisit(const cel::Expr* expr) {
21652230
visitor_->ValidateOrError(
21662231
!expr->call_expr().has_target() && expr->call_expr().args().size() >= 2,
@@ -2561,8 +2626,8 @@ absl::StatusOr<FlatExpression> FlatExprBuilder::CreateExpressionImpl(
25612626
// These objects are expected to remain scoped to one build call -- references
25622627
// to them shouldn't be persisted in any part of the result expression.
25632628
FlatExprVisitor visitor(resolver, options_, std::move(optimizers),
2564-
ast->reference_map(), GetTypeProvider(),
2565-
issue_collector, program_builder, extension_context,
2629+
ast->type_map(), GetTypeProvider(), issue_collector,
2630+
program_builder, extension_context,
25662631
enable_optional_types_);
25672632

25682633
if (options_.max_recursion_depth == -1 || options_.max_recursion_depth > 0) {

‎eval/eval/BUILD‎

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -314,13 +314,11 @@ cc_library(
314314
":direct_expression_step",
315315
":evaluator_core",
316316
":expression_step_base",
317-
"//common:expr",
317+
"//common:type",
318318
"//common:value",
319319
"//common:value_kind",
320-
"//eval/internal:errors",
321320
"//internal:status_macros",
322321
"//runtime:runtime_options",
323-
"@com_google_absl//absl/base:nullability",
324322
"@com_google_absl//absl/log:absl_check",
325323
"@com_google_absl//absl/log:absl_log",
326324
"@com_google_absl//absl/status",
@@ -791,7 +789,9 @@ cc_test(
791789
deps = [
792790
":attribute_trail",
793791
":cel_expression_flat_impl",
792+
":compiler_constant_step",
794793
":const_value_step",
794+
":create_map_step",
795795
":evaluator_core",
796796
":ident_step",
797797
":select_step",
@@ -800,11 +800,14 @@ cc_test(
800800
"//common:casting",
801801
"//common:expr",
802802
"//common:legacy_value",
803+
"//common:type",
803804
"//common:value",
804805
"//common:value_testing",
805806
"//eval/public:activation",
806807
"//eval/public:cel_attribute",
807808
"//eval/public:cel_value",
809+
"//eval/public:unknown_attribute_set",
810+
"//eval/public:unknown_set",
808811
"//eval/public/containers:container_backed_map_impl",
809812
"//eval/public/structs:cel_proto_wrapper",
810813
"//eval/public/structs:legacy_type_adapter",
@@ -831,6 +834,7 @@ cc_test(
831834
"@com_google_absl//absl/strings",
832835
"@com_google_cel_spec//proto/cel/expr:syntax_cc_proto",
833836
"@com_google_cel_spec//proto/cel/expr/conformance/proto3:test_all_types_cc_proto",
837+
"@com_google_protobuf//:protobuf",
834838
"@com_google_protobuf//:wrappers_cc_proto",
835839
],
836840
)

0 commit comments

Comments
 (0)