Skip to content

Allow mutation of AST containing expanded macro that does not exist in macro_call map #266

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
copybara-service merged 1 commit into from
Mar 5, 2024
Merged

Allow mutation of AST containing expanded macro that does not exist in macro_call map #266

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
7 changes: 7 additions & 0 deletions common/src/main/java/dev/cel/common/ast/CelExpr.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -525,6 +525,13 @@ public ImmutableList<CelExpr.Builder> 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);
Expand Down
87 changes: 81 additions & 6 deletions optimizer/src/main/java/dev/cel/optimizer/MutableAst.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -645,16 +645,18 @@ private CelSource normalizeMacroSource(
}));

// Update the macro call IDs and their call references
for (Entry<Long, CelExpr> macroCall : celSource.getMacroCalls().entrySet()) {
long macroId = macroCall.getKey();
for (Entry<Long, CelExpr> 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<CelExpr> callDescendants =
callNav.descendants().map(CelNavigableExpr::expr).collect(toImmutableList());

Expand All @@ -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.
Expand All @@ -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<CelExpr> 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.
*
* <p>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<CelExpr> 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,
Expand Down
42 changes: 42 additions & 0 deletions optimizer/src/test/java/dev/cel/optimizer/MutableAstTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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;

Expand Down Expand Up @@ -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<Long>) CEL.createProgram(CEL.check(mutatedAstWithList).getAst()).eval())
.containsExactly(2L);
}

@Test
public void globalCallExpr_replaceRoot() throws Exception {
// Tree shape (brackets are expr IDs):
Expand Down
18 changes: 9 additions & 9 deletions optimizer/src/test/resources/subexpression_unparsed.baseline
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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])
Expand All @@ -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]]
Expand Down Expand Up @@ -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])
Expand Down