From dfebe25d6228578a21506375400be4739ccfaf0d Mon Sep 17 00:00:00 2001 From: Sokwhan Huh Date: Tue, 5 Mar 2024 11:42:29 -0800 Subject: [PATCH] Fix constant folding with in operator involving comprehension identifiers PiperOrigin-RevId: 612916293 --- .../optimizers/ConstantFoldingOptimizer.java | 43 ++++++++++++++----- .../ConstantFoldingOptimizerTest.java | 10 ++++- 2 files changed, 41 insertions(+), 12 deletions(-) diff --git a/optimizer/src/main/java/dev/cel/optimizer/optimizers/ConstantFoldingOptimizer.java b/optimizer/src/main/java/dev/cel/optimizer/optimizers/ConstantFoldingOptimizer.java index 716cdafc6..f913c3d37 100644 --- a/optimizer/src/main/java/dev/cel/optimizer/optimizers/ConstantFoldingOptimizer.java +++ b/optimizer/src/main/java/dev/cel/optimizer/optimizers/ConstantFoldingOptimizer.java @@ -85,22 +85,21 @@ public OptimizationResult optimize(CelNavigableAst navigableAst, Cel cel) if (iterCount >= constantFoldingOptions.maxIterationLimit()) { throw new IllegalStateException("Max iteration count reached."); } - Optional foldableExpr = + Optional 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 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()); @@ -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: @@ -166,6 +165,30 @@ private static boolean canFold(CelNavigableExpr navigableExpr) { } } + private static boolean canFoldInOperator(CelNavigableExpr navigableExpr) { + ImmutableList 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; @@ -195,10 +218,10 @@ private static boolean isNestedComprehension(CelNavigableExpr expr) { } private Optional 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); @@ -209,11 +232,11 @@ private Optional 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 maybeAdaptEvaluatedResult(Object result) { diff --git a/optimizer/src/test/java/dev/cel/optimizer/optimizers/ConstantFoldingOptimizerTest.java b/optimizer/src/test/java/dev/cel/optimizer/optimizers/ConstantFoldingOptimizerTest.java index 2ed564cc9..5ccc41388 100644 --- a/optimizer/src/test/java/dev/cel/optimizer/optimizers/ConstantFoldingOptimizerTest.java +++ b/optimizer/src/test/java/dev/cel/optimizer/optimizers/ConstantFoldingOptimizerTest.java @@ -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; @@ -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(); @@ -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 { @@ -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 = @@ -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 =