Skip to content

Commit 769efcf

Browse files
l46kokcopybara-github
l46kok
authored and
copybara-github
committed
Fix ConstantFoldingOptimizer to not treat true && dyn_x as a tautology
true && bool_x continues to fold to bool_x PiperOrigin-RevId: 952324181
1 parent d4c8913 commit 769efcf

2 files changed

Lines changed: 133 additions & 52 deletions

File tree

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

Lines changed: 101 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121

2222
import com.google.auto.value.AutoValue;
2323
import com.google.common.collect.ImmutableList;
24+
import com.google.common.collect.ImmutableMap;
2425
import com.google.common.collect.ImmutableSet;
2526
import com.google.errorprone.annotations.CanIgnoreReturnValue;
2627
import dev.cel.bundle.Cel;
@@ -62,9 +63,12 @@
6263
import java.util.ArrayList;
6364
import java.util.Arrays;
6465
import java.util.Collection;
66+
import java.util.HashMap;
6567
import java.util.HashSet;
68+
import java.util.Iterator;
6669
import java.util.List;
6770
import java.util.Map;
71+
import java.util.Objects;
6872
import java.util.Optional;
6973
import org.jspecify.annotations.Nullable;
7074

@@ -73,6 +77,19 @@
7377
* calls and select statements with their evaluated result.
7478
*/
7579
public final class ConstantFoldingOptimizer implements CelAstOptimizer {
80+
private static final ImmutableSet<String> BOOLEAN_RETURN_OPERATORS =
81+
ImmutableSet.of(
82+
Operator.LOGICAL_AND.getFunction(),
83+
Operator.LOGICAL_OR.getFunction(),
84+
Operator.LOGICAL_NOT.getFunction(),
85+
Operator.EQUALS.getFunction(),
86+
Operator.NOT_EQUALS.getFunction(),
87+
Operator.LESS.getFunction(),
88+
Operator.LESS_EQUALS.getFunction(),
89+
Operator.GREATER.getFunction(),
90+
Operator.GREATER_EQUALS.getFunction(),
91+
Operator.IN.getFunction());
92+
7693
private static final ConstantFoldingOptimizer INSTANCE =
7794
new ConstantFoldingOptimizer(ConstantFoldingOptions.newBuilder().build());
7895

@@ -115,6 +132,25 @@ public OptimizationResult optimize(CelAbstractSyntaxTree ast, Cel cel)
115132
Cel optimizerEnv = builder.setResultType(SimpleType.DYN).build();
116133

117134
CelMutableAst mutableAst = CelMutableAst.fromCelAst(ast);
135+
136+
// HACK: The AstMutator strips type metadata during intermediate folds due to ID renumbering.
137+
// We pre-compute identifier types from the unmutated AST to safely evaluate boolean conditions
138+
// later.
139+
// TODO: Improve AstMutator to retain type metadata when possible.
140+
Map<String, CelType> mutableIdentTypes = new HashMap<>();
141+
Iterator<CelNavigableMutableExpr> identNodes =
142+
CelNavigableMutableAst.fromAst(mutableAst)
143+
.getRoot()
144+
.allNodes()
145+
.filter(node -> node.getKind().equals(Kind.IDENT))
146+
.iterator();
147+
while (identNodes.hasNext()) {
148+
CelNavigableMutableExpr node = identNodes.next();
149+
Optional<CelType> type = mutableAst.getType(node.id());
150+
type.ifPresent(celType -> mutableIdentTypes.put(node.expr().ident().name(), celType));
151+
}
152+
ImmutableMap<String, CelType> identTypes = ImmutableMap.copyOf(mutableIdentTypes);
153+
118154
int iterCount = 0;
119155
boolean continueFolding = true;
120156
while (continueFolding) {
@@ -123,7 +159,6 @@ public OptimizationResult optimize(CelAbstractSyntaxTree ast, Cel cel)
123159
}
124160
iterCount++;
125161
continueFolding = false;
126-
127162
ImmutableList<CelNavigableMutableExpr> foldableExprs =
128163
CelNavigableMutableAst.fromAst(mutableAst)
129164
.getRoot()
@@ -135,7 +170,7 @@ public OptimizationResult optimize(CelAbstractSyntaxTree ast, Cel cel)
135170

136171
Optional<CelMutableAst> mutatedResult;
137172
// Attempt to prune if it is a non-strict call
138-
mutatedResult = maybePruneBranches(mutableAst, foldableExpr.expr());
173+
mutatedResult = maybePruneBranches(mutableAst, identTypes, foldableExpr.expr());
139174
if (!mutatedResult.isPresent()) {
140175
// Evaluate the call then fold
141176
try {
@@ -210,6 +245,9 @@ private boolean canFold(CelNavigableMutableExpr navigableExpr) {
210245

211246
if (functionName.equals(Operator.EQUALS.getFunction())
212247
|| functionName.equals(Operator.NOT_EQUALS.getFunction())) {
248+
if (hasComprehensionVar(navigableExpr)) {
249+
return false;
250+
}
213251
if (mutableCall.args().stream()
214252
.anyMatch(node -> isExprConstantOfKind(node, CelConstant.Kind.BOOLEAN_VALUE))
215253
|| mutableCall.args().stream()
@@ -219,7 +257,7 @@ private boolean canFold(CelNavigableMutableExpr navigableExpr) {
219257
}
220258

221259
if (functionName.equals(Operator.IN.getFunction())) {
222-
return canFoldInOperator(navigableExpr);
260+
return !hasComprehensionVar(navigableExpr);
223261
}
224262

225263
// Default case: all call arguments must be constants. If the argument is a container (ex:
@@ -248,32 +286,31 @@ private static boolean isCallTimestampOrDuration(CelMutableCall call) {
248286
|| call.function().equals(DURATION.functionName());
249287
}
250288

251-
private static boolean canFoldInOperator(CelNavigableMutableExpr navigableExpr) {
252-
ImmutableList<CelNavigableMutableExpr> allIdents =
253-
navigableExpr
254-
.allNodes()
255-
.filter(node -> node.getKind().equals(Kind.IDENT))
256-
.collect(toImmutableList());
257-
for (CelNavigableMutableExpr identNode : allIdents) {
258-
CelNavigableMutableExpr parent = identNode.parent().orElse(null);
259-
while (parent != null) {
260-
if (parent.getKind().equals(Kind.COMPREHENSION)) {
261-
String identName = identNode.expr().ident().name();
262-
CelMutableComprehension parentComprehension = parent.expr().comprehension();
263-
if (parentComprehension.accuVar().equals(identName)
264-
|| parentComprehension.iterVar().equals(identName)
265-
|| parentComprehension.iterVar2().equals(identName)) {
266-
// Prevent folding a subexpression if it contains a variable declared by a
267-
// comprehension. The subexpression cannot be compiled without the full context of the
268-
// surrounding comprehension.
269-
return false;
270-
}
271-
}
272-
parent = parent.parent().orElse(null);
273-
}
274-
}
275-
276-
return true;
289+
private static boolean hasComprehensionVar(CelNavigableMutableExpr expr) {
290+
return expr.allNodes()
291+
.filter(node -> node.getKind().equals(Kind.IDENT))
292+
.anyMatch(
293+
identNode -> {
294+
String identName = identNode.expr().ident().name();
295+
CelNavigableMutableExpr curr = identNode;
296+
Optional<CelNavigableMutableExpr> maybeParent = curr.parent();
297+
while (maybeParent.isPresent()) {
298+
CelNavigableMutableExpr parent = maybeParent.get();
299+
if (parent.getKind().equals(Kind.COMPREHENSION)) {
300+
CelMutableComprehension compre = parent.expr().comprehension();
301+
if ((compre.accuVar().equals(identName)
302+
|| compre.iterVar().equals(identName)
303+
|| compre.iterVar2().equals(identName))
304+
&& curr.id() != compre.iterRange().id()
305+
&& curr.id() != compre.accuInit().id()) {
306+
return true;
307+
}
308+
}
309+
curr = parent;
310+
maybeParent = parent.parent();
311+
}
312+
return false;
313+
});
277314
}
278315

279316
private static boolean areChildrenArgConstant(CelNavigableMutableExpr expr) {
@@ -311,6 +348,9 @@ private Optional<CelMutableAst> maybeFold(
311348
CelMutableAst mutableAst,
312349
CelNavigableMutableExpr node)
313350
throws CelOptimizationException, CelEvaluationException {
351+
if (!node.getKind().equals(Kind.COMPREHENSION) && hasComprehensionVar(node)) {
352+
return Optional.empty();
353+
}
314354
Object result;
315355
try {
316356
result = evaluateExpr(cel, node);
@@ -465,7 +505,7 @@ private static boolean isCallToFunction(CelMutableExpr expr, String functionName
465505

466506
/** Inspects the non-strict calls to determine whether a branch can be removed. */
467507
private Optional<CelMutableAst> maybePruneBranches(
468-
CelMutableAst mutableAst, CelMutableExpr expr) {
508+
CelMutableAst mutableAst, Map<String, CelType> identTypes, CelMutableExpr expr) {
469509
if (!expr.getKind().equals(Kind.CALL)) {
470510
return Optional.empty();
471511
}
@@ -474,7 +514,7 @@ private Optional<CelMutableAst> maybePruneBranches(
474514
String function = call.function();
475515
if (function.equals(Operator.LOGICAL_AND.getFunction())
476516
|| function.equals(Operator.LOGICAL_OR.getFunction())) {
477-
return maybeShortCircuitCall(mutableAst, expr);
517+
return maybeShortCircuitCall(mutableAst, identTypes, expr);
478518
} else if (function.equals(Operator.CONDITIONAL.getFunction())) {
479519
CelMutableExpr cond = call.args().get(0);
480520
CelMutableExpr truthy = call.args().get(1);
@@ -518,24 +558,24 @@ private Optional<CelMutableAst> maybePruneBranches(
518558
|| function.equals(Operator.NOT_EQUALS.getFunction())) {
519559
CelMutableExpr lhs = call.args().get(0);
520560
CelMutableExpr rhs = call.args().get(1);
521-
boolean lhsIsBoolean = isExprConstantOfKind(lhs, CelConstant.Kind.BOOLEAN_VALUE);
522-
boolean rhsIsBoolean = isExprConstantOfKind(rhs, CelConstant.Kind.BOOLEAN_VALUE);
561+
boolean lhsIsBooleanConstant = isExprConstantOfKind(lhs, CelConstant.Kind.BOOLEAN_VALUE);
562+
boolean rhsIsBooleanConstant = isExprConstantOfKind(rhs, CelConstant.Kind.BOOLEAN_VALUE);
523563
boolean invertCondition = function.equals(Operator.NOT_EQUALS.getFunction());
524564
Optional<CelMutableExpr> replacementExpr = Optional.empty();
525565

526566
if (lhs.getKind().equals(Kind.CONSTANT) && rhs.getKind().equals(Kind.CONSTANT)) {
527567
// If both args are const, don't prune any branches and let maybeFold method evaluate this
528568
// subExpr
529569
return Optional.empty();
530-
} else if (lhsIsBoolean) {
570+
} else if (lhsIsBooleanConstant && evaluatesToBoolean(mutableAst, identTypes, rhs)) {
531571
boolean cond = invertCondition != lhs.constant().booleanValue();
532572
replacementExpr =
533573
Optional.of(
534574
cond
535575
? rhs
536576
: CelMutableExpr.ofCall(
537577
CelMutableCall.create(Operator.LOGICAL_NOT.getFunction(), rhs)));
538-
} else if (rhsIsBoolean) {
578+
} else if (rhsIsBooleanConstant && evaluatesToBoolean(mutableAst, identTypes, lhs)) {
539579
boolean cond = invertCondition != rhs.constant().booleanValue();
540580
replacementExpr =
541581
Optional.of(
@@ -552,7 +592,7 @@ private Optional<CelMutableAst> maybePruneBranches(
552592
}
553593

554594
private Optional<CelMutableAst> maybeShortCircuitCall(
555-
CelMutableAst mutableAst, CelMutableExpr expr) {
595+
CelMutableAst mutableAst, Map<String, CelType> identTypes, CelMutableExpr expr) {
556596
CelMutableCall call = expr.call();
557597
boolean shortCircuit = false;
558598
boolean skip = true;
@@ -583,14 +623,38 @@ private Optional<CelMutableAst> maybeShortCircuitCall(
583623
return Optional.of(astMutator.replaceSubtree(mutableAst, shortCircuitTarget, expr.id()));
584624
}
585625
if (newArgs.size() == 1) {
586-
return Optional.of(astMutator.replaceSubtree(mutableAst, newArgs.get(0), expr.id()));
626+
CelMutableExpr remainingArg = newArgs.get(0);
627+
if (evaluatesToBoolean(mutableAst, identTypes, remainingArg)) {
628+
return Optional.of(astMutator.replaceSubtree(mutableAst, remainingArg, expr.id()));
629+
}
630+
return Optional.empty();
587631
}
588632

589633
// TODO: Support folding variadic AND/ORs.
590634
throw new UnsupportedOperationException(
591635
"Folding variadic logical operator is not supported yet.");
592636
}
593637

638+
private boolean evaluatesToBoolean(
639+
CelMutableAst mutableAst, Map<String, CelType> identTypes, CelMutableExpr expr) {
640+
if (isExprConstantOfKind(expr, CelConstant.Kind.BOOLEAN_VALUE)) {
641+
return true;
642+
}
643+
// The AST's type map relies on the type-checker having explicitly populated the type for a
644+
// given node. However, during the optimization pipeline, mutated intermediate nodes might
645+
// temporarily lack type metadata. Standard CEL operators like &&, ||, and == inherently
646+
// always return a boolean, so checking the function name provides a reliable fallback when
647+
// the type map is incomplete.
648+
if (expr.getKind().equals(Kind.CALL)
649+
&& BOOLEAN_RETURN_OPERATORS.contains(expr.call().function())) {
650+
return true;
651+
}
652+
if (expr.getKind().equals(Kind.IDENT)) {
653+
return Objects.equals(identTypes.get(expr.ident().name()), SimpleType.BOOL);
654+
}
655+
return mutableAst.getType(expr.id()).map(SimpleType.BOOL::equals).orElse(false);
656+
}
657+
594658
private boolean isFoldedAggregateLiteral(CelMutableExpr expr) {
595659
if (expr.getKind().equals(Kind.CONSTANT)) {
596660
return true;

‎optimizer/src/test/java/dev/cel/optimizer/optimizers/ConstantFoldingOptimizerTest.java‎

Lines changed: 32 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -80,6 +80,7 @@ private static Cel setupEnv(CelBuilder celBuilder) {
8080
return celBuilder
8181
.addVar("x", SimpleType.DYN)
8282
.addVar("y", SimpleType.DYN)
83+
.addVar("bool_var", SimpleType.BOOL)
8384
.addVar("list_var", ListType.create(SimpleType.STRING))
8485
.addVar("map_var", MapType.create(SimpleType.STRING, SimpleType.STRING))
8586
.setStandardMacros(CelStandardMacro.STANDARD_MACROS)
@@ -127,17 +128,16 @@ private static Cel setupEnv(CelBuilder celBuilder) {
127128
@TestParameters("{source: 'false || false', expected: 'false'}")
128129
@TestParameters("{source: 'true && false || true', expected: 'true'}")
129130
@TestParameters("{source: 'false && true || false', expected: 'false'}")
130-
@TestParameters("{source: 'true && x', expected: 'x'}")
131-
@TestParameters("{source: 'x && true', expected: 'x'}")
131+
@TestParameters("{source: 'true && bool_var', expected: 'bool_var'}")
132+
@TestParameters("{source: 'bool_var && false', expected: 'false'}")
133+
@TestParameters("{source: 'bool_var && true', expected: 'bool_var'}")
134+
@TestParameters("{source: 'false || [1 + 2, x][0]', expected: 'false || [3, x][0]'}")
132135
@TestParameters("{source: 'false && x', expected: 'false'}")
133136
@TestParameters("{source: 'x && false', expected: 'false'}")
134137
@TestParameters("{source: 'true || x', expected: 'true'}")
135138
@TestParameters("{source: 'x || true', expected: 'true'}")
136-
@TestParameters("{source: 'false || x', expected: 'x'}")
137-
@TestParameters("{source: 'x || false', expected: 'x'}")
138-
@TestParameters("{source: 'true && x && true && x', expected: 'x && x'}")
139-
@TestParameters("{source: 'false || x || false || x', expected: 'x || x'}")
140-
@TestParameters("{source: 'false || x || false || y', expected: 'x || y'}")
139+
@TestParameters("{source: 'false || bool_var', expected: 'bool_var'}")
140+
@TestParameters("{source: 'bool_var || false', expected: 'bool_var'}")
141141
@TestParameters("{source: 'true ? x + 1 : x + 2', expected: 'x + 1'}")
142142
@TestParameters("{source: 'false ? x + 1 : x + 2', expected: 'x + 2'}")
143143
@TestParameters(
@@ -230,10 +230,10 @@ private static Cel setupEnv(CelBuilder celBuilder) {
230230
@TestParameters("{source: 'sets.contains([1], [1])', expected: 'true'}")
231231
@TestParameters(
232232
"{source: 'cel.bind(r0, [1, 2, 3], cel.bind(r1, 1 in r0, r1))', expected: 'true'}")
233-
@TestParameters("{source: 'x == true', expected: 'x'}")
234-
@TestParameters("{source: 'true == x', expected: 'x'}")
235-
@TestParameters("{source: 'x == false', expected: '!x'}")
236-
@TestParameters("{source: 'false == x', expected: '!x'}")
233+
@TestParameters("{source: 'bool_var == true', expected: 'bool_var'}")
234+
@TestParameters("{source: 'true == bool_var', expected: 'bool_var'}")
235+
@TestParameters("{source: 'bool_var == false', expected: '!bool_var'}")
236+
@TestParameters("{source: 'false == bool_var', expected: '!bool_var'}")
237237
@TestParameters("{source: 'true == false', expected: 'false'}")
238238
@TestParameters("{source: 'true == true', expected: 'true'}")
239239
@TestParameters("{source: 'false == true', expected: 'false'}")
@@ -257,10 +257,10 @@ private static Cel setupEnv(CelBuilder celBuilder) {
257257
@TestParameters("{source: 'false == false', expected: 'true'}")
258258
@TestParameters("{source: '10 == 42', expected: 'false'}")
259259
@TestParameters("{source: '42 == 42', expected: 'true'}")
260-
@TestParameters("{source: 'x != true', expected: '!x'}")
261-
@TestParameters("{source: 'true != x', expected: '!x'}")
262-
@TestParameters("{source: 'x != false', expected: 'x'}")
263-
@TestParameters("{source: 'false != x', expected: 'x'}")
260+
@TestParameters("{source: 'bool_var != true', expected: '!bool_var'}")
261+
@TestParameters("{source: 'true != bool_var', expected: '!bool_var'}")
262+
@TestParameters("{source: 'bool_var != false', expected: 'bool_var'}")
263+
@TestParameters("{source: 'false != bool_var', expected: 'bool_var'}")
264264
@TestParameters("{source: 'true != false', expected: 'true'}")
265265
@TestParameters("{source: 'true != true', expected: 'false'}")
266266
@TestParameters("{source: 'false != true', expected: 'true'}")
@@ -395,6 +395,7 @@ public void constantFold_protoMessageLiteral_success(String source, String expec
395395
@TestParameters(
396396
"{source: 'cel.bind(myMap, {\"foo\": \"bar\"}, myMap[?\"foo\"].optMap(x, x + \"baz\"))', "
397397
+ "expected: 'optional.of(\"barbaz\")'}")
398+
@TestParameters("{source: '(1 + 2 + 3 == x) && (x in [1, 2, x])', expected: '6 == x'}")
398399
public void constantFold_macros_macroCallMetadataPopulated(String source, String expected)
399400
throws Exception {
400401
Cel cel =
@@ -498,6 +499,22 @@ public void constantFold_macros_withoutMacroCallMetadata(String source) throws E
498499
@TestParameters("{source: '[true].exists(x, x == get_true())'}")
499500
@TestParameters("{source: 'get_list([1, 2]).map(x, x * 2)'}")
500501
@TestParameters("{source: '[(x - 1 > 3) ? (x - 1) : 5].exists(x, x - 1 > 3)'}")
502+
@TestParameters("{source: 'true && x'}")
503+
@TestParameters("{source: 'x && true'}")
504+
@TestParameters("{source: 'false || x'}")
505+
@TestParameters("{source: 'x || false'}")
506+
@TestParameters("{source: 'true && x && true && x'}")
507+
@TestParameters("{source: 'false || x || false || x'}")
508+
@TestParameters("{source: 'false || x || false || y'}")
509+
@TestParameters("{source: 'x == true'}")
510+
@TestParameters("{source: 'true == x'}")
511+
@TestParameters("{source: 'x == false'}")
512+
@TestParameters("{source: 'false == x'}")
513+
@TestParameters("{source: 'x != true'}")
514+
@TestParameters("{source: 'true != x'}")
515+
@TestParameters("{source: 'x != false'}")
516+
@TestParameters("{source: 'false != x'}")
517+
@TestParameters("{source: '[x].exists(item, item == true)'}")
501518
public void constantFold_noOp(String source) throws Exception {
502519
CelAbstractSyntaxTree ast = cel.compile(source).getAst();
503520

0 commit comments

Comments
 (0)