Skip to content

Commit a6e667d

Browse files
l46kokl46kok
committed
Add select inlining
1 parent 90d0931 commit a6e667d

3 files changed

Lines changed: 189 additions & 70 deletions

File tree

‎optimizer/src/main/java/dev/cel/optimizer/optimizers/BUILD.bazel‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -92,6 +92,8 @@ java_library(
9292
"//common/navigation",
9393
"//common/navigation:common",
9494
"//common/navigation:mutable_navigation",
95+
"//common:operator",
96+
"//common/values:values",
9597
"//common/types",
9698
"//common/types:type_providers",
9799
"//optimizer:ast_optimizer",

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

Lines changed: 87 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -7,15 +7,26 @@
77
import dev.cel.bundle.Cel;
88
import dev.cel.common.CelAbstractSyntaxTree;
99
import dev.cel.common.CelMutableAst;
10+
import dev.cel.common.ast.CelExpr;
1011
import dev.cel.common.ast.CelExpr.ExprKind.Kind;
1112
import dev.cel.common.ast.CelMutableExpr.CelMutableComprehension;
1213
import dev.cel.common.ast.CelMutableExprConverter;
1314
import dev.cel.common.navigation.CelNavigableMutableAst;
1415
import dev.cel.common.navigation.CelNavigableMutableExpr;
16+
import dev.cel.common.Operator;
17+
import dev.cel.common.ast.CelConstant;
1518
import dev.cel.optimizer.AstMutator;
1619
import dev.cel.optimizer.CelAstOptimizer;
1720
import dev.cel.optimizer.CelOptimizationException;
18-
21+
import java.util.Optional;
22+
import dev.cel.common.values.NullValue;
23+
import java.util.stream.Stream;
24+
25+
/**
26+
* Performs optimization for inlining variables within function calls and select
27+
* statements with
28+
* their associated AST.
29+
*/
1930
public final class InliningOptimizer implements CelAstOptimizer {
2031

2132
private final ImmutableList<InlineVariable> inlineVariables;
@@ -38,43 +49,99 @@ public OptimizationResult optimize(CelAbstractSyntaxTree ast, Cel cel)
3849
throws CelOptimizationException {
3950
CelMutableAst mutableAst = CelMutableAst.fromCelAst(ast);
4051
for (InlineVariable inlineVariable : inlineVariables) {
41-
ImmutableList<CelNavigableMutableExpr> inlinableExprs =
42-
CelNavigableMutableAst.fromAst(mutableAst)
43-
.getRoot()
44-
.allNodes()
45-
.filter(node -> node.getKind().equals(Kind.IDENT))
46-
.filter(node -> node.expr().ident().name().equals(inlineVariable.name()))
47-
.filter(expr -> canInline(expr, inlineVariable.name()))
48-
.collect(toImmutableList());
52+
ImmutableList<CelNavigableMutableExpr> inlinableExprs = CelNavigableMutableAst.fromAst(mutableAst)
53+
.getRoot()
54+
.allNodes()
55+
.filter(node -> canInline(node, inlineVariable.name()))
56+
.collect(toImmutableList());
4957

5058
for (CelNavigableMutableExpr inlinableExpr : inlinableExprs) {
51-
mutableAst = astMutator.replaceSubtree(mutableAst, CelMutableExprConverter.fromCelExpr(inlineVariable.ast().getExpr()), inlinableExpr.id());
59+
CelExpr replacementExpr = inlineVariable.ast().getExpr();
60+
61+
if (inlinableExpr.getKind().equals(Kind.SELECT) && inlinableExpr.expr().select().testOnly()) {
62+
if (replacementExpr.getKind().equals(Kind.SELECT)) {
63+
// Preserve testOnly property for Select replacements (has(A) -> has(B))
64+
replacementExpr = replacementExpr.toBuilder().setSelect(
65+
replacementExpr.select().toBuilder().setTestOnly(true).build()).build();
66+
} else {
67+
// Rewrite has(X) -> X != null for non-select replacements
68+
replacementExpr = CelExpr.newBuilder()
69+
.setId(0)
70+
.setCall(
71+
CelExpr.CelCall.newBuilder()
72+
.setFunction(Operator.NOT_EQUALS.getFunction())
73+
.addArgs(replacementExpr)
74+
.addArgs(CelExpr.newBuilder().setId(0).setConstant(CelConstant.ofValue(NullValue.NULL_VALUE))
75+
.build())
76+
.build())
77+
.build();
78+
}
79+
}
80+
81+
mutableAst = astMutator.replaceSubtree(mutableAst,
82+
CelMutableExprConverter.fromCelExpr(replacementExpr), inlinableExpr.id());
5283
}
5384
}
5485

5586
return OptimizationResult.create(mutableAst.toParsedAst());
5687
}
5788

58-
private static boolean canInline(CelNavigableMutableExpr expr, String identifier) {
59-
for (CelNavigableMutableExpr p = expr.parent().orElse(null); p != null; p = p.parent().orElse(null)) {
89+
private static boolean canInline(CelNavigableMutableExpr node, String identifier) {
90+
boolean matches = maybeToQualifiedName(node)
91+
.map(name -> name.equals(identifier))
92+
.orElse(false);
93+
94+
if (!matches) {
95+
return false;
96+
}
97+
98+
for (CelNavigableMutableExpr p = node.parent().orElse(null); p != null; p = p.parent().orElse(null)) {
6099
if (p.getKind() != Kind.COMPREHENSION) {
61100
continue;
62101
}
63102

64103
CelMutableComprehension comp = p.expr().comprehension();
104+
boolean shadows = Stream.of(comp.iterVar(), comp.iterVar2(), comp.accuVar())
105+
.anyMatch(identifier::equals);
65106

66-
if (comp.iterVar().equals(identifier) || comp.iterVar2().equals(identifier) || comp.accuVar().equals(identifier)) {
107+
if (shadows) {
67108
return false;
68109
}
69110
}
111+
70112
return true;
71113
}
72114

115+
private static Optional<String> maybeToQualifiedName(CelNavigableMutableExpr node) {
116+
if (node.getKind().equals(Kind.IDENT)) {
117+
return Optional.of(node.expr().ident().name());
118+
}
119+
120+
if (node.getKind().equals(Kind.SELECT)) {
121+
return node.children().findFirst()
122+
.flatMap(InliningOptimizer::maybeToQualifiedName)
123+
.map(operandName -> operandName + "." + node.expr().select().field());
124+
}
125+
126+
return Optional.empty();
127+
}
128+
129+
/**
130+
* Represents a variable to be inlined.
131+
*/
73132
@AutoValue
74133
public abstract static class InlineVariable {
75134
public abstract String name();
135+
76136
public abstract CelAbstractSyntaxTree ast();
77137

138+
/**
139+
* Creates a new {@link InlineVariable} with the given name and AST.
140+
*
141+
* <p>
142+
* The name must be a qualified name (e.g. "a.b.c") and cannot be an internal
143+
* variable (starting with @).
144+
*/
78145
public static InlineVariable of(String name, CelAbstractSyntaxTree ast) {
79146
if (name.startsWith("@")) {
80147
throw new IllegalArgumentException("Internal variables cannot be inlined: " + name);
@@ -83,9 +150,7 @@ public static InlineVariable of(String name, CelAbstractSyntaxTree ast) {
83150
}
84151
}
85152

86-
/**
87-
* TODO
88-
*/
153+
/** Options to configure how Inlining behaves. */
89154
@AutoValue
90155
public abstract static class InliningOptions {
91156
public abstract int maxIterationLimit();
@@ -95,14 +160,16 @@ public abstract static class InliningOptions {
95160
public abstract static class Builder {
96161

97162
/**
98-
* Limit the number of iteration while inlining variables. An exception is thrown if
163+
* Limit the number of iteration while inlining variables. An exception is
164+
* thrown if
99165
* the iteration count exceeds the set value.
100166
*/
101167
public abstract InliningOptions.Builder maxIterationLimit(int value);
102168

103169
public abstract InliningOptimizer.InliningOptions build();
104170

105-
Builder() {}
171+
Builder() {
172+
}
106173
}
107174

108175
/** Returns a new options builder with recommended defaults pre-configured. */
@@ -111,7 +178,8 @@ public static InliningOptimizer.InliningOptions.Builder newBuilder() {
111178
.maxIterationLimit(400);
112179
}
113180

114-
InliningOptions() {}
181+
InliningOptions() {
182+
}
115183
}
116184

117185
private InliningOptimizer(InliningOptions options, ImmutableList<InlineVariable> inlineVariables) {

‎optimizer/src/test/java/dev/cel/optimizer/optimizers/InliningOptimizerTest.java‎

Lines changed: 100 additions & 51 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
import static com.google.common.truth.Truth.assertThat;
44
import static org.junit.Assert.assertThrows;
55

6+
import com.google.testing.junit.testparameterinjector.TestParameter;
67
import com.google.testing.junit.testparameterinjector.TestParameterInjector;
78
import dev.cel.bundle.Cel;
89
import dev.cel.bundle.CelFactory;
@@ -12,6 +13,9 @@
1213
import dev.cel.common.ast.CelExpr;
1314
import dev.cel.common.types.SimpleType;
1415
import dev.cel.extensions.CelExtensions;
16+
import dev.cel.common.CelContainer;
17+
import dev.cel.common.types.StructTypeReference;
18+
import dev.cel.expr.conformance.proto3.TestAllTypes;
1519
import dev.cel.optimizer.CelOptimizationException;
1620
import dev.cel.optimizer.CelOptimizer;
1721
import dev.cel.optimizer.CelOptimizerFactory;
@@ -25,80 +29,125 @@
2529
@RunWith(TestParameterInjector.class)
2630
public class InliningOptimizerTest {
2731

32+
private static final Cel CEL = CelFactory.standardCelBuilder()
33+
.setStandardMacros(CelStandardMacro.STANDARD_MACROS)
34+
.setContainer(CelContainer.ofName("google.expr.proto3.test"))
35+
.addFileTypes(TestAllTypes.getDescriptor().getFile())
36+
.addCompilerLibraries(CelExtensions.bindings())
37+
.addVar("int_var_to_inline", SimpleType.INT)
38+
.addVar("a", SimpleType.DYN)
39+
.addVar("msg", StructTypeReference.create(TestAllTypes.getDescriptor().getFullName()))
40+
.addVar("unpacked_wrapper", SimpleType.STRING)
41+
.addVar("wrapper_var", StructTypeReference.create("google.protobuf.Int64Value"))
42+
.addVar("child", StructTypeReference.create(TestAllTypes.NestedMessage.getDescriptor().getFullName()))
43+
.addVar("shadowed_ident", SimpleType.INT)
44+
.addVar("x", SimpleType.DYN)
45+
.setOptions(CelOptions.current().populateMacroCalls(true).build())
46+
.build();
47+
2848
@Test
29-
public void inlineConstant() throws Exception {
30-
Cel cel = CelFactory.standardCelBuilder().addVar("ident_to_inline", SimpleType.INT).build();
31-
String expression = "ident_to_inline + 2 + ident_to_inline";
32-
CelAbstractSyntaxTree astToInline = cel.compile(expression).getAst();
33-
CelOptimizer optimizer = CelOptimizerFactory.standardCelOptimizerBuilder(cel)
49+
public void inlining_success(@TestParameter SuccessTestCase testCase) throws Exception {
50+
CelAbstractSyntaxTree astToInline = CEL.compile(testCase.source).getAst();
51+
CelAbstractSyntaxTree replacementAst = CEL.compile(testCase.replacement).getAst();
52+
53+
CelOptimizer optimizer = CelOptimizerFactory.standardCelOptimizerBuilder(CEL)
3454
.addAstOptimizers(InliningOptimizer.newInstance(
35-
InlineVariable.of("ident_to_inline", cel.compile("1").getAst())
36-
))
55+
InlineVariable.of(testCase.varName, replacementAst)))
3756
.build();
3857

3958
CelAbstractSyntaxTree optimized = optimizer.optimize(astToInline);
4059

4160
String unparsed = CelUnparserFactory.newUnparser().unparse(optimized);
42-
assertThat(unparsed).isEqualTo("1 + 2 + 1");
61+
assertThat(unparsed).isEqualTo(testCase.expected);
4362
}
4463

4564
@Test
46-
public void inline_doesNotInlineIterVar() throws Exception {
47-
Cel cel = CelFactory.standardCelBuilder()
48-
.setStandardMacros(CelStandardMacro.STANDARD_MACROS)
49-
.setOptions(CelOptions.current().populateMacroCalls(true).build())
50-
.addVar("shadowed_ident", SimpleType.INT)
51-
.build();
52-
String expression = "[0].exists(shadowed_ident, shadowed_ident == 0)";
53-
CelAbstractSyntaxTree astToInline = cel.compile(expression).getAst();
54-
CelOptimizer optimizer = CelOptimizerFactory.standardCelOptimizerBuilder(cel)
65+
public void inlining_noop(@TestParameter NoOpTestCase testCase) throws Exception {
66+
CelAbstractSyntaxTree astToInline = CEL.compile(testCase.source).getAst();
67+
CelAbstractSyntaxTree replacementAst = CEL.compile(testCase.replacement).getAst();
68+
69+
CelOptimizer optimizer = CelOptimizerFactory.standardCelOptimizerBuilder(CEL)
5570
.addAstOptimizers(InliningOptimizer.newInstance(
56-
InlineVariable.of("shadowed_ident", cel.compile("1").getAst())
57-
))
71+
InlineVariable.of(testCase.varName, replacementAst)))
5872
.build();
5973

6074
CelAbstractSyntaxTree optimized = optimizer.optimize(astToInline);
6175

6276
String unparsed = CelUnparserFactory.newUnparser().unparse(optimized);
63-
assertThat(unparsed).isEqualTo("[0].exists(shadowed_ident, shadowed_ident == 0)");
77+
assertThat(unparsed).isEqualTo(testCase.source);
6478
}
6579

66-
@Test
67-
public void inline_bindMacro_doesNotInlineVarName() throws Exception {
68-
Cel cel = CelFactory.standardCelBuilder()
69-
.setStandardMacros(CelStandardMacro.STANDARD_MACROS)
70-
.addCompilerLibraries(CelExtensions.bindings())
71-
.setOptions(CelOptions.current().populateMacroCalls(true).build())
72-
.addVar("shadowed_ident", SimpleType.INT)
73-
.build();
74-
String expression = "cel.bind(shadowed_ident, 2, shadowed_ident + 1)";
75-
CelAbstractSyntaxTree astToInline = cel.compile(expression).getAst();
76-
CelOptimizer optimizer = CelOptimizerFactory.standardCelOptimizerBuilder(cel)
77-
.addAstOptimizers(InliningOptimizer.newInstance(
78-
InlineVariable.of("shadowed_ident", cel.compile("1").getAst())
79-
))
80-
.build();
80+
private enum SuccessTestCase {
81+
CONSTANT(
82+
"int_var_to_inline + 2 + int_var_to_inline",
83+
"int_var_to_inline",
84+
"1",
85+
"1 + 2 + 1"),
86+
REPEATED(
87+
"a + [a]",
88+
"a",
89+
"dyn([1, 2])",
90+
"dyn([1, 2]) + [dyn([1, 2])]"),
91+
SELECT(
92+
"has(msg.single_any)",
93+
"msg.single_any",
94+
"msg.single_int64_wrapper",
95+
"has(msg.single_int64_wrapper)"),
96+
PRESENCE(
97+
"has(msg.single_int64_wrapper)",
98+
"msg.single_int64_wrapper",
99+
"wrapper_var",
100+
"wrapper_var != null"),
101+
NESTED(
102+
"msg.standalone_message.bb",
103+
"msg.standalone_message",
104+
"child",
105+
"child.bb"),
106+
;
81107

82-
CelAbstractSyntaxTree optimized = optimizer.optimize(astToInline);
108+
private final String source;
109+
private final String varName;
110+
private final String replacement;
111+
private final String expected;
83112

84-
String unparsed = CelUnparserFactory.newUnparser().unparse(optimized);
85-
assertThat(unparsed).isEqualTo("cel.bind(shadowed_ident, 2, shadowed_ident + 1)");
113+
SuccessTestCase(String source, String varName, String replacement, String expected) {
114+
this.source = source;
115+
this.varName = varName;
116+
this.replacement = replacement;
117+
this.expected = expected;
118+
}
119+
}
120+
121+
private enum NoOpTestCase {
122+
NO_INLINE_ITER_VAR(
123+
"[0].exists(shadowed_ident, shadowed_ident == 0)",
124+
"shadowed_ident",
125+
"1"),
126+
NO_INLINE_BIND_VAR(
127+
"cel.bind(shadowed_ident, 2, shadowed_ident + 1)",
128+
"shadowed_ident",
129+
"1"),
130+
;
131+
132+
private final String source;
133+
private final String varName;
134+
private final String replacement;
135+
136+
NoOpTestCase(String source, String varName, String replacement) {
137+
this.source = source;
138+
this.varName = varName;
139+
this.replacement = replacement;
140+
}
86141
}
87142

88143
@Test
89144
public void inline_exceededIterationLimit_throws() throws Exception {
90-
Cel cel = CelFactory.standardCelBuilder()
91-
.setStandardMacros(CelStandardMacro.STANDARD_MACROS)
92-
.addCompilerLibraries(CelExtensions.bindings())
93-
.addVar("ident_to_replace", SimpleType.INT)
94-
.build();
95-
String expression = "ident_to_replace + ident_to_replace + ident_to_replace";
96-
CelAbstractSyntaxTree astToInline = cel.compile(expression).getAst();
97-
CelOptimizer optimizer = CelOptimizerFactory.standardCelOptimizerBuilder(cel)
145+
String expression = "int_var_to_inline + int_var_to_inline + int_var_to_inline";
146+
CelAbstractSyntaxTree astToInline = CEL.compile(expression).getAst();
147+
CelOptimizer optimizer = CelOptimizerFactory.standardCelOptimizerBuilder(CEL)
98148
.addAstOptimizers(InliningOptimizer.newInstance(
99149
InliningOptions.newBuilder().maxIterationLimit(2).build(),
100-
InlineVariable.of("ident_to_replace", cel.compile("1").getAst())
101-
))
150+
InlineVariable.of("int_var_to_inline", CEL.compile("1").getAst())))
102151
.build();
103152

104153
CelOptimizationException e = assertThrows(CelOptimizationException.class, () -> optimizer.optimize(astToInline));
@@ -107,8 +156,8 @@ public void inline_exceededIterationLimit_throws() throws Exception {
107156

108157
@Test
109158
public void inlineVariableDecl_internalVar_throws() throws Exception {
110-
IllegalArgumentException e = assertThrows(IllegalArgumentException.class, () ->
111-
InlineVariable.of("@internal_var", CelAbstractSyntaxTree.newParsedAst(CelExpr.ofNotSet(0L), CelSource.newBuilder().build())));
112-
assertThat(e).hasMessageThat().contains("Internal variables cannot be inlined: @internal_var" );
159+
IllegalArgumentException e = assertThrows(IllegalArgumentException.class, () -> InlineVariable.of("@internal_var",
160+
CelAbstractSyntaxTree.newParsedAst(CelExpr.ofNotSet(0L), CelSource.newBuilder().build())));
161+
assertThat(e).hasMessageThat().contains("Internal variables cannot be inlined: @internal_var");
113162
}
114163
}

0 commit comments

Comments
 (0)