Skip to content

Fix constant folding with in operator involving comprehension identifiers #263

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
1 commit merged into from
Mar 5, 2024
Merged

Fix constant folding with in operator involving comprehension identifiers #263

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
43 changes: 33 additions & 10 deletions optimizer/src/main/java/dev/cel/optimizer/optimizers/ConstantFoldingOptimizer.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -85,22 +85,21 @@ public OptimizationResult optimize(CelNavigableAst navigableAst, Cel cel)
if (iterCount >= constantFoldingOptions.maxIterationLimit()) {
throw new IllegalStateException("Max iteration count reached.");
}
Optional<CelExpr> foldableExpr =
Optional<CelNavigableExpr> foldableExpr =
navigableAst
.getRoot()
.allNodes()
.filter(ConstantFoldingOptimizer::canFold)
.map(CelNavigableExpr::expr)
.filter(expr -> !visitedExprs.contains(expr))
.filter(node -> !visitedExprs.contains(node.expr()))
.findAny();
if (!foldableExpr.isPresent()) {
break;
}
visitedExprs.add(foldableExpr.get());
visitedExprs.add(foldableExpr.get().expr());

Optional<CelAbstractSyntaxTree> mutatedAst;
// Attempt to prune if it is a non-strict call
mutatedAst = maybePruneBranches(navigableAst.getAst(), foldableExpr.get());
mutatedAst = maybePruneBranches(navigableAst.getAst(), foldableExpr.get().expr());
if (!mutatedAst.isPresent()) {
// Evaluate the call then fold
mutatedAst = maybeFold(cel, navigableAst.getAst(), foldableExpr.get());
Expand Down Expand Up @@ -150,7 +149,7 @@ private static boolean canFold(CelNavigableExpr navigableExpr) {
}

if (functionName.equals(Operator.IN.getFunction())) {
return true;
return canFoldInOperator(navigableExpr);
}

// Default case: all call arguments must be constants. If the argument is a container (ex:
Expand All @@ -166,6 +165,30 @@ private static boolean canFold(CelNavigableExpr navigableExpr) {
}
}

private static boolean canFoldInOperator(CelNavigableExpr navigableExpr) {
ImmutableList<CelNavigableExpr> allIdents =
navigableExpr
.allNodes()
.filter(node -> node.getKind().equals(Kind.IDENT))
.collect(toImmutableList());
for (CelNavigableExpr identNode : allIdents) {
CelNavigableExpr parent = identNode.parent().orElse(null);
while (parent != null) {
if (parent.getKind().equals(Kind.COMPREHENSION)) {
if (parent.expr().comprehension().accuVar().equals(identNode.expr().ident().name())) {
// Prevent folding a subexpression if it contains a variable declared by a
// comprehension. The subexpression cannot be compiled without the full context of the
// surrounding comprehension.
return false;
}
}
parent = parent.parent().orElse(null);
}
}

return true;
}

private static boolean areChildrenArgConstant(CelNavigableExpr expr) {
if (expr.getKind().equals(Kind.CONSTANT)) {
return true;
Expand Down Expand Up @@ -195,10 +218,10 @@ private static boolean isNestedComprehension(CelNavigableExpr expr) {
}

private Optional<CelAbstractSyntaxTree> maybeFold(
Cel cel, CelAbstractSyntaxTree ast, CelExpr expr) throws CelOptimizationException {
Cel cel, CelAbstractSyntaxTree ast, CelNavigableExpr node) throws CelOptimizationException {
Object result;
try {
result = CelExprUtil.evaluateExpr(cel, expr);
result = CelExprUtil.evaluateExpr(cel, node.expr());
} catch (CelValidationException | CelEvaluationException e) {
throw new CelOptimizationException(
"Constant folding failure. Failed to evaluate subtree due to: " + e.getMessage(), e);
Expand All @@ -209,11 +232,11 @@ private Optional<CelAbstractSyntaxTree> maybeFold(
// ex2: optional.ofNonZeroValue(5) -> optional.of(5)
if (result instanceof Optional<?>) {
Optional<?> optResult = ((Optional<?>) result);
return maybeRewriteOptional(optResult, ast, expr);
return maybeRewriteOptional(optResult, ast, node.expr());
}

return maybeAdaptEvaluatedResult(result)
.map(celExpr -> mutableAst.replaceSubtree(ast, celExpr, expr.id()));
.map(celExpr -> mutableAst.replaceSubtree(ast, celExpr, node.id()));
}

private Optional<CelExpr> maybeAdaptEvaluatedResult(Object result) {
Expand Down
10 changes: 8 additions & 2 deletions optimizer/src/test/java/dev/cel/optimizer/optimizers/ConstantFoldingOptimizerTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,7 @@
import dev.cel.common.types.ListType;
import dev.cel.common.types.MapType;
import dev.cel.common.types.SimpleType;
import dev.cel.extensions.CelExtensions;
import dev.cel.extensions.CelOptionalLibrary;
import dev.cel.optimizer.CelOptimizationException;
import dev.cel.optimizer.CelOptimizer;
Expand All @@ -48,7 +49,7 @@ public class ConstantFoldingOptimizerTest {
.addVar("map_var", MapType.create(SimpleType.STRING, SimpleType.STRING))
.addMessageTypes(TestAllTypes.getDescriptor())
.setContainer("dev.cel.testing.testdata.proto3")
.addCompilerLibraries(CelOptionalLibrary.INSTANCE)
.addCompilerLibraries(CelExtensions.bindings(), CelOptionalLibrary.INSTANCE)
.addRuntimeLibraries(CelOptionalLibrary.INSTANCE)
.build();

Expand Down Expand Up @@ -159,6 +160,8 @@ public class ConstantFoldingOptimizerTest {
"{source: '{\"a\": dyn([1, 2]), \"b\": x}', expected: '{\"a\": [1, 2], \"b\": x}'}")
@TestParameters("{source: 'map_var[?\"key\"]', expected: 'map_var[?\"key\"]'}")
@TestParameters("{source: '\"abc\" in list_var', expected: '\"abc\" in list_var'}")
@TestParameters(
"{source: 'cel.bind(r0, [1, 2, 3], cel.bind(r1, 1 in r0, r1))', expected: 'true'}")
// TODO: Support folding lists with mixed types. This requires mutable lists.
// @TestParameters("{source: 'dyn([1]) + [1.0]'}")
public void constantFold_success(String source, String expected) throws Exception {
Expand Down Expand Up @@ -198,6 +201,9 @@ public void constantFold_success(String source, String expected) throws Exceptio
@TestParameters(
"{source: '[{}, {\"a\": 1}, {\"b\": 2}].filter(m, has(x.a))', expected:"
+ " '[{}, {\"a\": 1}, {\"b\": 2}].filter(m, has(x.a))'}")
@TestParameters(
"{source: 'cel.bind(r0, [1, 2, 3], cel.bind(r1, 1 in r0 && 2 in x, r1))', expected:"
+ " 'cel.bind(r0, [1, 2, 3], cel.bind(r1, 1 in r0 && 2 in x, r1))'}")
public void constantFold_macros_macroCallMetadataPopulated(String source, String expected)
throws Exception {
Cel cel =
Expand All @@ -207,7 +213,7 @@ public void constantFold_macros_macroCallMetadataPopulated(String source, String
.addMessageTypes(TestAllTypes.getDescriptor())
.setStandardMacros(CelStandardMacro.STANDARD_MACROS)
.setOptions(CelOptions.current().populateMacroCalls(true).build())
.addCompilerLibraries(CelOptionalLibrary.INSTANCE)
.addCompilerLibraries(CelExtensions.bindings(), CelOptionalLibrary.INSTANCE)
.addRuntimeLibraries(CelOptionalLibrary.INSTANCE)
.build();
CelOptimizer celOptimizer =
Expand Down