Skip to content

Commit 0959e0e

Browse files
l46kokcopybara-github
l46kok
authored and
copybara-github
committed
No public description
PiperOrigin-RevId: 1006470468
1 parent b7096de commit 0959e0e

3 files changed

Lines changed: 254 additions & 19 deletions

File tree

‎optimizer/src/main/java/dev/cel/optimizer/optimizers/ConstantFoldingOptimizer.java‎

Lines changed: 153 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,7 @@
3838
import dev.cel.common.ast.CelMutableExpr.CelMutableCall;
3939
import dev.cel.common.ast.CelMutableExpr.CelMutableList;
4040
import dev.cel.common.ast.CelMutableExpr.CelMutableMap;
41+
import dev.cel.common.ast.CelMutableExpr.CelMutableSelect;
4142
import dev.cel.common.ast.CelMutableExpr.CelMutableStruct;
4243
import dev.cel.common.ast.CelMutableExprConverter;
4344
import dev.cel.common.internal.DateTimeHelpers;
@@ -55,7 +56,6 @@
5556
import dev.cel.common.values.StructValue;
5657
import dev.cel.extensions.CelOptionalLibrary.Function;
5758
import dev.cel.optimizer.AstMutator;
58-
import dev.cel.optimizer.AstMutator.SubtreeReplacement;
5959
import dev.cel.optimizer.CelAstOptimizer;
6060
import dev.cel.optimizer.CelOptimizationException;
6161
import dev.cel.runtime.CelAttributePattern;
@@ -73,6 +73,7 @@
7373
import java.util.Map;
7474
import java.util.Objects;
7575
import java.util.Optional;
76+
import java.util.Set;
7677
import org.jspecify.annotations.Nullable;
7778

7879
/**
@@ -141,7 +142,7 @@ public OptimizationResult optimize(CelAbstractSyntaxTree ast, Cel cel)
141142
ImmutableMap<String, CelType> identTypes = precomputeIdentTypes(initialMutableAst);
142143

143144
CelMutableAst mutableAst =
144-
foldConstants(optimizerEnv, valueProvider, identTypes, initialMutableAst);
145+
foldConstants(cel, optimizerEnv, valueProvider, identTypes, initialMutableAst);
145146
mutableAst = pruneOptionalElements(mutableAst);
146147

147148
if (mutableAst == initialMutableAst) {
@@ -172,23 +173,25 @@ private static ImmutableMap<String, CelType> precomputeIdentTypes(CelMutableAst
172173
}
173174

174175
private CelMutableAst foldConstants(
176+
Cel cel,
175177
Cel optimizerEnv,
176178
@Nullable CelValueProvider valueProvider,
177179
ImmutableMap<String, CelType> identTypes,
178180
CelMutableAst mutableAst)
179181
throws CelOptimizationException {
180182
for (int iterCount = 0; iterCount < constantFoldingOptions.maxIterationLimit(); iterCount++) {
181-
Optional<SubtreeReplacement> replacement =
182-
findNextFoldableSubtree(optimizerEnv, valueProvider, identTypes, mutableAst);
183-
if (!replacement.isPresent()) {
183+
Optional<CelMutableAst> nextAst =
184+
findNextFoldableSubtree(cel, optimizerEnv, valueProvider, identTypes, mutableAst);
185+
if (!nextAst.isPresent()) {
184186
return mutableAst;
185187
}
186-
mutableAst = astMutator.replaceSubtree(mutableAst, replacement.get());
188+
mutableAst = nextAst.get();
187189
}
188190
throw new IllegalStateException("Max iteration count reached.");
189191
}
190192

191-
private Optional<SubtreeReplacement> findNextFoldableSubtree(
193+
private Optional<CelMutableAst> findNextFoldableSubtree(
194+
Cel cel,
192195
Cel optimizerEnv,
193196
@Nullable CelValueProvider valueProvider,
194197
ImmutableMap<String, CelType> identTypes,
@@ -199,19 +202,40 @@ private Optional<SubtreeReplacement> findNextFoldableSubtree(
199202
navAst.getRoot().allNodes(TraversalOrder.PRE_ORDER).filter(this::canFold).iterator();
200203
while (foldableExprs.hasNext()) {
201204
CelNavigableMutableExpr foldableExpr = foldableExprs.next();
202-
Optional<CelMutableExpr> pruned =
205+
Optional<CelMutableExpr> candidate =
203206
maybePruneBranches(mutableAst, identTypes, foldableExpr.expr());
204-
if (pruned.isPresent()) {
205-
return Optional.of(SubtreeReplacement.of(foldableExpr.id(), pruned.get()));
207+
if (!candidate.isPresent()) {
208+
candidate = maybeFold(optimizerEnv, valueProvider, foldableExpr);
206209
}
207-
Optional<CelMutableExpr> folded = maybeFold(optimizerEnv, valueProvider, foldableExpr);
208-
if (folded.isPresent()) {
209-
return Optional.of(SubtreeReplacement.of(foldableExpr.id(), folded.get()));
210+
if (candidate.isPresent()) {
211+
Optional<CelMutableAst> replacedAst =
212+
tryReplaceSubtree(cel, optimizerEnv, mutableAst, foldableExpr, candidate.get());
213+
if (replacedAst.isPresent()) {
214+
return replacedAst;
215+
}
210216
}
211217
}
212218
return Optional.empty();
213219
}
214220

221+
private Optional<CelMutableAst> tryReplaceSubtree(
222+
Cel cel,
223+
Cel optimizerEnv,
224+
CelMutableAst mutableAst,
225+
CelNavigableMutableExpr foldableExpr,
226+
CelMutableExpr candidate) {
227+
// Folding an expression containing a dynamic identifier can narrow its static type from dyn to
228+
// a concrete type (e.g. `false || [3, x][0]` -> `false || 3`), which may cause the enclosing
229+
// expression or expected result type to fail type-checking.
230+
CelMutableAst candidateAst =
231+
astMutator.replaceSubtree(
232+
CelMutableAst.fromCelAst(mutableAst.toParsedAst()), candidate, foldableExpr.id());
233+
Cel envToCheck = foldableExpr.parent().isPresent() ? optimizerEnv : cel;
234+
return envToCheck.check(candidateAst.toParsedAst()).hasError()
235+
? Optional.empty()
236+
: Optional.of(candidateAst);
237+
}
238+
215239
private boolean canFold(CelNavigableMutableExpr navigableExpr) {
216240
switch (navigableExpr.getKind()) {
217241
case CALL:
@@ -269,19 +293,55 @@ private boolean canFold(CelNavigableMutableExpr navigableExpr) {
269293
return !CelNavigableExprUtil.hasComprehensionVariable(navigableExpr);
270294
}
271295

272-
// Default case: all call arguments must be constants. If the argument is a container (ex:
273-
// list, map), then its arguments must be a constant.
274-
return areChildrenArgConstant(navigableExpr);
296+
return areChildrenArgConstant(navigableExpr)
297+
|| maybePruneContainerLookup(navigableExpr.expr()).isPresent();
275298
case SELECT:
276299
CelNavigableMutableExpr operand = navigableExpr.children().collect(onlyElement());
277-
return areChildrenArgConstant(operand);
300+
return areChildrenArgConstant(operand)
301+
|| maybePruneContainerLookup(navigableExpr.expr()).isPresent();
278302
case COMPREHENSION:
279303
return !isNestedComprehension(navigableExpr) && containsFoldableFunctionOnly(navigableExpr);
280304
default:
281305
return false;
282306
}
283307
}
284308

309+
private static boolean isErrorFree(CelMutableExpr expr) {
310+
switch (expr.getKind()) {
311+
case CONSTANT:
312+
case IDENT:
313+
return true;
314+
case LIST:
315+
CelMutableList list = expr.list();
316+
return list.optionalIndices().isEmpty()
317+
&& list.elements().stream().allMatch(ConstantFoldingOptimizer::isErrorFree);
318+
case MAP:
319+
return isErrorFreeMap(expr.map());
320+
default:
321+
return false;
322+
}
323+
}
324+
325+
private static boolean isErrorFreeMap(CelMutableMap map) {
326+
Set<CelConstant> keysSeen = new HashSet<>();
327+
for (CelMutableMap.Entry entry : map.entries()) {
328+
if (entry.optionalEntry()
329+
|| !entry.key().getKind().equals(Kind.CONSTANT)
330+
|| !isErrorFree(entry.value())
331+
|| !keysSeen.add(normalizeMapKey(entry.key().constant()))) {
332+
return false;
333+
}
334+
}
335+
return true;
336+
}
337+
338+
private static CelConstant normalizeMapKey(CelConstant key) {
339+
if (key.getKind().equals(CelConstant.Kind.UINT64_VALUE) && key.uint64Value().longValue() >= 0) {
340+
return CelConstant.ofValue(key.uint64Value().longValue());
341+
}
342+
return key;
343+
}
344+
285345
private boolean containsFoldableFunctionOnly(CelNavigableMutableExpr navigableExpr) {
286346
return navigableExpr
287347
.allNodes()
@@ -334,7 +394,13 @@ private Optional<CelMutableExpr> maybeFold(
334394
Object result;
335395
try {
336396
result = evaluateExpr(cel, node);
337-
} catch (CelEvaluationException | CelValidationException e) {
397+
} catch (CelValidationException | CelEvaluationException e) {
398+
if (node.getKind().equals(Kind.COMPREHENSION)) {
399+
// Comprehension may fail to evaluate due to other unknowns the expression references. We'll
400+
// simply opt to not fold the subtree and move on.
401+
return Optional.empty();
402+
}
403+
338404
throw new CelOptimizationException(
339405
"Constant folding failure. Failed to evaluate subtree due to: " + e.getMessage(), e);
340406
}
@@ -479,7 +545,7 @@ private static boolean isCallToFunction(CelMutableExpr expr, String functionName
479545
private Optional<CelMutableExpr> maybePruneBranches(
480546
CelMutableAst mutableAst, Map<String, CelType> identTypes, CelMutableExpr expr) {
481547
if (!expr.getKind().equals(Kind.CALL)) {
482-
return Optional.empty();
548+
return maybePruneContainerLookup(expr);
483549
}
484550

485551
CelMutableCall call = expr.call();
@@ -575,6 +641,74 @@ private Optional<CelMutableExpr> maybePruneBranches(
575641
return replacementExpr;
576642
}
577643

644+
return maybePruneContainerLookup(expr);
645+
}
646+
647+
private static Optional<CelMutableExpr> maybePruneContainerLookup(CelMutableExpr expr) {
648+
if (expr.getKind().equals(Kind.SELECT)) {
649+
CelMutableSelect select = expr.select();
650+
if (select.operand().getKind().equals(Kind.MAP) && isErrorFreeMap(select.operand().map())) {
651+
Optional<CelMutableExpr> val =
652+
findMapEntryByConstantKey(select.operand().map(), CelConstant.ofValue(select.field()));
653+
return select.testOnly()
654+
? Optional.of(CelMutableExpr.ofConstant(CelConstant.ofValue(val.isPresent())))
655+
: val;
656+
}
657+
return Optional.empty();
658+
}
659+
if (!expr.getKind().equals(Kind.CALL)) {
660+
return Optional.empty();
661+
}
662+
CelMutableCall call = expr.call();
663+
String function = call.function();
664+
if (function.equals(Operator.INDEX.getFunction())) {
665+
return findContainerElement(call.args().get(0), call.args().get(1), /* optional= */ false);
666+
}
667+
if (function.equals(Operator.OPTIONAL_INDEX.getFunction())
668+
|| function.equals(Operator.OPTIONAL_SELECT.getFunction())) {
669+
return findContainerElement(call.args().get(0), call.args().get(1), /* optional= */ true);
670+
}
671+
return Optional.empty();
672+
}
673+
674+
private static Optional<CelMutableExpr> findContainerElement(
675+
CelMutableExpr target, CelMutableExpr key, boolean optional) {
676+
Optional<CelMutableExpr> elem;
677+
if (target.getKind().equals(Kind.LIST)
678+
&& isErrorFree(target)
679+
&& isExprConstantOfKind(key, CelConstant.Kind.INT64_VALUE)) {
680+
long idx = key.constant().int64Value();
681+
List<CelMutableExpr> elements = target.list().elements();
682+
elem =
683+
idx >= 0 && idx < elements.size()
684+
? Optional.of(elements.get((int) idx))
685+
: Optional.empty();
686+
} else if (target.getKind().equals(Kind.MAP)
687+
&& isErrorFreeMap(target.map())
688+
&& key.getKind().equals(Kind.CONSTANT)) {
689+
elem = findMapEntryByConstantKey(target.map(), key.constant());
690+
} else {
691+
return Optional.empty();
692+
}
693+
if (!optional) {
694+
return elem;
695+
}
696+
return Optional.of(
697+
elem.map(
698+
val ->
699+
CelMutableExpr.ofCall(
700+
CelMutableCall.create(Function.OPTIONAL_OF.getFunction(), val)))
701+
.orElseGet(ConstantFoldingOptimizer::newOptionalNoneExpr));
702+
}
703+
704+
private static Optional<CelMutableExpr> findMapEntryByConstantKey(
705+
CelMutableMap map, CelConstant targetKey) {
706+
CelConstant normalizedTarget = normalizeMapKey(targetKey);
707+
for (CelMutableMap.Entry entry : map.entries()) {
708+
if (normalizeMapKey(entry.key().constant()).equals(normalizedTarget)) {
709+
return Optional.of(entry.value());
710+
}
711+
}
578712
return Optional.empty();
579713
}
580714

‎optimizer/src/test/java/dev/cel/optimizer/optimizers/BUILD.bazel‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -49,6 +49,8 @@ java_library(
4949
"@maven//:com_google_testparameterinjector_test_parameter_injector",
5050
"//:java_truth",
5151
"@maven//:com_google_truth_extensions_truth_proto_extension",
52+
"//verifier",
53+
"//verifier:verifier_factory",
5254
"@cel_spec//proto/cel/expr:syntax_java_proto",
5355
"@cel_spec//proto/cel/expr/conformance/proto2:test_all_types_java_proto",
5456
"@cel_spec//proto/cel/expr/conformance/proto3:test_all_types_java_proto",

0 commit comments

Comments
 (0)