Skip to content

Commit 2dd8da3

Browse files
l46kokcopybara-github
authored andcommitted
Add a validation pass for ID uniqueness in optimizers
PiperOrigin-RevId: 963040423
1 parent 94ba8ad commit 2dd8da3

3 files changed

Lines changed: 243 additions & 1 deletion

File tree

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -57,6 +57,8 @@ java_library(
5757
"//bundle:cel",
5858
"//common:cel_ast",
5959
"//common:compiler_common",
60+
"//common/ast",
61+
"//common/navigation",
6062
"@maven//:com_google_guava_guava",
6163
],
6264
)

optimizer/src/main/java/dev/cel/optimizer/CelOptimizerImpl.java

Lines changed: 58 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,8 +20,14 @@
2020
import dev.cel.bundle.Cel;
2121
import dev.cel.common.CelAbstractSyntaxTree;
2222
import dev.cel.common.CelValidationException;
23+
import dev.cel.common.ast.CelExpr;
24+
import dev.cel.common.ast.CelExpr.ExprKind.Kind;
25+
import dev.cel.common.navigation.CelNavigableAst;
26+
import dev.cel.common.navigation.CelNavigableExpr;
2327
import dev.cel.optimizer.CelAstOptimizer.OptimizationResult;
2428
import java.util.Arrays;
29+
import java.util.HashMap;
30+
import java.util.Map;
2531

2632
final class CelOptimizerImpl implements CelOptimizer {
2733
private final Cel cel;
@@ -52,6 +58,7 @@ public CelAbstractSyntaxTree optimize(CelAbstractSyntaxTree ast) throws CelOptim
5258
.build();
5359
}
5460
optimizedAst = celOptimizerEnv.check(result.optimizedAst()).getAst();
61+
assertAstIdCorrectness(optimizedAst);
5562
}
5663
} catch (CelValidationException e) {
5764
throw new CelOptimizationException(
@@ -63,6 +70,57 @@ public CelAbstractSyntaxTree optimize(CelAbstractSyntaxTree ast) throws CelOptim
6370
return optimizedAst;
6471
}
6572

73+
private static void assertAstIdCorrectness(CelAbstractSyntaxTree ast) {
74+
Map<Long, CelExpr> allExprs = new HashMap<>();
75+
CelNavigableAst.fromAst(ast)
76+
.getRoot()
77+
.allNodes()
78+
.forEach(
79+
navExpr -> {
80+
CelExpr expr = navExpr.expr();
81+
CelExpr existing = allExprs.put(expr.id(), expr);
82+
if (existing != null) {
83+
throw new IllegalStateException(
84+
String.format("Duplicate expr ID %d detected in the AST.", expr.id()));
85+
}
86+
});
87+
88+
for (CelExpr macroCall : ast.getSource().getMacroCalls().values()) {
89+
if (macroCall.id() != 0) {
90+
throw new IllegalStateException(
91+
String.format("Expected macro call root ID to be 0, but was %d.", macroCall.id()));
92+
}
93+
CelNavigableExpr.fromExpr(macroCall)
94+
.descendants()
95+
.forEach(
96+
navExpr -> {
97+
CelExpr node = navExpr.expr();
98+
CelExpr e = allExprs.get(node.id());
99+
// A node may not exist in the AST if it is a synthetic macro node or was eliminated
100+
// during optimization passes.
101+
if (e == null) {
102+
return;
103+
}
104+
105+
if (e.exprKind().getKind().equals(Kind.COMPREHENSION)) {
106+
if (!node.exprKind().getKind().equals(Kind.NOT_SET)) {
107+
throw new IllegalStateException(
108+
String.format(
109+
"Expected macro call node %d to be NOT_SET for comprehension, but"
110+
+ " was %s.",
111+
node.id(), node.exprKind().getKind()));
112+
}
113+
} else if (!node.exprKind().getKind().equals(e.exprKind().getKind())) {
114+
throw new IllegalStateException(
115+
String.format(
116+
"Macro call node %d kind mismatch: expected %s (from AST), but was %s"
117+
+ " (in macro call).",
118+
node.id(), e.exprKind().getKind(), node.exprKind().getKind()));
119+
}
120+
});
121+
}
122+
}
123+
66124
/** Create a new builder for constructing a {@link CelOptimizer} instance. */
67125
static CelOptimizerImpl.Builder newBuilder(Cel cel) {
68126
return new CelOptimizerImpl.Builder(cel);

optimizer/src/test/java/dev/cel/optimizer/CelOptimizerImplTest.java

Lines changed: 183 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,23 +17,32 @@
1717
import static com.google.common.truth.Truth.assertThat;
1818
import static org.junit.Assert.assertThrows;
1919

20+
import com.google.common.collect.ImmutableList;
2021
import dev.cel.bundle.Cel;
2122
import dev.cel.bundle.CelFactory;
2223
import dev.cel.common.CelAbstractSyntaxTree;
24+
import dev.cel.common.CelOptions;
2325
import dev.cel.common.CelSource;
2426
import dev.cel.common.CelValidationException;
27+
import dev.cel.common.ast.CelConstant;
2528
import dev.cel.common.ast.CelExpr;
2629
import dev.cel.optimizer.CelAstOptimizer.OptimizationResult;
30+
import dev.cel.parser.CelStandardMacro;
2731
import java.util.ArrayList;
2832
import java.util.List;
33+
import java.util.Optional;
2934
import org.junit.Test;
3035
import org.junit.runner.RunWith;
3136
import org.junit.runners.JUnit4;
3237

3338
@RunWith(JUnit4.class)
3439
public class CelOptimizerImplTest {
3540

36-
private static final Cel CEL = CelFactory.standardCelBuilder().build();
41+
private static final Cel CEL =
42+
CelFactory.standardCelBuilder()
43+
.setOptions(CelOptions.current().populateMacroCalls(true).build())
44+
.setStandardMacros(CelStandardMacro.STANDARD_MACROS)
45+
.build();
3746

3847
@Test
3948
public void constructCelOptimizer_success() {
@@ -131,4 +140,177 @@ public void optimizedAst_failsToTypeCheck_throwsException() {
131140
+ " 'undeclared_ident' (in container '')");
132141
assertThat(e).hasCauseThat().isInstanceOf(CelValidationException.class);
133142
}
143+
144+
@Test
145+
public void optimize_duplicateExprId_throwsException() {
146+
CelOptimizer celOptimizer =
147+
CelOptimizerImpl.newBuilder(CEL)
148+
.addAstOptimizers(
149+
(navigableAst, cel) ->
150+
OptimizationResult.create(
151+
CelAbstractSyntaxTree.newParsedAst(
152+
CelExpr.ofCall(
153+
1,
154+
Optional.empty(),
155+
"_+_",
156+
ImmutableList.of(
157+
CelExpr.ofConstant(1, CelConstant.ofValue(1L)),
158+
CelExpr.ofConstant(2, CelConstant.ofValue(2L)))),
159+
CelSource.newBuilder().build())))
160+
.build();
161+
162+
CelOptimizationException e =
163+
assertThrows(
164+
CelOptimizationException.class,
165+
() -> celOptimizer.optimize(CEL.compile("1 + 2").getAst()));
166+
167+
assertThat(e)
168+
.hasMessageThat()
169+
.isEqualTo("Optimization failure: Duplicate expr ID 1 detected in the AST.");
170+
assertThat(e).hasCauseThat().isInstanceOf(IllegalStateException.class);
171+
}
172+
173+
@Test
174+
public void optimize_macroCallRootIdNonZero_throwsException() {
175+
CelOptimizer celOptimizer =
176+
CelOptimizerImpl.newBuilder(CEL)
177+
.addAstOptimizers(
178+
(navigableAst, cel) ->
179+
OptimizationResult.create(
180+
CelAbstractSyntaxTree.newParsedAst(
181+
CelExpr.ofConstant(1, CelConstant.ofValue(1L)),
182+
CelSource.newBuilder()
183+
.addMacroCalls(
184+
1L,
185+
CelExpr.ofCall(
186+
10L,
187+
Optional.empty(),
188+
"has",
189+
ImmutableList.of(
190+
CelExpr.ofConstant(1L, CelConstant.ofValue(1L)))))
191+
.build())))
192+
.build();
193+
194+
CelOptimizationException e =
195+
assertThrows(
196+
CelOptimizationException.class, () -> celOptimizer.optimize(CEL.compile("1").getAst()));
197+
198+
assertThat(e)
199+
.hasMessageThat()
200+
.isEqualTo("Optimization failure: Expected macro call root ID to be 0, but was 10.");
201+
assertThat(e).hasCauseThat().isInstanceOf(IllegalStateException.class);
202+
}
203+
204+
@Test
205+
public void optimize_macroCallKindMismatch_throwsException() {
206+
CelOptimizer celOptimizer =
207+
CelOptimizerImpl.newBuilder(CEL)
208+
.addAstOptimizers(
209+
(navigableAst, cel) ->
210+
OptimizationResult.create(
211+
CelAbstractSyntaxTree.newParsedAst(
212+
CelExpr.ofConstant(1, CelConstant.ofValue(1L)),
213+
CelSource.newBuilder()
214+
.addMacroCalls(
215+
1L,
216+
CelExpr.ofCall(
217+
0L,
218+
Optional.empty(),
219+
"has",
220+
ImmutableList.of(CelExpr.ofIdent(1L, "x"))))
221+
.build())))
222+
.build();
223+
224+
CelOptimizationException e =
225+
assertThrows(
226+
CelOptimizationException.class, () -> celOptimizer.optimize(CEL.compile("1").getAst()));
227+
228+
assertThat(e)
229+
.hasMessageThat()
230+
.isEqualTo(
231+
"Optimization failure: Macro call node 1 kind mismatch: expected CONSTANT (from AST),"
232+
+ " but was IDENT (in macro call).");
233+
assertThat(e).hasCauseThat().isInstanceOf(IllegalStateException.class);
234+
}
235+
236+
@Test
237+
public void optimize_macroCallComprehensionKindNotSetMismatch_throwsException() throws Exception {
238+
CelAbstractSyntaxTree astWithComprehension = CEL.compile("[1].all(x, x > 0)").getAst();
239+
long compId = astWithComprehension.getExpr().id();
240+
241+
CelOptimizer celOptimizer =
242+
CelOptimizerImpl.newBuilder(CEL)
243+
.addAstOptimizers(
244+
(navigableAst, cel) ->
245+
OptimizationResult.create(
246+
CelAbstractSyntaxTree.newParsedAst(
247+
astWithComprehension.getExpr(),
248+
CelSource.newBuilder()
249+
.addMacroCalls(
250+
compId,
251+
CelExpr.ofCall(
252+
0L,
253+
Optional.empty(),
254+
"all",
255+
ImmutableList.of(
256+
CelExpr.ofIdent(compId, "not_set_expected"))))
257+
.build())))
258+
.build();
259+
260+
CelOptimizationException e =
261+
assertThrows(
262+
CelOptimizationException.class, () -> celOptimizer.optimize(astWithComprehension));
263+
264+
assertThat(e)
265+
.hasMessageThat()
266+
.isEqualTo(
267+
String.format(
268+
"Optimization failure: Expected macro call node %d to be NOT_SET for comprehension,"
269+
+ " but was IDENT.",
270+
compId));
271+
assertThat(e).hasCauseThat().isInstanceOf(IllegalStateException.class);
272+
}
273+
274+
@Test
275+
public void optimize_macroCallComprehensionKindNotSet_success() throws Exception {
276+
CelAbstractSyntaxTree astWithComprehension = CEL.compile("[1].all(x, x > 0)").getAst();
277+
long compId = astWithComprehension.getExpr().id();
278+
279+
CelOptimizer celOptimizer =
280+
CelOptimizerImpl.newBuilder(CEL)
281+
.addAstOptimizers(
282+
(navigableAst, cel) ->
283+
OptimizationResult.create(
284+
CelAbstractSyntaxTree.newParsedAst(
285+
astWithComprehension.getExpr(),
286+
CelSource.newBuilder()
287+
.addMacroCalls(
288+
compId,
289+
CelExpr.ofCall(
290+
0L,
291+
Optional.empty(),
292+
"all",
293+
ImmutableList.of(CelExpr.ofNotSet(compId))))
294+
.build())))
295+
.build();
296+
297+
CelAbstractSyntaxTree optimizedAst = celOptimizer.optimize(astWithComprehension);
298+
299+
assertThat(optimizedAst).isNotNull();
300+
}
301+
302+
@Test
303+
public void optimize_validMacroCalls_success() throws Exception {
304+
CelAbstractSyntaxTree ast = CEL.compile("[1, 2, 3].all(x, x > 0)").getAst();
305+
306+
CelOptimizer celOptimizer =
307+
CelOptimizerImpl.newBuilder(CEL)
308+
.addAstOptimizers((navigableAst, cel) -> OptimizationResult.create(navigableAst))
309+
.build();
310+
311+
CelAbstractSyntaxTree optimizedAst = celOptimizer.optimize(ast);
312+
313+
assertThat(optimizedAst).isNotNull();
314+
assertThat(optimizedAst.getSource().getMacroCalls()).hasSize(1);
315+
}
134316
}

0 commit comments

Comments
 (0)