Skip to content

Commit 4c95b4f

Browse files
jnthntatumcopybara-github
authored andcommitted
Refactor short-circuit planning logic.
Split out specializations for optional or and &&/||. They are less uniform after introducing variadics. PiperOrigin-RevId: 945783246
1 parent de51eee commit 4c95b4f

1 file changed

Lines changed: 81 additions & 125 deletions

File tree

eval/compiler/flat_expr_builder.cc

Lines changed: 81 additions & 125 deletions
Original file line numberDiff line numberDiff line change
@@ -198,14 +198,7 @@ class CondVisitor {
198198
virtual void PostVisitTarget(const cel::Expr* expr) {}
199199
};
200200

201-
enum class BinaryCond {
202-
kAnd = 0,
203-
kOr,
204-
kOptionalOr,
205-
kOptionalOrValue,
206-
};
207-
208-
// Visitor managing the "&&" and "||" operatiions.
201+
// Visitor managing the "&&" and "||" (boolean logic) operations.
209202
// Implements short-circuiting if enabled.
210203
//
211204
// With short-circuiting enabled, generates a program like:
@@ -218,20 +211,41 @@ enum class BinaryCond {
218211
// | i + 3 | BooleanOperator | Op(arg1, arg2) |
219212
// | i + 4 | <rest of program> | arg1 | Op(arg1, arg2) |
220213
// +-------------+------------------------+------------------------+
221-
class BinaryCondVisitor : public CondVisitor {
214+
class LogicalCondVisitor : public CondVisitor {
222215
public:
223-
explicit BinaryCondVisitor(FlatExprVisitor* visitor, BinaryCond cond,
224-
bool short_circuiting)
225-
: visitor_(visitor), cond_(cond), short_circuiting_(short_circuiting) {}
216+
explicit LogicalCondVisitor(FlatExprVisitor* visitor, bool is_or,
217+
bool short_circuiting)
218+
: visitor_(visitor), is_or_(is_or), short_circuiting_(short_circuiting) {}
226219

227220
void PreVisit(const cel::Expr* expr) override;
228221
void PostVisitArg(int arg_num, const cel::Expr* expr) override;
229222
void PostVisit(const cel::Expr* expr) override;
223+
224+
private:
225+
FlatExprVisitor* visitor_;
226+
const bool is_or_;
227+
std::vector<Jump> jump_steps_;
228+
bool short_circuiting_;
229+
};
230+
231+
// Visitor managing optional "or" and "orValue" operations.
232+
// Implements short-circuiting if enabled.
233+
class OptionalOrCondVisitor : public CondVisitor {
234+
public:
235+
explicit OptionalOrCondVisitor(FlatExprVisitor* visitor, bool is_or_value,
236+
bool short_circuiting)
237+
: visitor_(visitor),
238+
is_or_value_(is_or_value),
239+
short_circuiting_(short_circuiting) {}
240+
241+
void PreVisit(const cel::Expr* expr) override;
242+
void PostVisitArg(int arg_num, const cel::Expr* expr) override {}
230243
void PostVisitTarget(const cel::Expr* expr) override;
244+
void PostVisit(const cel::Expr* expr) override;
231245

232246
private:
233247
FlatExprVisitor* visitor_;
234-
const BinaryCond cond_;
248+
const bool is_or_value_;
235249
std::vector<Jump> jump_steps_;
236250
bool short_circuiting_;
237251
};
@@ -997,11 +1011,11 @@ class FlatExprVisitor : public cel::AstVisitor {
9971011

9981012
std::unique_ptr<CondVisitor> cond_visitor;
9991013
if (call_expr.function() == cel::builtin::kAnd) {
1000-
cond_visitor = std::make_unique<BinaryCondVisitor>(
1001-
this, BinaryCond::kAnd, options_.short_circuiting);
1014+
cond_visitor = std::make_unique<LogicalCondVisitor>(
1015+
this, /*is_or=*/false, options_.short_circuiting);
10021016
} else if (call_expr.function() == cel::builtin::kOr) {
1003-
cond_visitor = std::make_unique<BinaryCondVisitor>(
1004-
this, BinaryCond::kOr, options_.short_circuiting);
1017+
cond_visitor = std::make_unique<LogicalCondVisitor>(
1018+
this, /*is_or=*/true, options_.short_circuiting);
10051019
} else if (call_expr.function() == cel::builtin::kTernary) {
10061020
if (options_.short_circuiting) {
10071021
cond_visitor = std::make_unique<TernaryCondVisitor>(this);
@@ -1011,13 +1025,13 @@ class FlatExprVisitor : public cel::AstVisitor {
10111025
} else if (enable_optional_types_ &&
10121026
call_expr.function() == kOptionalOrFn &&
10131027
call_expr.has_target() && call_expr.args().size() == 1) {
1014-
cond_visitor = std::make_unique<BinaryCondVisitor>(
1015-
this, BinaryCond::kOptionalOr, options_.short_circuiting);
1028+
cond_visitor = std::make_unique<OptionalOrCondVisitor>(
1029+
this, /*is_or_value=*/false, options_.short_circuiting);
10161030
} else if (enable_optional_types_ &&
10171031
call_expr.function() == kOptionalOrValueFn &&
10181032
call_expr.has_target() && call_expr.args().size() == 1) {
1019-
cond_visitor = std::make_unique<BinaryCondVisitor>(
1020-
this, BinaryCond::kOptionalOrValue, options_.short_circuiting);
1033+
cond_visitor = std::make_unique<OptionalOrCondVisitor>(
1034+
this, /*is_or_value=*/true, options_.short_circuiting);
10211035
} else if (IsBlock(&call_expr)) {
10221036
// cel.@block
10231037
if (block_.has_value()) {
@@ -2147,91 +2161,64 @@ FlatExprVisitor::HandleHeterogeneousEqualityIn(const cel::Expr& expr,
21472161
return CallHandlerResult::kIntercepted;
21482162
}
21492163

2150-
void BinaryCondVisitor::PreVisit(const cel::Expr* expr) {
2151-
switch (cond_) {
2152-
case BinaryCond::kAnd:
2153-
ABSL_FALLTHROUGH_INTENDED;
2154-
case BinaryCond::kOr:
2155-
visitor_->ValidateOrError(
2156-
!expr->call_expr().has_target() &&
2157-
expr->call_expr().args().size() >= 2,
2158-
"Invalid argument count for a binary function call.");
2159-
break;
2160-
case BinaryCond::kOptionalOr:
2161-
ABSL_FALLTHROUGH_INTENDED;
2162-
case BinaryCond::kOptionalOrValue:
2163-
visitor_->ValidateOrError(expr->call_expr().has_target() &&
2164-
expr->call_expr().args().size() == 1,
2165-
"Invalid argument count for or/orValue call.");
2166-
break;
2167-
}
2164+
void LogicalCondVisitor::PreVisit(const cel::Expr* expr) {
2165+
visitor_->ValidateOrError(
2166+
!expr->call_expr().has_target() && expr->call_expr().args().size() >= 2,
2167+
"Invalid argument count for a binary function call.");
21682168
}
21692169

2170-
void BinaryCondVisitor::PostVisitArg(int arg_num, const cel::Expr* expr) {
2170+
void LogicalCondVisitor::PostVisitArg(int arg_num, const cel::Expr* expr) {
21712171
if (visitor_->PlanRecursiveProgram()) {
21722172
return;
21732173
}
21742174
const int last_arg_index = expr->call_expr().args().size() - 1;
2175-
if (cond_ == BinaryCond::kAnd || cond_ == BinaryCond::kOr) {
2176-
if (arg_num > 0) {
2177-
switch (cond_) {
2178-
case BinaryCond::kAnd:
2179-
visitor_->AddStep(CreateAndStep(expr->id()));
2180-
break;
2181-
case BinaryCond::kOr:
2182-
visitor_->AddStep(CreateOrStep(expr->id()));
2183-
break;
2184-
default:
2185-
break;
2186-
}
2187-
if (short_circuiting_ && !jump_steps_.empty()) {
2188-
visitor_->SetProgressStatusIfError(
2189-
jump_steps_.back().set_target(visitor_->GetCurrentIndex()));
2190-
}
2175+
if (arg_num > 0) {
2176+
if (is_or_) {
2177+
visitor_->AddStep(CreateOrStep(expr->id()));
2178+
} else {
2179+
visitor_->AddStep(CreateAndStep(expr->id()));
21912180
}
2192-
if (short_circuiting_ && arg_num < last_arg_index) {
2193-
std::unique_ptr<JumpStepBase> jump_step;
2194-
switch (cond_) {
2195-
case BinaryCond::kAnd:
2196-
jump_step = CreateCondJumpStep(false, {}, expr->id());
2197-
break;
2198-
case BinaryCond::kOr:
2199-
jump_step = CreateCondJumpStep(true, {}, expr->id());
2200-
break;
2201-
default:
2202-
ABSL_UNREACHABLE();
2203-
}
2204-
ProgramStepIndex index = visitor_->GetCurrentIndex();
2205-
if (JumpStepBase* jump_step_ptr = visitor_->AddStep(std::move(jump_step));
2206-
jump_step_ptr) {
2207-
jump_steps_.push_back(Jump(index, jump_step_ptr));
2208-
}
2181+
if (short_circuiting_ && !jump_steps_.empty()) {
2182+
visitor_->SetProgressStatusIfError(
2183+
jump_steps_.back().set_target(visitor_->GetCurrentIndex()));
22092184
}
22102185
}
2186+
if (short_circuiting_ && arg_num < last_arg_index) {
2187+
std::unique_ptr<JumpStepBase> jump_step =
2188+
is_or_ ? CreateCondJumpStep(true, {}, expr->id())
2189+
: CreateCondJumpStep(false, {}, expr->id());
2190+
ProgramStepIndex index = visitor_->GetCurrentIndex();
2191+
if (JumpStepBase* jump_step_ptr = visitor_->AddStep(std::move(jump_step));
2192+
jump_step_ptr) {
2193+
jump_steps_.push_back(Jump(index, jump_step_ptr));
2194+
}
2195+
}
2196+
}
2197+
2198+
void LogicalCondVisitor::PostVisit(const cel::Expr* expr) {
2199+
if (visitor_->PlanRecursiveProgram()) {
2200+
visitor_->MakeShortcircuitRecursive(expr, is_or_);
2201+
}
2202+
}
2203+
2204+
void OptionalOrCondVisitor::PreVisit(const cel::Expr* expr) {
2205+
visitor_->ValidateOrError(
2206+
expr->call_expr().has_target() && expr->call_expr().args().size() == 1,
2207+
"Invalid argument count for or/orValue call.");
22112208
}
22122209

2213-
void BinaryCondVisitor::PostVisitTarget(const cel::Expr* expr) {
2210+
void OptionalOrCondVisitor::PostVisitTarget(const cel::Expr* expr) {
22142211
if (visitor_->PlanRecursiveProgram()) {
22152212
return;
22162213
}
2217-
if (short_circuiting_ && (cond_ == BinaryCond::kOptionalOr ||
2218-
cond_ == BinaryCond::kOptionalOrValue)) {
2214+
if (short_circuiting_) {
22192215
// If first branch evaluation result is enough to determine output,
22202216
// jump over the second branch and provide result of the first argument as
22212217
// final output.
22222218
// Retain a pointer to the jump step so we can update the target after
22232219
// planning the second argument.
2224-
std::unique_ptr<JumpStepBase> jump_step;
2225-
switch (cond_) {
2226-
case BinaryCond::kOptionalOr:
2227-
jump_step = CreateOptionalHasValueJumpStep(false, expr->id());
2228-
break;
2229-
case BinaryCond::kOptionalOrValue:
2230-
jump_step = CreateOptionalHasValueJumpStep(true, expr->id());
2231-
break;
2232-
default:
2233-
ABSL_UNREACHABLE();
2234-
}
2220+
std::unique_ptr<JumpStepBase> jump_step =
2221+
CreateOptionalHasValueJumpStep(is_or_value_, expr->id());
22352222
ProgramStepIndex index = visitor_->GetCurrentIndex();
22362223
if (JumpStepBase* jump_step_ptr = visitor_->AddStep(std::move(jump_step));
22372224
jump_step_ptr) {
@@ -2240,48 +2227,17 @@ void BinaryCondVisitor::PostVisitTarget(const cel::Expr* expr) {
22402227
}
22412228
}
22422229

2243-
void BinaryCondVisitor::PostVisit(const cel::Expr* expr) {
2230+
void OptionalOrCondVisitor::PostVisit(const cel::Expr* expr) {
22442231
if (visitor_->PlanRecursiveProgram()) {
2245-
switch (cond_) {
2246-
case BinaryCond::kAnd:
2247-
visitor_->MakeShortcircuitRecursive(expr, /*is_or=*/false);
2248-
break;
2249-
case BinaryCond::kOr:
2250-
visitor_->MakeShortcircuitRecursive(expr, /*is_or=*/true);
2251-
break;
2252-
case BinaryCond::kOptionalOr:
2253-
visitor_->MakeOptionalShortcircuit(expr,
2254-
/*is_or_value=*/false);
2255-
break;
2256-
case BinaryCond::kOptionalOrValue:
2257-
visitor_->MakeOptionalShortcircuit(expr,
2258-
/*is_or_value=*/true);
2259-
break;
2260-
default:
2261-
ABSL_UNREACHABLE();
2262-
}
2232+
visitor_->MakeOptionalShortcircuit(expr, is_or_value_);
22632233
return;
22642234
}
22652235

2266-
if (cond_ == BinaryCond::kOptionalOr ||
2267-
cond_ == BinaryCond::kOptionalOrValue) {
2268-
switch (cond_) {
2269-
case BinaryCond::kOptionalOr:
2270-
visitor_->AddStep(
2271-
CreateOptionalOrStep(/*is_or_value=*/false, expr->id()));
2272-
break;
2273-
case BinaryCond::kOptionalOrValue:
2274-
visitor_->AddStep(
2275-
CreateOptionalOrStep(/*is_or_value=*/true, expr->id()));
2276-
break;
2277-
default:
2278-
ABSL_UNREACHABLE();
2279-
}
2280-
if (short_circuiting_) {
2281-
for (auto& jump : jump_steps_) {
2282-
visitor_->SetProgressStatusIfError(
2283-
jump.set_target(visitor_->GetCurrentIndex()));
2284-
}
2236+
visitor_->AddStep(CreateOptionalOrStep(is_or_value_, expr->id()));
2237+
if (short_circuiting_) {
2238+
for (auto& jump : jump_steps_) {
2239+
visitor_->SetProgressStatusIfError(
2240+
jump.set_target(visitor_->GetCurrentIndex()));
22852241
}
22862242
}
22872243
}

0 commit comments

Comments
 (0)