From 4aabfd07cfad9001a7c196319bd5da178f390c30 Mon Sep 17 00:00:00 2001 From: Sokwhan Huh Date: Tue, 5 Mar 2024 14:56:22 -0800 Subject: [PATCH] Allow mutation of AST containing expanded macro that does not exist in macro_call map PiperOrigin-RevId: 612979566 --- .../main/java/dev/cel/common/ast/CelExpr.java | 7 ++ .../java/dev/cel/optimizer/MutableAst.java | 87 +++++++++++++++++-- .../dev/cel/optimizer/MutableAstTest.java | 42 +++++++++ .../resources/subexpression_unparsed.baseline | 18 ++-- 4 files changed, 139 insertions(+), 15 deletions(-) diff --git a/common/src/main/java/dev/cel/common/ast/CelExpr.java b/common/src/main/java/dev/cel/common/ast/CelExpr.java index c4e53906a..3aa3037b5 100644 --- a/common/src/main/java/dev/cel/common/ast/CelExpr.java +++ b/common/src/main/java/dev/cel/common/ast/CelExpr.java @@ -525,6 +525,13 @@ public ImmutableList getArgsBuilders() { return mutableArgs.stream().map(CelExpr::toBuilder).collect(toImmutableList()); } + @CanIgnoreReturnValue + public Builder clearArgs() { + mutableArgs.clear(); + return this; + } + + @CanIgnoreReturnValue public Builder setArg(int index, CelExpr arg) { checkNotNull(arg); mutableArgs.set(index, arg); diff --git a/optimizer/src/main/java/dev/cel/optimizer/MutableAst.java b/optimizer/src/main/java/dev/cel/optimizer/MutableAst.java index bf813eef4..225f7a385 100644 --- a/optimizer/src/main/java/dev/cel/optimizer/MutableAst.java +++ b/optimizer/src/main/java/dev/cel/optimizer/MutableAst.java @@ -645,16 +645,18 @@ private CelSource normalizeMacroSource( })); // Update the macro call IDs and their call references - for (Entry macroCall : celSource.getMacroCalls().entrySet()) { - long macroId = macroCall.getKey(); + for (Entry existingMacroCall : celSource.getMacroCalls().entrySet()) { + long macroId = existingMacroCall.getKey(); long callId = idGenerator.generate(macroId); if (!allExprs.containsKey(callId)) { continue; } - CelExpr.Builder newCall = renumberExprIds(idGenerator, macroCall.getValue().toBuilder()); - CelNavigableExpr callNav = CelNavigableExpr.fromExpr(newCall.build()); + CelExpr.Builder newMacroCallExpr = + renumberExprIds(idGenerator, existingMacroCall.getValue().toBuilder()); + + CelNavigableExpr callNav = CelNavigableExpr.fromExpr(newMacroCallExpr.build()); ImmutableList callDescendants = callNav.descendants().map(CelNavigableExpr::expr).collect(toImmutableList()); @@ -665,10 +667,24 @@ private CelSource normalizeMacroSource( CelExpr mutatedExpr = allExprs.get(callChild.id()); if (!callChild.equals(mutatedExpr)) { - newCall = mutateExpr((arg) -> arg, newCall, mutatedExpr.toBuilder(), callChild.id()); + newMacroCallExpr = + mutateExpr( + NO_OP_ID_GENERATOR, newMacroCallExpr, mutatedExpr.toBuilder(), callChild.id()); + } + } + + if (exprIdToReplace > 0) { + long replacedId = idGenerator.generate(exprIdToReplace); + boolean isListExprBeingReplaced = + allExprs.containsKey(replacedId) + && allExprs.get(replacedId).exprKind().getKind().equals(Kind.CREATE_LIST); + if (isListExprBeingReplaced) { + unwrapListArgumentsInMacroCallExpr( + allExprs.get(callId).comprehension(), newMacroCallExpr); } } - sourceBuilder.addMacroCalls(callId, newCall.build()); + + sourceBuilder.addMacroCalls(callId, newMacroCallExpr.build()); } // Replace comprehension nodes with a NOT_SET reference to reduce AST size. @@ -688,11 +704,70 @@ private CelSource normalizeMacroSource( node.id()); macroCall.setValue(mutatedNode.build()); }); + + // Prune any NOT_SET (comprehension) nodes that no longer exist in the main AST + // This can occur from pulling out a nested comprehension into a separate cel.block index + CelNavigableExpr.fromExpr(macroCallExpr) + .allNodes() + .filter(node -> node.getKind().equals(Kind.NOT_SET)) + .map(CelNavigableExpr::id) + .filter(id -> !allExprs.containsKey(id)) + .forEach( + id -> { + ImmutableList newCallArgs = + macroCallExpr.call().args().stream() + .filter(node -> node.id() != id) + .collect(toImmutableList()); + CelCall.Builder call = + macroCallExpr.call().toBuilder().clearArgs().addArgs(newCallArgs); + + macroCall.setValue(macroCallExpr.toBuilder().setCall(call.build()).build()); + }); } return sourceBuilder.build(); } + /** + * Unwraps the arguments in the extraneous list_expr which is present in the AST but does not + * exist in the macro call map. `map`, `filter` are examples of such. + * + *

This method inspects the comprehension's accumulator initializer to infer that the list_expr + * solely exists to match the expected result type of the macro call signature. + * + * @param comprehension Comprehension in the main AST to extract the macro call arguments from + * (loop step). + * @param newMacroCallExpr (Output parameter) Modified macro call expression with the call + * arguments unwrapped. + */ + private static void unwrapListArgumentsInMacroCallExpr( + CelComprehension comprehension, CelExpr.Builder newMacroCallExpr) { + CelExpr accuInit = comprehension.accuInit(); + if (!accuInit.exprKind().getKind().equals(Kind.CREATE_LIST) + || !accuInit.createList().elements().isEmpty()) { + // Does not contain an extraneous list. + return; + } + + CelExpr loopStepExpr = comprehension.loopStep(); + ImmutableList args = loopStepExpr.call().args(); + if (args.size() != 2) { + throw new IllegalArgumentException( + String.format( + "Expected exactly 2 arguments but got %d instead on expr id: %d", + args.size(), loopStepExpr.id())); + } + + CelCall newMacroCall = newMacroCallExpr.call(); + newMacroCallExpr.setCall( + newMacroCallExpr.call().toBuilder() + .clearArgs() + .addArgs( + newMacroCall.args().get(0)) // iter_var is first argument of the call by convention + .addArgs(args.get(1).createList().elements()) + .build()); + } + private CelExpr.Builder mutateExpr( ExprIdGenerator idGenerator, CelExpr.Builder root, diff --git a/optimizer/src/test/java/dev/cel/optimizer/MutableAstTest.java b/optimizer/src/test/java/dev/cel/optimizer/MutableAstTest.java index af0e1241e..25cc48f9f 100644 --- a/optimizer/src/test/java/dev/cel/optimizer/MutableAstTest.java +++ b/optimizer/src/test/java/dev/cel/optimizer/MutableAstTest.java @@ -18,6 +18,7 @@ import static com.google.common.truth.Truth.assertThat; import static org.junit.Assert.assertThrows; +import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; import com.google.testing.junit.testparameterinjector.TestParameterInjector; import com.google.testing.junit.testparameterinjector.TestParameters; @@ -47,6 +48,7 @@ import dev.cel.parser.CelUnparserFactory; import dev.cel.parser.Operator; import dev.cel.testing.testdata.proto3.TestAllTypesProto.TestAllTypes; +import java.util.List; import org.junit.Test; import org.junit.runner.RunWith; @@ -439,6 +441,46 @@ public void replaceSubtree_macroReplacedWithConstExpr_macroCallCleared() throws assertThat(CEL.createProgram(CEL.check(mutatedAst).getAst()).eval()).isEqualTo(1); } + @Test + @SuppressWarnings("unchecked") // Test only + public void replaceSubtree_replaceExtraneousListCreatedByMacro_unparseSuccess() throws Exception { + // Certain macros such as `map` or `filter` generates an extraneous list_expr in the loop step's + // argument that does not exist in the original expression. + // For example, the loop step of this expression looks like: + // CALL [10] { + // function: _+_ + // args: { + // IDENT [8] { + // name: __result__ + // } + // CREATE_LIST [9] { + // elements: { + // CONSTANT [5] { value: 1 } + // } + // } + // } + // } + CelAbstractSyntaxTree ast = CEL.compile("[1].map(x, 1)").getAst(); + + // These two mutation are equivalent. + CelAbstractSyntaxTree mutatedAstWithList = + MUTABLE_AST.replaceSubtree( + ast, + CelExpr.ofCreateListExpr( + 0, + ImmutableList.of(CelExpr.newBuilder().setConstant(CelConstant.ofValue(2L)).build()), + ImmutableList.of()), + 9L); + CelAbstractSyntaxTree mutatedAstWithConstant = + MUTABLE_AST.replaceSubtree( + ast, CelExpr.newBuilder().setConstant(CelConstant.ofValue(2L)).build(), 5L); + + assertThat(CEL_UNPARSER.unparse(mutatedAstWithList)).isEqualTo("[1].map(x, 2)"); + assertThat(CEL_UNPARSER.unparse(mutatedAstWithConstant)).isEqualTo("[1].map(x, 2)"); + assertThat((List) CEL.createProgram(CEL.check(mutatedAstWithList).getAst()).eval()) + .containsExactly(2L); + } + @Test public void globalCallExpr_replaceRoot() throws Exception { // Tree shape (brackets are expr IDs): diff --git a/optimizer/src/test/resources/subexpression_unparsed.baseline b/optimizer/src/test/resources/subexpression_unparsed.baseline index 28d55996b..83a434ac4 100644 --- a/optimizer/src/test/resources/subexpression_unparsed.baseline +++ b/optimizer/src/test/resources/subexpression_unparsed.baseline @@ -310,10 +310,10 @@ Result: true [BLOCK_COMMON_SUBEXPR_ONLY]: cel.@block([[1, 2, 3], [2, 3, 4]], @index0.map(@c0:0, @index0.map(@c1:0, @c1:0 + 1)) == [@index1, @index1, @index1]) [BLOCK_RECURSION_DEPTH_1]: cel.@block([[1, 2, 3], [2, 3, 4], @c1:0 + 1, [@index2], @x1:0 + @index3, [@index1, @index1, @index1]], @index0.map(@c0:0, @index0.map(@c1:0, @index2)) == @index5) [BLOCK_RECURSION_DEPTH_2]: cel.@block([[1, 2, 3], [2, 3, 4], [@c1:0 + 1], @index0.map(@c1:0, @c1:0 + 1), @x0:0 + [@index3], @index0.map(@c0:0, @index3)], @index5 == [@index1, @index1, @index1]) -[BLOCK_RECURSION_DEPTH_3]: Unparse Error: java.lang.IllegalArgumentException: unexpected expr kind: NOT_SET +[BLOCK_RECURSION_DEPTH_3]: cel.@block([[1, 2, 3], [2, 3, 4], @x1:0 + [@c1:0 + 1], [@index0.map(@c1:0, @c1:0 + 1)]], @index0.map(@c0:0) == [@index1, @index1, @index1]) [BLOCK_RECURSION_DEPTH_4]: cel.@block([[1, 2, 3], [2, 3, 4], @index0.map(@c1:0, @c1:0 + 1)], @index0.map(@c0:0, @index2) == [@index1, @index1, @index1]) -[BLOCK_RECURSION_DEPTH_5]: Unparse Error: java.lang.IllegalArgumentException: unexpected expr kind: NOT_SET -[BLOCK_RECURSION_DEPTH_6]: Unparse Error: java.lang.IllegalArgumentException: unexpected expr kind: NOT_SET +[BLOCK_RECURSION_DEPTH_5]: cel.@block([[1, 2, 3], [2, 3, 4], [@index0.map(@c1:0, @c1:0 + 1)]], @index0.map(@c0:0) == [@index1, @index1, @index1]) +[BLOCK_RECURSION_DEPTH_6]: cel.@block([[1, 2, 3], [2, 3, 4], @x0:0 + [@index0.map(@c1:0, @c1:0 + 1)]], @index0.map(@c0:0) == [@index1, @index1, @index1]) [BLOCK_RECURSION_DEPTH_7]: cel.@block([[1, 2, 3], [2, 3, 4], @index0.map(@c0:0, @index0.map(@c1:0, @c1:0 + 1))], @index2 == [@index1, @index1, @index1]) [BLOCK_RECURSION_DEPTH_8]: cel.@block([[1, 2, 3], [2, 3, 4]], @index0.map(@c0:0, @index0.map(@c1:0, @c1:0 + 1)) == [@index1, @index1, @index1]) [BLOCK_RECURSION_DEPTH_9]: cel.@block([[1, 2, 3], [2, 3, 4]], @index0.map(@c0:0, @index0.map(@c1:0, @c1:0 + 1)) == [@index1, @index1, @index1]) @@ -326,10 +326,10 @@ Result: true [BLOCK_COMMON_SUBEXPR_ONLY]: [1, 2].map(@c0:0, [1, 2, 3].filter(@c1:0, @c1:0 == @c0:0)) == [[1], [2]] [BLOCK_RECURSION_DEPTH_1]: cel.@block([[1, 2], [1, 2, 3], @c1:0 == @c0:0, [@c1:0], @x1:0 + @index3, @index2 ? @index4 : @x1:0, [1], [2], [@index6, @index7]], @index0.map(@c0:0, @index1.filter(@c1:0, @index2)) == @index8) [BLOCK_RECURSION_DEPTH_2]: cel.@block([@x1:0 + [@c1:0], (@c1:0 == @c0:0) ? @index0 : @x1:0, [1, 2, 3].filter(@c1:0, @c1:0 == @c0:0), @x0:0 + [@index2], [1, 2].map(@c0:0, @index2), [[1], [2]]], @index4 == @index5) -[BLOCK_RECURSION_DEPTH_3]: Unparse Error: java.lang.IllegalArgumentException: unexpected expr kind: NOT_SET +[BLOCK_RECURSION_DEPTH_3]: cel.@block([(@c1:0 == @c0:0) ? (@x1:0 + [@c1:0]) : @x1:0, [[1, 2, 3].filter(@c1:0, @c1:0 == @c0:0)]], [1, 2].map(@c0:0) == [[1], [2]]) [BLOCK_RECURSION_DEPTH_4]: cel.@block([[1, 2, 3].filter(@c1:0, @c1:0 == @c0:0)], [1, 2].map(@c0:0, @index0) == [[1], [2]]) -[BLOCK_RECURSION_DEPTH_5]: Unparse Error: java.lang.IllegalArgumentException: unexpected expr kind: NOT_SET -[BLOCK_RECURSION_DEPTH_6]: Unparse Error: java.lang.IllegalArgumentException: unexpected expr kind: NOT_SET +[BLOCK_RECURSION_DEPTH_5]: cel.@block([[[1, 2, 3].filter(@c1:0, @c1:0 == @c0:0)]], [1, 2].map(@c0:0) == [[1], [2]]) +[BLOCK_RECURSION_DEPTH_6]: cel.@block([@x0:0 + [[1, 2, 3].filter(@c1:0, @c1:0 == @c0:0)]], [1, 2].map(@c0:0) == [[1], [2]]) [BLOCK_RECURSION_DEPTH_7]: cel.@block([[1, 2].map(@c0:0, [1, 2, 3].filter(@c1:0, @c1:0 == @c0:0))], @index0 == [[1], [2]]) [BLOCK_RECURSION_DEPTH_8]: [1, 2].map(@c0:0, [1, 2, 3].filter(@c1:0, @c1:0 == @c0:0)) == [[1], [2]] [BLOCK_RECURSION_DEPTH_9]: [1, 2].map(@c0:0, [1, 2, 3].filter(@c1:0, @c1:0 == @c0:0)) == [[1], [2]] @@ -374,10 +374,10 @@ Result: true [BLOCK_COMMON_SUBEXPR_ONLY]: cel.@block([[1, 2], [3, 4], [@index1, @index1]], @index0.map(@c0:0, @index0.map(@c1:0, @index1)) == [@index2, @index2]) [BLOCK_RECURSION_DEPTH_1]: cel.@block([[1, 2], [3, 4], [@index1, @index1], [@index1], @x1:0 + @index3, [@index2, @index2]], @index0.map(@c0:0, @index0.map(@c1:0, @index1)) == @index5) [BLOCK_RECURSION_DEPTH_2]: cel.@block([[[3, 4], [3, 4]], [1, 2], [[3, 4]], @index1.map(@c1:0, [3, 4]), @x0:0 + [@index3], @index1.map(@c0:0, @index3)], @index5 == [@index0, @index0]) -[BLOCK_RECURSION_DEPTH_3]: Unparse Error: java.lang.IllegalArgumentException: unexpected expr kind: NOT_SET +[BLOCK_RECURSION_DEPTH_3]: cel.@block([[[3, 4], [3, 4]], [1, 2], @x1:0 + [[3, 4]], [@index1.map(@c1:0, [3, 4])]], @index1.map(@c0:0) == [@index0, @index0]) [BLOCK_RECURSION_DEPTH_4]: cel.@block([[[3, 4], [3, 4]], [1, 2], @index1.map(@c1:0, [3, 4])], @index1.map(@c0:0, @index2) == [@index0, @index0]) -[BLOCK_RECURSION_DEPTH_5]: Unparse Error: java.lang.IllegalArgumentException: unexpected expr kind: NOT_SET -[BLOCK_RECURSION_DEPTH_6]: Unparse Error: java.lang.IllegalArgumentException: unexpected expr kind: NOT_SET +[BLOCK_RECURSION_DEPTH_5]: cel.@block([[[3, 4], [3, 4]], [1, 2], [@index1.map(@c1:0, [3, 4])]], @index1.map(@c0:0) == [@index0, @index0]) +[BLOCK_RECURSION_DEPTH_6]: cel.@block([[[3, 4], [3, 4]], [1, 2], @x0:0 + [@index1.map(@c1:0, [3, 4])]], @index1.map(@c0:0) == [@index0, @index0]) [BLOCK_RECURSION_DEPTH_7]: cel.@block([[[3, 4], [3, 4]], [1, 2], @index1.map(@c0:0, @index1.map(@c1:0, [3, 4]))], @index2 == [@index0, @index0]) [BLOCK_RECURSION_DEPTH_8]: cel.@block([[[3, 4], [3, 4]], [1, 2]], @index1.map(@c0:0, @index1.map(@c1:0, [3, 4])) == [@index0, @index0]) [BLOCK_RECURSION_DEPTH_9]: cel.@block([[[3, 4], [3, 4]], [1, 2]], @index1.map(@c0:0, @index1.map(@c1:0, [3, 4])) == [@index0, @index0])