Skip to content

Commit c4f9b58

Browse files
jnthntatumcopybara-github
authored andcommitted
Fix for slot calculation for block expressions.
With this change, all comprehension variables inside cel.@block are assigned to a dedicated slot. It's hard to identify when they can be safely reused with lazy evaluation and likely doesn't provide enough benefit to try to reuse them more aggressively. PiperOrigin-RevId: 694717283
1 parent 2f43982 commit c4f9b58

File tree

2 files changed

+15
-3
lines changed

2 files changed

+15
-3
lines changed

eval/compiler/flat_expr_builder.cc

Lines changed: 14 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -886,6 +886,7 @@ class FlatExprVisitor : public cel::AstVisitor {
886886
block.bindings_set.insert(&list_expr_element.expr());
887887
}
888888
block.index = index_manager().ReserveSlots(block.size);
889+
block.slot_count = block.size;
889890
block.expr = &expr;
890891
block.bindings = &call_expr.args()[0];
891892
block.bound = &call_expr.args()[1];
@@ -1203,8 +1204,8 @@ class FlatExprVisitor : public cel::AstVisitor {
12031204
BlockInfo& block = *block_;
12041205
if (block.expr == &expr) {
12051206
block.in = false;
1206-
index_manager().ReleaseSlots(block.size);
1207-
AddStep(CreateClearSlotsStep(block.index, block.size, -1));
1207+
index_manager().ReleaseSlots(block.slot_count);
1208+
AddStep(CreateClearSlotsStep(block.index, block.slot_count, -1));
12081209
return;
12091210
}
12101211
}
@@ -1269,6 +1270,7 @@ class FlatExprVisitor : public cel::AstVisitor {
12691270

12701271
size_t iter_slot, accu_slot, slot_count;
12711272
bool is_bind = IsBind(&comprehension);
1273+
12721274
if (is_bind) {
12731275
accu_slot = iter_slot = index_manager_.ReserveSlots(1);
12741276
slot_count = 1;
@@ -1277,6 +1279,14 @@ class FlatExprVisitor : public cel::AstVisitor {
12771279
accu_slot = iter_slot + 1;
12781280
slot_count = 2;
12791281
}
1282+
1283+
if (block_.has_value()) {
1284+
BlockInfo& block = *block_;
1285+
if (block.in) {
1286+
block.slot_count += slot_count;
1287+
slot_count = 0;
1288+
}
1289+
}
12801290
// If this is in the scope of an optimized bind accu-init, account the slots
12811291
// to the outermost bind-init scope.
12821292
//
@@ -1707,6 +1717,8 @@ class FlatExprVisitor : public cel::AstVisitor {
17071717
// Starting slot index for `cel.@block`. We occupy he slot indices `index`
17081718
// through `index + size + (var_size * 2)`.
17091719
size_t index = 0;
1720+
// The total number of slots needed for evaluating the bound expressions.
1721+
size_t slot_count = 0;
17101722
// The current slot index we are processing, any index references must be
17111723
// less than this to be valid.
17121724
size_t current_index = 0;

runtime/internal/BUILD

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,7 @@
1313
# limitations under the License.
1414

1515
package(
16-
# Under active development, not yet being released.
16+
# Internals for cel/runtime.
1717
default_visibility = ["//visibility:public"],
1818
)
1919

0 commit comments

Comments
 (0)