Skip to content

Commit eba8e2a

Browse files
jnthntatumcopybara-github
authored andcommitted
Propagate embedder context to functions declared as contextual.
PiperOrigin-RevId: 846811964
1 parent 707a4ce commit eba8e2a

10 files changed

Lines changed: 154 additions & 93 deletions

eval/compiler/instrumentation_test.cc

Lines changed: 30 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -102,9 +102,9 @@ TEST_F(InstrumentationTest, Basic) {
102102
env_->MutableMessageFactory(), &arena_);
103103
cel::Activation activation;
104104

105-
ASSERT_OK_AND_ASSIGN(
106-
auto value,
107-
plan.EvaluateWithCallback(activation, EvaluationListener(), state));
105+
ASSERT_OK_AND_ASSIGN(auto value, plan.EvaluateWithCallback(
106+
activation, /*embedder_context=*/nullptr,
107+
EvaluationListener(), state));
108108

109109
// AST for the test expression:
110110
// + <4>
@@ -147,9 +147,9 @@ TEST_F(InstrumentationTest, BasicWithConstFolding) {
147147
env_->MutableMessageFactory(), &arena_);
148148
cel::Activation activation;
149149

150-
ASSERT_OK_AND_ASSIGN(
151-
auto value,
152-
plan.EvaluateWithCallback(activation, EvaluationListener(), state));
150+
ASSERT_OK_AND_ASSIGN(auto value, plan.EvaluateWithCallback(
151+
activation, /*embedder_context=*/nullptr,
152+
EvaluationListener(), state));
153153

154154
// AST for the test expression:
155155
// + <4>
@@ -187,16 +187,17 @@ TEST_F(InstrumentationTest, AndShortCircuit) {
187187
activation.InsertOrAssignValue("a", cel::BoolValue(true));
188188
activation.InsertOrAssignValue("b", cel::BoolValue(false));
189189

190-
ASSERT_OK_AND_ASSIGN(
191-
auto value,
192-
plan.EvaluateWithCallback(activation, EvaluationListener(), state));
190+
ASSERT_OK_AND_ASSIGN(auto value, plan.EvaluateWithCallback(
191+
activation, /*embedder_context=*/nullptr,
192+
EvaluationListener(), state));
193193

194194
EXPECT_THAT(expr_ids, ElementsAre(1, 2, 3));
195195

196196
activation.InsertOrAssignValue("a", cel::BoolValue(false));
197197

198-
ASSERT_OK_AND_ASSIGN(value, plan.EvaluateWithCallback(
199-
activation, EvaluationListener(), state));
198+
ASSERT_OK_AND_ASSIGN(
199+
value, plan.EvaluateWithCallback(activation, /*embedder_context=*/nullptr,
200+
EvaluationListener(), state));
200201

201202
EXPECT_THAT(expr_ids, ElementsAre(1, 2, 3, 1, 3));
202203
}
@@ -228,16 +229,17 @@ TEST_F(InstrumentationTest, OrShortCircuit) {
228229
activation.InsertOrAssignValue("a", cel::BoolValue(false));
229230
activation.InsertOrAssignValue("b", cel::BoolValue(true));
230231

231-
ASSERT_OK_AND_ASSIGN(
232-
auto value,
233-
plan.EvaluateWithCallback(activation, EvaluationListener(), state));
232+
ASSERT_OK_AND_ASSIGN(auto value, plan.EvaluateWithCallback(
233+
activation, /*embedder_context=*/nullptr,
234+
EvaluationListener(), state));
234235

235236
EXPECT_THAT(expr_ids, ElementsAre(1, 2, 3));
236237
expr_ids.clear();
237238
activation.InsertOrAssignValue("a", cel::BoolValue(true));
238239

239-
ASSERT_OK_AND_ASSIGN(value, plan.EvaluateWithCallback(
240-
activation, EvaluationListener(), state));
240+
ASSERT_OK_AND_ASSIGN(
241+
value, plan.EvaluateWithCallback(activation, /*embedder_context=*/nullptr,
242+
EvaluationListener(), state));
241243

242244
EXPECT_THAT(expr_ids, ElementsAre(1, 3));
243245
}
@@ -270,9 +272,9 @@ TEST_F(InstrumentationTest, Ternary) {
270272
activation.InsertOrAssignValue("a", cel::IntValue(1));
271273
activation.InsertOrAssignValue("b", cel::IntValue(2));
272274

273-
ASSERT_OK_AND_ASSIGN(
274-
auto value,
275-
plan.EvaluateWithCallback(activation, EvaluationListener(), state));
275+
ASSERT_OK_AND_ASSIGN(auto value, plan.EvaluateWithCallback(
276+
activation, /*embedder_context=*/nullptr,
277+
EvaluationListener(), state));
276278

277279
// AST
278280
// ?:() <2>
@@ -283,8 +285,9 @@ TEST_F(InstrumentationTest, Ternary) {
283285

284286
activation.InsertOrAssignValue("c", cel::BoolValue(false));
285287

286-
ASSERT_OK_AND_ASSIGN(value, plan.EvaluateWithCallback(
287-
activation, EvaluationListener(), state));
288+
ASSERT_OK_AND_ASSIGN(
289+
value, plan.EvaluateWithCallback(activation, /*embedder_context=*/nullptr,
290+
EvaluationListener(), state));
288291

289292
EXPECT_THAT(expr_ids, ElementsAre(1, 4, 2));
290293
expr_ids.clear();
@@ -317,9 +320,9 @@ TEST_F(InstrumentationTest, OptimizedStepsNotEvaluated) {
317320
env_->MutableMessageFactory(), &arena_);
318321
cel::Activation activation;
319322

320-
ASSERT_OK_AND_ASSIGN(
321-
auto value,
322-
plan.EvaluateWithCallback(activation, EvaluationListener(), state));
323+
ASSERT_OK_AND_ASSIGN(auto value, plan.EvaluateWithCallback(
324+
activation, /*embedder_context=*/nullptr,
325+
EvaluationListener(), state));
323326

324327
EXPECT_THAT(expr_ids, ElementsAre(1, 2));
325328
EXPECT_TRUE(value.Is<cel::BoolValue>() && value.GetBool().NativeValue());
@@ -346,9 +349,9 @@ TEST_F(InstrumentationTest, NoopSkipped) {
346349
activation.InsertOrAssignValue("a", cel::IntValue(1));
347350
activation.InsertOrAssignValue("b", cel::IntValue(2));
348351

349-
ASSERT_OK_AND_ASSIGN(
350-
auto value,
351-
plan.EvaluateWithCallback(activation, EvaluationListener(), state));
352+
ASSERT_OK_AND_ASSIGN(auto value, plan.EvaluateWithCallback(
353+
activation, /*embedder_context=*/nullptr,
354+
EvaluationListener(), state));
352355

353356
// AST
354357
// ?:() <2>

eval/eval/cel_expression_flat_impl.cc

Lines changed: 7 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -82,10 +82,11 @@ absl::StatusOr<CelValue> CelExpressionFlatImpl::Trace(
8282
state->state().Reset();
8383
cel::interop_internal::AdapterActivationImpl modern_activation(activation);
8484

85-
CEL_ASSIGN_OR_RETURN(
86-
cel::Value value,
87-
flat_expression_.EvaluateWithCallback(
88-
modern_activation, AdaptListener(callback), state->state()));
85+
CEL_ASSIGN_OR_RETURN(cel::Value value,
86+
flat_expression_.EvaluateWithCallback(
87+
modern_activation,
88+
/*embedder_context=*/nullptr,
89+
AdaptListener(callback), state->state()));
8990

9091
return cel::interop_internal::ModernValueToLegacyValueOrDie(state->arena(),
9192
value);
@@ -128,7 +129,8 @@ absl::StatusOr<CelValue> CelExpressionRecursiveImpl::Trace(
128129
ExecutionFrameBase execution_frame(
129130
modern_activation, AdaptListener(callback), flat_expression_.options(),
130131
flat_expression_.type_provider(), env_->descriptor_pool.get(),
131-
env_->MutableMessageFactory(), arena, slots);
132+
env_->MutableMessageFactory(), arena,
133+
/*embedder_context=*/nullptr, slots);
132134

133135
cel::Value result;
134136
AttributeTrail trail;

eval/eval/comprehension_step_test.cc

Lines changed: 40 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -228,10 +228,11 @@ class DirectComprehensionTest : public testing::Test {
228228
TEST_F(DirectComprehensionTest, PropagateRangeNonOkStatus) {
229229
cel::RuntimeOptions options;
230230

231-
ExecutionFrameBase frame(
232-
empty_activation_, /*callback=*/nullptr, options, type_provider_,
233-
cel::internal::GetTestingDescriptorPool(),
234-
cel::internal::GetTestingMessageFactory(), &arena_, slots_);
231+
ExecutionFrameBase frame(empty_activation_, /*callback=*/nullptr, options,
232+
type_provider_,
233+
cel::internal::GetTestingDescriptorPool(),
234+
cel::internal::GetTestingMessageFactory(), &arena_,
235+
/*embedder_context=*/nullptr, slots_);
235236

236237
auto range_step = std::make_unique<MockDirectStep>();
237238
MockDirectStep* mock = range_step.get();
@@ -257,10 +258,11 @@ TEST_F(DirectComprehensionTest, PropagateRangeNonOkStatus) {
257258
TEST_F(DirectComprehensionTest, PropagateAccuInitNonOkStatus) {
258259
cel::RuntimeOptions options;
259260

260-
ExecutionFrameBase frame(
261-
empty_activation_, /*callback=*/nullptr, options, type_provider_,
262-
cel::internal::GetTestingDescriptorPool(),
263-
cel::internal::GetTestingMessageFactory(), &arena_, slots_);
261+
ExecutionFrameBase frame(empty_activation_, /*callback=*/nullptr, options,
262+
type_provider_,
263+
cel::internal::GetTestingDescriptorPool(),
264+
cel::internal::GetTestingMessageFactory(), &arena_,
265+
/*embedder_context=*/nullptr, slots_);
264266

265267
auto accu_init = std::make_unique<MockDirectStep>();
266268
MockDirectStep* mock = accu_init.get();
@@ -288,10 +290,11 @@ TEST_F(DirectComprehensionTest, PropagateAccuInitNonOkStatus) {
288290
TEST_F(DirectComprehensionTest, PropagateLoopNonOkStatus) {
289291
cel::RuntimeOptions options;
290292

291-
ExecutionFrameBase frame(
292-
empty_activation_, /*callback=*/nullptr, options, type_provider_,
293-
cel::internal::GetTestingDescriptorPool(),
294-
cel::internal::GetTestingMessageFactory(), &arena_, slots_);
293+
ExecutionFrameBase frame(empty_activation_, /*callback=*/nullptr, options,
294+
type_provider_,
295+
cel::internal::GetTestingDescriptorPool(),
296+
cel::internal::GetTestingMessageFactory(), &arena_,
297+
/*embedder_context=*/nullptr, slots_);
295298

296299
auto loop_step = std::make_unique<MockDirectStep>();
297300
MockDirectStep* mock = loop_step.get();
@@ -319,10 +322,11 @@ TEST_F(DirectComprehensionTest, PropagateLoopNonOkStatus) {
319322
TEST_F(DirectComprehensionTest, PropagateConditionNonOkStatus) {
320323
cel::RuntimeOptions options;
321324

322-
ExecutionFrameBase frame(
323-
empty_activation_, /*callback=*/nullptr, options, type_provider_,
324-
cel::internal::GetTestingDescriptorPool(),
325-
cel::internal::GetTestingMessageFactory(), &arena_, slots_);
325+
ExecutionFrameBase frame(empty_activation_, /*callback=*/nullptr, options,
326+
type_provider_,
327+
cel::internal::GetTestingDescriptorPool(),
328+
cel::internal::GetTestingMessageFactory(), &arena_,
329+
/*embedder_context=*/nullptr, slots_);
326330

327331
auto condition = std::make_unique<MockDirectStep>();
328332
MockDirectStep* mock = condition.get();
@@ -350,10 +354,11 @@ TEST_F(DirectComprehensionTest, PropagateConditionNonOkStatus) {
350354
TEST_F(DirectComprehensionTest, PropagateResultNonOkStatus) {
351355
cel::RuntimeOptions options;
352356

353-
ExecutionFrameBase frame(
354-
empty_activation_, /*callback=*/nullptr, options, type_provider_,
355-
cel::internal::GetTestingDescriptorPool(),
356-
cel::internal::GetTestingMessageFactory(), &arena_, slots_);
357+
ExecutionFrameBase frame(empty_activation_, /*callback=*/nullptr, options,
358+
type_provider_,
359+
cel::internal::GetTestingDescriptorPool(),
360+
cel::internal::GetTestingMessageFactory(), &arena_,
361+
/*embedder_context=*/nullptr, slots_);
357362

358363
auto result_step = std::make_unique<MockDirectStep>();
359364
MockDirectStep* mock = result_step.get();
@@ -381,10 +386,11 @@ TEST_F(DirectComprehensionTest, PropagateResultNonOkStatus) {
381386
TEST_F(DirectComprehensionTest, Shortcircuit) {
382387
cel::RuntimeOptions options;
383388

384-
ExecutionFrameBase frame(
385-
empty_activation_, /*callback=*/nullptr, options, type_provider_,
386-
cel::internal::GetTestingDescriptorPool(),
387-
cel::internal::GetTestingMessageFactory(), &arena_, slots_);
389+
ExecutionFrameBase frame(empty_activation_, /*callback=*/nullptr, options,
390+
type_provider_,
391+
cel::internal::GetTestingDescriptorPool(),
392+
cel::internal::GetTestingMessageFactory(), &arena_,
393+
/*embedder_context=*/nullptr, slots_);
388394

389395
auto loop_step = std::make_unique<MockDirectStep>();
390396
MockDirectStep* mock = loop_step.get();
@@ -416,10 +422,11 @@ TEST_F(DirectComprehensionTest, Shortcircuit) {
416422
TEST_F(DirectComprehensionTest, IterationLimit) {
417423
cel::RuntimeOptions options;
418424
options.comprehension_max_iterations = 2;
419-
ExecutionFrameBase frame(
420-
empty_activation_, /*callback=*/nullptr, options, type_provider_,
421-
cel::internal::GetTestingDescriptorPool(),
422-
cel::internal::GetTestingMessageFactory(), &arena_, slots_);
425+
ExecutionFrameBase frame(empty_activation_, /*callback=*/nullptr, options,
426+
type_provider_,
427+
cel::internal::GetTestingDescriptorPool(),
428+
cel::internal::GetTestingMessageFactory(), &arena_,
429+
/*embedder_context=*/nullptr, slots_);
423430

424431
auto loop_step = std::make_unique<MockDirectStep>();
425432
MockDirectStep* mock = loop_step.get();
@@ -451,10 +458,11 @@ TEST_F(DirectComprehensionTest, IterationLimit) {
451458
TEST_F(DirectComprehensionTest, Exhaustive) {
452459
cel::RuntimeOptions options;
453460

454-
ExecutionFrameBase frame(
455-
empty_activation_, /*callback=*/nullptr, options, type_provider_,
456-
cel::internal::GetTestingDescriptorPool(),
457-
cel::internal::GetTestingMessageFactory(), &arena_, slots_);
461+
ExecutionFrameBase frame(empty_activation_, /*callback=*/nullptr, options,
462+
type_provider_,
463+
cel::internal::GetTestingDescriptorPool(),
464+
cel::internal::GetTestingMessageFactory(), &arena_,
465+
/*embedder_context=*/nullptr, slots_);
458466

459467
auto loop_step = std::make_unique<MockDirectStep>();
460468
MockDirectStep* mock = loop_step.get();

eval/eval/evaluator_core.cc

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -164,12 +164,13 @@ FlatExpressionEvaluatorState FlatExpression::MakeEvaluatorState(
164164
}
165165

166166
absl::StatusOr<cel::Value> FlatExpression::EvaluateWithCallback(
167-
const cel::ActivationInterface& activation, EvaluationListener listener,
168-
FlatExpressionEvaluatorState& state) const {
167+
const cel::ActivationInterface& activation,
168+
const cel::EmbedderContext* absl_nullable embedder_context,
169+
EvaluationListener listener, FlatExpressionEvaluatorState& state) const {
169170
state.Reset();
170171

171172
ExecutionFrame frame(subexpressions_, activation, options_, state,
172-
std::move(listener));
173+
std::move(listener), embedder_context);
173174

174175
return frame.Evaluate(frame.callback());
175176
}

0 commit comments

Comments
 (0)