Skip to content

Commit 846e0fc

Browse files
l46kokcopybara-github
authored andcommitted
Add canonicalization for two-variable comprehensions
PiperOrigin-RevId: 961136764
1 parent 29fd271 commit 846e0fc

8 files changed

Lines changed: 1002 additions & 475 deletions

File tree

common/src/main/java/dev/cel/common/navigation/CelNavigableExprUtil.java

Lines changed: 56 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,58 @@
2727
@CheckReturnValue
2828
public final class CelNavigableExprUtil {
2929

30+
/**
31+
* Returns the nearest enclosing comprehension that declares {@code variableName} in scope for
32+
* {@code expr}, or {@code Optional.empty()} if none exists.
33+
*
34+
* <p>A comprehension declares {@code variableName} in scope for {@code expr} if {@code
35+
* variableName} matches {@code iterVar}, {@code iterVar2}, or {@code accuVar}, and {@code expr}
36+
* resides within the branch where that variable is active:
37+
*
38+
* <ul>
39+
* <li>In {@code loopCondition} and {@code loopStep}: {@code iterVar}, {@code iterVar2}, and
40+
* {@code accuVar} are in scope.
41+
* <li>In {@code result}: only {@code accuVar} is in scope.
42+
* <li>In {@code iterRange} and {@code accuInit}: none of the comprehension variables are in
43+
* scope.
44+
* </ul>
45+
*/
46+
@SuppressWarnings("ReferenceEquality") // Disambiguates mutable child branches
47+
public static <E extends Expression, T extends BaseNavigableExpr<E>>
48+
Optional<T> findDeclaringComprehension(T expr, String variableName) {
49+
checkNotNull(expr);
50+
checkNotNull(variableName);
51+
if (variableName.isEmpty()) {
52+
return Optional.empty();
53+
}
54+
T curr = expr;
55+
Optional<T> maybeParent = curr.parent();
56+
while (maybeParent.isPresent()) {
57+
T parent = maybeParent.get();
58+
if (parent.getKind() == Kind.COMPREHENSION) {
59+
Expression.Comprehension<?> comp = parent.expr().comprehension();
60+
Expression currExpr = curr.expr();
61+
62+
if (currExpr != comp.iterRange() && currExpr != comp.accuInit()) {
63+
if (currExpr == comp.result()) {
64+
if (comp.accuVar().equals(variableName)) {
65+
return Optional.of(parent);
66+
}
67+
} else {
68+
if (comp.iterVar().equals(variableName)
69+
|| comp.iterVar2().equals(variableName)
70+
|| comp.accuVar().equals(variableName)) {
71+
return Optional.of(parent);
72+
}
73+
}
74+
}
75+
}
76+
curr = parent;
77+
maybeParent = parent.parent();
78+
}
79+
return Optional.empty();
80+
}
81+
3082
/**
3183
* Returns true if {@code variableName} is in scope and shadowed by an enclosing comprehension
3284
* above {@code expr}.
@@ -56,7 +108,7 @@ public final class CelNavigableExprUtil {
56108
* </ul>
57109
*/
58110
public static boolean isVariableShadowed(BaseNavigableExpr<?> expr, String variableName) {
59-
return areVariablesShadowed(expr, Collections.singleton(variableName));
111+
return findDeclaringComprehension(expr, variableName).isPresent();
60112
}
61113

62114
/**
@@ -72,38 +124,14 @@ public static boolean isVariableShadowed(BaseNavigableExpr<?> expr, String varia
72124
* At {@code y > 0}, {@code areVariablesShadowed(node, ImmutableSet.of("x", "z"))} is {@code true}
73125
* because {@code x} is in scope from the outer comprehension.
74126
*/
75-
@SuppressWarnings("ReferenceEquality") // Required to disambiguate child branches
76127
public static boolean areVariablesShadowed(
77128
BaseNavigableExpr<?> expr, Collection<String> variableNames) {
78129
checkNotNull(expr);
79130
checkNotNull(variableNames);
80-
if (variableNames.isEmpty()) {
81-
return false;
82-
}
83-
BaseNavigableExpr<?> curr = expr;
84-
Optional<? extends BaseNavigableExpr<?>> maybeParent = curr.parent();
85-
while (maybeParent.isPresent()) {
86-
BaseNavigableExpr<?> parent = maybeParent.get();
87-
if (parent.getKind() == Kind.COMPREHENSION) {
88-
Expression.Comprehension<?> comp = parent.expr().comprehension();
89-
Expression currExpr = curr.expr();
90-
91-
if (currExpr != comp.iterRange() && currExpr != comp.accuInit()) {
92-
if (currExpr == comp.result()) {
93-
if (variableNames.contains(comp.accuVar())) {
94-
return true;
95-
}
96-
} else {
97-
if (variableNames.contains(comp.iterVar())
98-
|| variableNames.contains(comp.iterVar2())
99-
|| variableNames.contains(comp.accuVar())) {
100-
return true;
101-
}
102-
}
103-
}
131+
for (String varName : variableNames) {
132+
if (findDeclaringComprehension(expr, varName).isPresent()) {
133+
return true;
104134
}
105-
curr = parent;
106-
maybeParent = parent.parent();
107135
}
108136
return false;
109137
}

common/src/test/java/dev/cel/common/navigation/CelNavigableExprUtilTest.java

Lines changed: 130 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -195,6 +195,76 @@ public void isVariableShadowed_comprehensionResultBranch() throws Exception {
195195
.isFalse();
196196
}
197197

198+
@Test
199+
public void isVariableShadowed_twoVarComprehension_resultBranch() throws Exception {
200+
CelAbstractSyntaxTree ast =
201+
COMPILER.compile("{'k1': 1, 'k2': 2}.all(k, v, k != '' && v > 0)").getAst();
202+
CelNavigableMutableAst navigableAst =
203+
CelNavigableMutableAst.fromAst(CelMutableAst.fromCelAst(ast));
204+
205+
CelNavigableMutableExpr comprehensionNode =
206+
navigableAst
207+
.getRoot()
208+
.allNodes()
209+
.filter(node -> node.getKind() == Kind.COMPREHENSION)
210+
.findFirst()
211+
.get();
212+
213+
CelMutableComprehension comprehension = comprehensionNode.expr().comprehension();
214+
long resultId = comprehension.result().id();
215+
216+
CelNavigableMutableExpr resultNode =
217+
comprehensionNode.allNodes().filter(node -> node.id() == resultId).findFirst().get();
218+
219+
// In result branch of two-var comprehension: accuVar is in scope, but iterVar and iterVar2 are
220+
// not
221+
assertThat(CelNavigableExprUtil.isVariableShadowed(resultNode, comprehension.accuVar()))
222+
.isTrue();
223+
assertThat(CelNavigableExprUtil.isVariableShadowed(resultNode, comprehension.iterVar()))
224+
.isFalse();
225+
assertThat(CelNavigableExprUtil.isVariableShadowed(resultNode, comprehension.iterVar2()))
226+
.isFalse();
227+
}
228+
229+
@Test
230+
public void isVariableShadowed_accuInit_notShadowed() {
231+
CelMutableExpr iterRange = CelMutableExpr.ofList(0, CelMutableList.create());
232+
CelMutableExpr accuInitIdent = CelMutableExpr.ofIdent(1, "x");
233+
CelMutableExpr loopCond = CelMutableExpr.ofConstant(2, CelConstant.ofValue(true));
234+
CelMutableExpr loopStep = CelMutableExpr.ofConstant(3, CelConstant.ofValue(true));
235+
CelMutableExpr result = CelMutableExpr.ofIdent(4, "accu");
236+
237+
CelMutableExpr comp =
238+
CelMutableExpr.ofComprehension(
239+
5,
240+
CelMutableComprehension.create(
241+
"x", iterRange, "accu", accuInitIdent, loopCond, loopStep, result));
242+
243+
CelNavigableMutableExpr root = CelNavigableMutableExpr.fromExpr(comp);
244+
CelNavigableMutableExpr navAccuInit =
245+
root.allNodes().filter(node -> node.id() == 1).findFirst().get();
246+
247+
assertThat(CelNavigableExprUtil.isVariableShadowed(navAccuInit, "x")).isFalse();
248+
assertThat(CelNavigableExprUtil.isVariableShadowed(navAccuInit, "accu")).isFalse();
249+
}
250+
251+
@Test
252+
public void findDeclaringComprehension_emptyVariableName_returnsEmpty() throws Exception {
253+
CelAbstractSyntaxTree ast = COMPILER.compile("[1, 2].all(x, x > 0)").getAst();
254+
CelNavigableAst navigableAst = CelNavigableAst.fromAst(ast);
255+
256+
CelNavigableExpr identX =
257+
navigableAst
258+
.getRoot()
259+
.allNodes()
260+
.filter(node -> node.expr().identOrDefault().name().equals("x"))
261+
.findFirst()
262+
.get();
263+
264+
assertThat(CelNavigableExprUtil.findDeclaringComprehension(identX, "")).isEmpty();
265+
assertThat(CelNavigableExprUtil.isVariableShadowed(identX, "")).isFalse();
266+
}
267+
198268
@Test
199269
public void areVariablesShadowed_multipleVariables() throws Exception {
200270
CelAbstractSyntaxTree ast = COMPILER.compile("[1, 2].all(x, x > 0)").getAst();
@@ -345,4 +415,64 @@ public void isVariableShadowed_zeroedOutIds_scopedCorrectly() {
345415
assertThat(CelNavigableExprUtil.isVariableShadowed(navResult, "x")).isFalse();
346416
assertThat(CelNavigableExprUtil.isVariableShadowed(navResult, "accu")).isTrue();
347417
}
418+
419+
@Test
420+
public void
421+
findDeclaringComprehension_nestedComprehensions_resolvesToInnermostDeclaringComprehension()
422+
throws Exception {
423+
CelAbstractSyntaxTree ast =
424+
COMPILER
425+
.compile("[1, 2].all(x, {'k': 1}.exists(k, v, x > 0 && k != '' && v > 0))")
426+
.getAst();
427+
CelNavigableAst navigableAst = CelNavigableAst.fromAst(ast);
428+
429+
CelNavigableExpr outerComp =
430+
navigableAst
431+
.getRoot()
432+
.allNodes()
433+
.filter(
434+
node ->
435+
node.getKind() == Kind.COMPREHENSION
436+
&& node.expr().comprehension().iterVar().equals("x"))
437+
.findFirst()
438+
.get();
439+
440+
CelNavigableExpr innerComp =
441+
navigableAst
442+
.getRoot()
443+
.allNodes()
444+
.filter(
445+
node ->
446+
node.getKind() == Kind.COMPREHENSION
447+
&& node.expr().comprehension().iterVar().equals("k"))
448+
.findFirst()
449+
.get();
450+
451+
CelNavigableExpr identX =
452+
navigableAst
453+
.getRoot()
454+
.allNodes()
455+
.filter(node -> node.expr().identOrDefault().name().equals("x"))
456+
.findFirst()
457+
.get();
458+
CelNavigableExpr identK =
459+
navigableAst
460+
.getRoot()
461+
.allNodes()
462+
.filter(node -> node.expr().identOrDefault().name().equals("k"))
463+
.findFirst()
464+
.get();
465+
CelNavigableExpr identV =
466+
navigableAst
467+
.getRoot()
468+
.allNodes()
469+
.filter(node -> node.expr().identOrDefault().name().equals("v"))
470+
.findFirst()
471+
.get();
472+
473+
assertThat(CelNavigableExprUtil.findDeclaringComprehension(identX, "x")).hasValue(outerComp);
474+
assertThat(CelNavigableExprUtil.findDeclaringComprehension(identK, "k")).hasValue(innerComp);
475+
assertThat(CelNavigableExprUtil.findDeclaringComprehension(identV, "v")).hasValue(innerComp);
476+
assertThat(CelNavigableExprUtil.findDeclaringComprehension(identX, "unknown")).isEmpty();
477+
}
348478
}

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -140,11 +140,13 @@ java_library(
140140
"//common/ast",
141141
"//common/ast:mutable_expr",
142142
"//common/navigation:common",
143+
"//common/navigation:expr_util",
143144
"//common/navigation:mutable_navigation",
144145
"//common/values:cel_byte_string",
145146
"//optimizer:ast_optimizer",
146147
"//optimizer:mutable_ast",
147148
"@maven//:com_google_guava_guava",
149+
"@maven//:org_jspecify_jspecify",
148150
],
149151
)
150152

0 commit comments

Comments
 (0)