Skip to content

Commit 7271bc5

Browse files
l46kokcopybara-github
authored andcommitted
Prevent ConstantFoldingOptimizer to fold x in [x] for dyn/double typed variables
PiperOrigin-RevId: 957324884
1 parent 8d150b2 commit 7271bc5

3 files changed

Lines changed: 117 additions & 12 deletions

File tree

optimizer/src/main/java/dev/cel/optimizer/optimizers/ConstantFoldingOptimizer.java

Lines changed: 76 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -46,6 +46,8 @@
4646
import dev.cel.common.navigation.TraversalOrder;
4747
import dev.cel.common.types.CelType;
4848
import dev.cel.common.types.CelTypeProvider;
49+
import dev.cel.common.types.NullableType;
50+
import dev.cel.common.types.OptionalType;
4951
import dev.cel.common.types.SimpleType;
5052
import dev.cel.common.types.StructType;
5153
import dev.cel.common.values.CelValue;
@@ -541,16 +543,37 @@ private Optional<CelMutableAst> maybePruneBranches(
541543

542544
CelMutableExpr needle = call.args().get(0);
543545
if (needle.getKind().equals(Kind.CONSTANT) || needle.getKind().equals(Kind.IDENT)) {
544-
Object needleValue =
545-
needle.getKind().equals(Kind.CONSTANT) ? needle.constant() : needle.ident();
546546
for (CelMutableExpr elem : haystack.elements()) {
547-
if ((elem.getKind().equals(Kind.CONSTANT) && elem.constant().equals(needleValue))
548-
|| (elem.getKind().equals(Kind.IDENT) && elem.ident().equals(needleValue))) {
549-
return Optional.of(
550-
astMutator.replaceSubtree(
551-
mutableAst.expr(),
552-
CelMutableExpr.ofConstant(CelConstant.ofValue(true)),
553-
expr.id()));
547+
if ((elem.getKind().equals(Kind.CONSTANT)
548+
&& needle.getKind().equals(Kind.CONSTANT)
549+
&& elem.constant().equals(needle.constant()))
550+
|| (elem.getKind().equals(Kind.IDENT)
551+
&& needle.getKind().equals(Kind.IDENT)
552+
&& elem.ident().equals(needle.ident()))) {
553+
if (needle.getKind().equals(Kind.CONSTANT)) {
554+
if (needle.constant().getKind().equals(CelConstant.Kind.DOUBLE_VALUE)
555+
&& Double.isNaN(needle.constant().doubleValue())) {
556+
continue;
557+
}
558+
return Optional.of(
559+
astMutator.replaceSubtree(
560+
mutableAst.expr(),
561+
CelMutableExpr.ofConstant(CelConstant.ofValue(true)),
562+
expr.id()));
563+
}
564+
565+
Optional<CelType> needleType = mutableAst.getType(needle.id());
566+
if (!needleType.isPresent()) {
567+
needleType = Optional.ofNullable(identTypes.get(needle.ident().name()));
568+
}
569+
570+
if (needleType.isPresent() && isSafeForExactEquality(needleType.get())) {
571+
return Optional.of(
572+
astMutator.replaceSubtree(
573+
mutableAst.expr(),
574+
CelMutableExpr.ofConstant(CelConstant.ofValue(true)),
575+
expr.id()));
576+
}
554577
}
555578
}
556579
}
@@ -948,6 +971,50 @@ private static boolean isExprConstantOfKind(CelMutableExpr expr, CelConstant.Kin
948971
return expr.getKind().equals(Kind.CONSTANT) && expr.constant().getKind().equals(constantKind);
949972
}
950973

974+
private static boolean isSafeForExactEquality(@Nullable CelType celType) {
975+
if (celType == null) {
976+
return false;
977+
}
978+
979+
if (celType instanceof NullableType) {
980+
return isSafeForExactEquality(((NullableType) celType).targetType());
981+
}
982+
983+
switch (celType.kind()) {
984+
case BOOL:
985+
case INT:
986+
case UINT:
987+
case STRING:
988+
case BYTES:
989+
case DURATION:
990+
case TIMESTAMP:
991+
case NULL_TYPE:
992+
case TYPE:
993+
return true;
994+
995+
case LIST:
996+
return !celType.parameters().isEmpty()
997+
&& isSafeForExactEquality(celType.parameters().get(0));
998+
999+
case MAP:
1000+
return celType.parameters().size() >= 2
1001+
&& isSafeForExactEquality(celType.parameters().get(0))
1002+
&& isSafeForExactEquality(celType.parameters().get(1));
1003+
1004+
case OPAQUE:
1005+
if ((celType instanceof OptionalType
1006+
|| celType.name().equals(OptionalType.NAME)
1007+
|| celType.name().equals("optional"))
1008+
&& !celType.parameters().isEmpty()) {
1009+
return isSafeForExactEquality(celType.parameters().get(0));
1010+
}
1011+
return false;
1012+
1013+
default:
1014+
return false;
1015+
}
1016+
}
1017+
9511018
private ConstantFoldingOptimizer(ConstantFoldingOptions constantFoldingOptions) {
9521019
this.constantFoldingOptions = constantFoldingOptions;
9531020
this.astMutator = AstMutator.newInstance(constantFoldingOptions.maxIterationLimit());

optimizer/src/test/java/dev/cel/optimizer/optimizers/ConstantFoldingOptimizerTest.java

Lines changed: 38 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,7 @@
3232
import dev.cel.common.CelOverloadDecl;
3333
import dev.cel.common.types.ListType;
3434
import dev.cel.common.types.MapType;
35+
import dev.cel.common.types.OptionalType;
3536
import dev.cel.common.types.SimpleType;
3637
import dev.cel.common.types.StructTypeReference;
3738
import dev.cel.expr.conformance.proto2.TestAllTypes.NestedMessage;
@@ -80,6 +81,18 @@ private static Cel setupEnv(CelBuilder celBuilder) {
8081
return celBuilder
8182
.addVar("x", SimpleType.DYN)
8283
.addVar("y", SimpleType.DYN)
84+
.addVar("dyn_x", SimpleType.DYN)
85+
.addVar("int_x", SimpleType.INT)
86+
.addVar("double_x", SimpleType.DOUBLE)
87+
.addVar("bool_x", SimpleType.BOOL)
88+
.addVar("string_x", SimpleType.STRING)
89+
.addVar("int_list_x", ListType.create(SimpleType.INT))
90+
.addVar("double_list_x", ListType.create(SimpleType.DOUBLE))
91+
.addVar("dyn_list_x", ListType.create(SimpleType.DYN))
92+
.addVar("map_string_int_x", MapType.create(SimpleType.STRING, SimpleType.INT))
93+
.addVar("map_string_double_x", MapType.create(SimpleType.STRING, SimpleType.DOUBLE))
94+
.addVar("optional_int_x", OptionalType.create(SimpleType.INT))
95+
.addVar("optional_double_x", OptionalType.create(SimpleType.DOUBLE))
8396
.addVar("bool_var", SimpleType.BOOL)
8497
.addVar("list_var", ListType.create(SimpleType.STRING))
8598
.addVar("map_var", MapType.create(SimpleType.STRING, SimpleType.STRING))
@@ -151,7 +164,29 @@ private static Cel setupEnv(CelBuilder celBuilder) {
151164
@TestParameters("{source: '5 in [1, 1 + 2, 1 + (2 + 3)]', expected: 'false'}")
152165
@TestParameters("{source: '5 in [1, x, y, 5]', expected: 'true'}")
153166
@TestParameters("{source: '!(5 in [1, x, y, 5])', expected: 'false'}")
154-
@TestParameters("{source: 'x in [1, x, y, 5]', expected: 'true'}")
167+
@TestParameters("{source: 'x in [1, x, y, 5]', expected: 'x in [1, x, y, 5]'}")
168+
@TestParameters("{source: 'dyn_x in [1, 2, dyn_x]', expected: 'dyn_x in [1, 2, dyn_x]'}")
169+
@TestParameters("{source: 'int_x in [1, 2, int_x]', expected: 'true'}")
170+
@TestParameters("{source: 'bool_x in [true, false, bool_x]', expected: 'true'}")
171+
@TestParameters("{source: 'string_x in [\"a\", \"b\", string_x]', expected: 'true'}")
172+
@TestParameters(
173+
"{source: 'double_x in [1.0, 2.0, double_x]', expected: 'double_x in [1.0, 2.0, double_x]'}")
174+
@TestParameters("{source: 'int_list_x in [[1], [2], int_list_x]', expected: 'true'}")
175+
@TestParameters(
176+
"{source: 'double_list_x in [[1.0], double_list_x]', expected: 'double_list_x in [[1.0],"
177+
+ " double_list_x]'}")
178+
@TestParameters(
179+
"{source: 'dyn_list_x in [[1], dyn_list_x]', expected: 'dyn_list_x in [[1], dyn_list_x]'}")
180+
@TestParameters(
181+
"{source: 'map_string_int_x in [{\"a\": 1}, map_string_int_x]', expected: 'true'}")
182+
@TestParameters(
183+
"{source: 'map_string_double_x in [{\"a\": 1.0}, map_string_double_x]', expected:"
184+
+ " 'map_string_double_x in [{\"a\": 1.0}, map_string_double_x]'}")
185+
@TestParameters(
186+
"{source: 'optional_int_x in [optional.of(1), optional_int_x]', expected: 'true'}")
187+
@TestParameters(
188+
"{source: 'optional_double_x in [optional.of(1.0), optional_double_x]', expected:"
189+
+ " 'optional_double_x in [optional.of(1.0), optional_double_x]'}")
155190
@TestParameters("{source: 'x in [1, 1 + 2, 1 + (2 + 3)]', expected: 'x in [1, 3, 6]'}")
156191
@TestParameters("{source: 'duration(string(7 * 24) + ''h'')', expected: 'duration(\"168h\")'}")
157192
@TestParameters("{source: '[1, ?optional.of(3)]', expected: '[1, 3]'}")
@@ -395,7 +430,8 @@ public void constantFold_protoMessageLiteral_success(String source, String expec
395430
@TestParameters(
396431
"{source: 'cel.bind(myMap, {\"foo\": \"bar\"}, myMap[?\"foo\"].optMap(x, x + \"baz\"))', "
397432
+ "expected: 'optional.of(\"barbaz\")'}")
398-
@TestParameters("{source: '(1 + 2 + 3 == x) && (x in [1, 2, x])', expected: '6 == x'}")
433+
@TestParameters(
434+
"{source: '(1 + 2 + 3 == x) && (x in [1, 2, x])', expected: '6 == x && x in [1, 2, x]'}")
399435
public void constantFold_macros_macroCallMetadataPopulated(String source, String expected)
400436
throws Exception {
401437
Cel cel =

verifier/src/test/java/dev/cel/verifier/CelVerifierZ3ImplTest.java

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1683,7 +1683,8 @@ private enum EquivalenceTestCase {
16831683
JSON_VALUE_OPTIONAL_NULL_VALUE_OF("google.protobuf.Value{?null_value: optional.of(0)}", "null"),
16841684
OPTIONAL_INDEX_LIST_UNWRAPPING("optional.of([1, 2, 3])[?0]", "optional.of(1)"),
16851685
OPTIONAL_INDEX_MAP_UNWRAPPING("optional.of({'a': 1})[?'a']", "optional.of(1)"),
1686-
OPTIONAL_INDEX_UNWRAPPING_NONE("optional.none()[?0]", "optional.none()");
1686+
OPTIONAL_INDEX_UNWRAPPING_NONE("optional.none()[?0]", "optional.none()"),
1687+
INT_IN_LIST_IDENTITY_EQUIVALENT("x in [1, 2, x]", "true");
16871688

16881689
private final String exprA;
16891690
private final String exprB;
@@ -1729,6 +1730,7 @@ private enum EquivalenceViolationTestCase {
17291730
OPTIONAL_VALUE_VIOLATION("optional.of(x).value()", "y"),
17301731
LIST_OPTIONAL_ELEMENTS_COLLISION("[1, ?opt_var]", "[1, opt_var]"),
17311732
CROSS_NUMERIC_EQUALITY_INT_DYN_VIOLATION("1 == request", "false"),
1733+
DYN_IN_LIST_NOT_EQUIVALENT_TO_TRUE("dyn_var in [1, 2, dyn_var]", "true"),
17321734
OPTIONAL_SELECTION_VS_DIRECT_ERROR(
17331735
"{'a': 1}.?missing_key", "optional.of({'a': 1}.missing_key)"),
17341736
OPTIONAL_NESTED_NONE_VS_FLAT_NONE("{'a': optional.none()}.?a", "optional.none()"),

0 commit comments

Comments
 (0)