Skip to content

Commit a02600a

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 0bd9173 commit a02600a

2 files changed

Lines changed: 158 additions & 53 deletions

File tree

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

Lines changed: 126 additions & 38 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,28 @@ 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
571+
&& (!constantFoldingOptions.enableSafeLogicalOptimization()
572+
|| evaluatesToBoolean(mutableAst, identTypes, rhs))) {
531573
boolean cond = invertCondition != lhs.constant().booleanValue();
532574
replacementExpr =
533575
Optional.of(
534576
cond
535577
? rhs
536578
: CelMutableExpr.ofCall(
537579
CelMutableCall.create(Operator.LOGICAL_NOT.getFunction(), rhs)));
538-
} else if (rhsIsBoolean) {
580+
} else if (rhsIsBooleanConstant
581+
&& (!constantFoldingOptions.enableSafeLogicalOptimization()
582+
|| evaluatesToBoolean(mutableAst, identTypes, lhs))) {
539583
boolean cond = invertCondition != rhs.constant().booleanValue();
540584
replacementExpr =
541585
Optional.of(
@@ -552,7 +596,7 @@ private Optional<CelMutableAst> maybePruneBranches(
552596
}
553597

554598
private Optional<CelMutableAst> maybeShortCircuitCall(
555-
CelMutableAst mutableAst, CelMutableExpr expr) {
599+
CelMutableAst mutableAst, Map<String, CelType> identTypes, CelMutableExpr expr) {
556600
CelMutableCall call = expr.call();
557601
boolean shortCircuit = false;
558602
boolean skip = true;
@@ -583,14 +627,39 @@ private Optional<CelMutableAst> maybeShortCircuitCall(
583627
return Optional.of(astMutator.replaceSubtree(mutableAst, shortCircuitTarget, expr.id()));
584628
}
585629
if (newArgs.size() == 1) {
586-
return Optional.of(astMutator.replaceSubtree(mutableAst, newArgs.get(0), expr.id()));
630+
CelMutableExpr remainingArg = newArgs.get(0);
631+
if (!constantFoldingOptions.enableSafeLogicalOptimization()
632+
|| evaluatesToBoolean(mutableAst, identTypes, remainingArg)) {
633+
return Optional.of(astMutator.replaceSubtree(mutableAst, remainingArg, expr.id()));
634+
}
635+
return Optional.empty();
587636
}
588637

589638
// TODO: Support folding variadic AND/ORs.
590639
throw new UnsupportedOperationException(
591640
"Folding variadic logical operator is not supported yet.");
592641
}
593642

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

812881
public abstract ImmutableSet<String> foldableFunctions();
813882

883+
/**
884+
* Returns true if safe logical optimization is enabled. When enabled, logical (&&, ||) and
885+
* equality (==, !=) expression optimizations strictly verify that pruned sub-expressions
886+
* evaluate to booleans.
887+
*/
888+
public abstract boolean enableSafeLogicalOptimization();
889+
814890
/** Builder for configuring the {@link ConstantFoldingOptions}. */
815891
@AutoValue.Builder
816892
public abstract static class Builder {
@@ -823,6 +899,17 @@ public abstract static class Builder {
823899
*/
824900
public abstract Builder maxIterationLimit(int value);
825901

902+
/**
903+
* Enables or disables safe logical optimization. When enabled (default: {@code true}),
904+
* constant folding on logical (&&, ||) and equality (==, !=) operators strictly checks
905+
* whether sub-expressions evaluate to booleans before pruning them. Disabling this flag
906+
* restores legacy aggressive folding behavior.
907+
*
908+
* <p>Note: Disabling this flag should only be done temporarily for migration purposes, with
909+
* the goal of eventually enabling it for safety.
910+
*/
911+
public abstract Builder enableSafeLogicalOptimization(boolean value);
912+
826913
/**
827914
* Adds a collection of custom functions that will be a candidate for constant folding. By
828915
* default, standard functions are foldable.
@@ -850,7 +937,8 @@ public Builder addFoldableFunctions(String... functions) {
850937
/** Returns a new options builder with recommended defaults pre-configured. */
851938
public static Builder newBuilder() {
852939
return new AutoValue_ConstantFoldingOptimizer_ConstantFoldingOptions.Builder()
853-
.maxIterationLimit(400);
940+
.maxIterationLimit(400)
941+
.enableSafeLogicalOptimization(true);
854942
}
855943

856944
ConstantFoldingOptions() {}

0 commit comments

Comments
 (0)