Skip to content

Map root identifiers to slots and memoize. #2367

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
copybara-service wants to merge 1 commit into
base: master
Choose a base branch
from
Open

Map root identifiers to slots and memoize. #2367

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 13 additions & 0 deletions eval/compiler/BUILD
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -149,8 +149,11 @@ cc_test(
"//common:function_descriptor",
"//common:kind",
"//common:value",
"//eval/eval:cel_expression_flat_impl",
"//eval/eval:evaluator_core",
"//eval/public:activation",
"//eval/public:ast_rewrite",
"//eval/public:base_activation",
"//eval/public:builtin_func_registrar",
"//eval/public:cel_attribute",
"//eval/public:cel_expr_builder_factory",
Expand Down Expand Up @@ -204,6 +207,11 @@ cc_test(
":cel_expression_builder_flat_impl",
":comprehension_vulnerability_check",
":flat_expr_builder",
"//base:ast",
"//common:source",
"//common:value",
"//common:value_testing",
"//eval/eval:evaluator_core",
"//eval/public:activation",
"//eval/public:builtin_func_registrar",
"//eval/public:cel_attribute",
Expand All @@ -213,10 +221,15 @@ cc_test(
"//eval/public/containers:container_backed_list_impl",
"//eval/public/testing:matchers",
"//eval/testutil:test_message_cc_proto",
"//extensions:bindings_ext",
"//internal:testing",
"//parser",
"//parser:parser_interface",
"//runtime:activation",
"//runtime:runtime_options",
"//runtime/internal:runtime_env",
"//runtime/internal:runtime_env_testing",
"//testutil:test_macros",
"@com_google_absl//absl/status",
"@com_google_absl//absl/strings",
"@com_google_cel_spec//proto/cel/expr:syntax_cc_proto",
Expand Down
9 changes: 5 additions & 4 deletions eval/compiler/constant_folding.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -75,8 +75,9 @@ class ConstantFoldingExtension : public ProgramOptimizer {
const TypeProvider& type_provider)
: shared_arena_(std::move(shared_arena)),
shared_message_factory_(std::move(shared_message_factory)),
state_(kDefaultStackLimit, kComprehensionSlotCount, type_provider,
descriptor_pool, message_factory, arena) {}
state_(kDefaultStackLimit, kComprehensionSlotCount,
/*iterator_stack_size=*/0, type_provider, descriptor_pool,
message_factory, arena) {}

absl::Status OnPreVisit(google::api::expr::runtime::PlannerContext& context,
const Expr& node) override;
Expand Down Expand Up @@ -200,18 +201,18 @@ absl::Status ConstantFoldingExtension::OnPostVisit(PlannerContext& context,
// This subexpression is already optimized out or suppressed.
return absl::OkStatus();
}
// copy string to managed handle if backed by the original program.
// Ensure that string constants are copied to the constant folding arena
Value value;
if (node.has_const_expr()) {
CEL_ASSIGN_OR_RETURN(value,
ConvertConstant(node.const_expr(), state_.arena()));
} else {
ExecutionFrame frame(subplan, empty_, context.options(), state_);
state_.Reset();
// Update stack size to accommodate sub expression.
// This only results in a vector resize if the new maxsize is greater than
// the current capacity.
state_.value_stack().SetMaxSize(subplan.size());
ExecutionFrame frame(subplan, empty_, context.options(), state_);

auto result = frame.Evaluate();
// If this would be a runtime error, then don't adjust the program plan, but
Expand Down
100 changes: 95 additions & 5 deletions eval/compiler/flat_expr_builder.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -70,6 +70,7 @@
#include "eval/eval/evaluator_core.h"
#include "eval/eval/expression_step_logic.h"
#include "eval/eval/function_step.h"
#include "eval/eval/ident_step.h"
#include "eval/eval/optional_or_step.h"
#include "eval/eval/select_step.h"
#include "eval/eval/shadowable_value_step.h"
Expand Down Expand Up @@ -105,7 +106,12 @@ class FlatExprVisitor;
// Helper for bookkeeping variables mapped to indexes.
class IndexManager {
public:
IndexManager() : next_free_slot_(0), max_slot_count_(0) {}
IndexManager()
: next_free_slot_(0),
max_slot_count_(0),
iter_stack_count_(0),
max_iter_stack_count_(0),
lazy_iter_stack_count_(0) {}

size_t ReserveSlots(size_t n) {
size_t result = next_free_slot_;
Expand All @@ -121,11 +127,30 @@ class IndexManager {
return next_free_slot_;
}

void ReserveIterStack() {
++iter_stack_count_;
max_iter_stack_count_ = std::max(max_iter_stack_count_, iter_stack_count_);
}
void ReleaseIterStack() { --iter_stack_count_; }
void AddLazyIterStack() { ++lazy_iter_stack_count_; }

size_t max_slot_count() const { return max_slot_count_; }

// Returns the maximum number of iterator stack slots required to evaluate the
// expression by tracking the maximum comprehension nesting in the main
// expression and dedicating a slot for any comprehensions evaluated lazily.
// This is not as tight as possible, but is a correct upper bound estimate and
// is tight for small expressions.
size_t IterStackRequirement() const {
return max_iter_stack_count_ + lazy_iter_stack_count_;
}

private:
size_t next_free_slot_;
size_t max_slot_count_;
size_t iter_stack_count_;
size_t max_iter_stack_count_;
size_t lazy_iter_stack_count_;
};

// Helper for computing jump offsets.
Expand Down Expand Up @@ -868,8 +893,12 @@ class FlatExprVisitor : public cel::AstVisitor {
return;
}

if (PlanningSuppressed()) {
return;
}
absl::string_view ident_name = absl::StripPrefix(ident_expr.name(), ".");
AddStep(ExpressionStep::MakeIdentifierStep(ident_name, expr.id()));
size_t ident_index = AddOrUpdateIdentifier(ident_name);
AddStep(ExpressionStep::MakeIdentifierStep(ident_index, expr.id()));
}

void PreVisitSelect(const cel::Expr& expr,
Expand Down Expand Up @@ -1145,21 +1174,36 @@ class FlatExprVisitor : public cel::AstVisitor {
slot_count = 0;
}
}

// If this is in the scope of an optimized bind accu-init, account the slots
// to the outermost bind-init scope.
//
// The init expression is effectively inlined at the first usage in the
// critical path (which is unknown at plan time), so the used slots need to
// be dedicated for the entire scope of that bind.
bool in_bind_init = false;
for (ComprehensionStackRecord& record : comprehension_stack_) {
if (record.in_accu_init && record.is_optimizable_bind) {
record.slot_count += slot_count;
slot_count = 0;
in_bind_init = true;
break;
}
// If no bind init subexpression, account normally.
}

// Iterator frames are releasable if not in a bind (no iterator push), and
// not in a lazy subexpression (which may evaluate in different implicit
// scopes).
const bool is_stack_releasable =
!is_bind && !in_bind_init && (!block_.has_value() || !block_->in);

if (is_stack_releasable) {
index_manager_.ReserveIterStack();
} else if (!is_bind) {
index_manager_.AddLazyIterStack();
}

comprehension_stack_.push_back(
{&expr, &comprehension, iter_slot, iter2_slot, accu_slot, slot_count,
/*subexpression=*/-1,
Expand All @@ -1172,6 +1216,7 @@ class FlatExprVisitor : public cel::AstVisitor {
options_.enable_comprehension_mutable_map,
options_.short_circuiting),
/*.is_optimizable_bind=*/is_bind,
/*.is_stack_releasable=*/is_stack_releasable,
/*.iter_var_in_scope=*/false,
/*.iter_var2_in_scope=*/false,
/*.accu_var_in_scope=*/false,
Expand All @@ -1198,6 +1243,10 @@ class FlatExprVisitor : public cel::AstVisitor {
record.visitor->PostVisit(&expr);

index_manager_.ReleaseSlots(record.slot_count);
if (record.is_stack_releasable) {
ABSL_DCHECK(!record.is_optimizable_bind);
index_manager_.ReleaseIterStack();
}
comprehension_stack_.pop_back();
}

Expand Down Expand Up @@ -1499,6 +1548,19 @@ class FlatExprVisitor : public cel::AstVisitor {

size_t slot_count() const { return index_manager_.max_slot_count(); }

size_t iter_stack_count() const {
return index_manager_.IterStackRequirement();
}

IdentifierTable BuildIdentifierTable(size_t& total_slot_count) && {
IdentifierTable table;
table.reserve(identifiers_.size());
for (std::string& name : identifiers_) {
table.push_back(Identifier{std::move(name), total_slot_count++});
}
return table;
}

void AddOptimizer(std::unique_ptr<ProgramOptimizer> optimizer) {
program_optimizers_.push_back(std::move(optimizer));
}
Expand Down Expand Up @@ -1530,6 +1592,7 @@ class FlatExprVisitor : public cel::AstVisitor {
bool is_optimizable_list_append;
bool is_optimizable_map_insert;
bool is_optimizable_bind;
bool is_stack_releasable;
bool iter_var_in_scope;
bool iter_var2_in_scope;
bool accu_var_in_scope;
Expand Down Expand Up @@ -1649,6 +1712,8 @@ class FlatExprVisitor : public cel::AstVisitor {

void MaybeResolveType(const cel::Expr& expr);

size_t AddOrUpdateIdentifier(absl::string_view name);

const Resolver& resolver_;
const cel::TypeProvider& type_provider_;
absl::Status progress_status_;
Expand Down Expand Up @@ -1677,6 +1742,8 @@ class FlatExprVisitor : public cel::AstVisitor {
ProgramBuilder& program_builder_;
PlannerContext& extension_context_;
IndexManager index_manager_;
std::vector<std::string> identifiers_;
absl::flat_hash_map<std::string, size_t> identifier_map_;

bool enable_optional_types_;
std::optional<FlatExprVisitor::BlockInfo> block_;
Expand Down Expand Up @@ -1823,6 +1890,16 @@ void FlatExprVisitor::MaybeResolveType(const cel::Expr& expr) {
SetResolvedType(expr, *type);
}

size_t FlatExprVisitor::AddOrUpdateIdentifier(absl::string_view name) {
auto it = identifier_map_.find(name);
if (it == identifier_map_.end()) {
size_t index = identifiers_.size();
it = identifier_map_.emplace(name, index).first;
identifiers_.push_back(std::string(name));
}
return it->second;
}

void LogicalCondVisitor::PreVisit(const cel::Expr* expr) {
visitor_->ValidateOrError(
!expr->call_expr().has_target() && expr->call_expr().args().size() >= 2,
Expand Down Expand Up @@ -2280,7 +2357,12 @@ absl::StatusOr<FlatExpression> FlatExprBuilder::CreateExpressionImpl(
return visitor.progress_status();
}

if (visitor.slot_count() > std::numeric_limits<uint32_t>::max() ||
size_t total_slot_count = visitor.slot_count();
size_t iterator_stack_size = visitor.iter_stack_count();
IdentifierTable identifiers =
std::move(visitor).BuildIdentifierTable(total_slot_count);

if (total_slot_count > std::numeric_limits<uint32_t>::max() ||
program_builder.ExtractedSubexpressionCount() >
std::numeric_limits<uint32_t>::max()) {
// Impractical to trigger (we'd run out of memory first), but assuming this
Expand All @@ -2296,10 +2378,18 @@ absl::StatusOr<FlatExpression> FlatExprBuilder::CreateExpressionImpl(
ExecutionPath execution_path;
std::vector<ExecutionPathView> subexpressions =
FlattenExpressionTable(program_builder, execution_path);
// Iterate over execution path and update identifier steps with resolved slot.
for (ExpressionStep& step : execution_path) {
if (IdentStepInfo* ident_step = GetIfIdentStep(step);
ident_step != nullptr) {
ident_step->slot_index = identifiers[ident_step->ident_index].slot_index;
}
}

return FlatExpression(std::move(execution_path), std::move(subexpressions),
visitor.slot_count(), GetTypeProvider(), options_,
std::move(arena));
total_slot_count, iterator_stack_size,
GetTypeProvider(), options_, std::move(arena),
std::move(identifiers));
}
const cel::TypeProvider& FlatExprBuilder::GetTypeProvider() const {
return use_legacy_type_provider_
Expand Down
Loading
Loading