Skip to content

Commit 702fc81

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 702fc81

3 files changed

Lines changed: 154 additions & 12 deletions

File tree

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

Lines changed: 67 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -46,6 +46,7 @@
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.OptionalType;
4950
import dev.cel.common.types.SimpleType;
5051
import dev.cel.common.types.StructType;
5152
import dev.cel.common.values.CelValue;
@@ -541,16 +542,37 @@ private Optional<CelMutableAst> maybePruneBranches(
541542

542543
CelMutableExpr needle = call.args().get(0);
543544
if (needle.getKind().equals(Kind.CONSTANT) || needle.getKind().equals(Kind.IDENT)) {
544-
Object needleValue =
545-
needle.getKind().equals(Kind.CONSTANT) ? needle.constant() : needle.ident();
546545
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()));
546+
if ((elem.getKind().equals(Kind.CONSTANT)
547+
&& needle.getKind().equals(Kind.CONSTANT)
548+
&& elem.constant().equals(needle.constant()))
549+
|| (elem.getKind().equals(Kind.IDENT)
550+
&& needle.getKind().equals(Kind.IDENT)
551+
&& elem.ident().equals(needle.ident()))) {
552+
if (needle.getKind().equals(Kind.CONSTANT)) {
553+
if (needle.constant().getKind().equals(CelConstant.Kind.DOUBLE_VALUE)
554+
&& Double.isNaN(needle.constant().doubleValue())) {
555+
continue;
556+
}
557+
return Optional.of(
558+
astMutator.replaceSubtree(
559+
mutableAst.expr(),
560+
CelMutableExpr.ofConstant(CelConstant.ofValue(true)),
561+
expr.id()));
562+
}
563+
564+
CelType needleType =
565+
mutableAst
566+
.getType(needle.id())
567+
.orElseGet(() -> identTypes.get(needle.ident().name()));
568+
569+
if (needleType != null && isSafeForExactEquality(needleType)) {
570+
return Optional.of(
571+
astMutator.replaceSubtree(
572+
mutableAst.expr(),
573+
CelMutableExpr.ofConstant(CelConstant.ofValue(true)),
574+
expr.id()));
575+
}
554576
}
555577
}
556578
}
@@ -948,6 +970,42 @@ private static boolean isExprConstantOfKind(CelMutableExpr expr, CelConstant.Kin
948970
return expr.getKind().equals(Kind.CONSTANT) && expr.constant().getKind().equals(constantKind);
949971
}
950972

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

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

Lines changed: 84 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,8 @@
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.NullableType;
36+
import dev.cel.common.types.OptionalType;
3537
import dev.cel.common.types.SimpleType;
3638
import dev.cel.common.types.StructTypeReference;
3739
import dev.cel.expr.conformance.proto2.TestAllTypes.NestedMessage;
@@ -80,6 +82,28 @@ private static Cel setupEnv(CelBuilder celBuilder) {
8082
return celBuilder
8183
.addVar("x", SimpleType.DYN)
8284
.addVar("y", SimpleType.DYN)
85+
.addVar("dyn_x", SimpleType.DYN)
86+
.addVar("int_x", SimpleType.INT)
87+
.addVar("double_x", SimpleType.DOUBLE)
88+
.addVar("bool_x", SimpleType.BOOL)
89+
.addVar("string_x", SimpleType.STRING)
90+
.addVar("int_list_x", ListType.create(SimpleType.INT))
91+
.addVar("double_list_x", ListType.create(SimpleType.DOUBLE))
92+
.addVar("dyn_list_x", ListType.create(SimpleType.DYN))
93+
.addVar("map_string_int_x", MapType.create(SimpleType.STRING, SimpleType.INT))
94+
.addVar("map_string_double_x", MapType.create(SimpleType.STRING, SimpleType.DOUBLE))
95+
.addVar("optional_int_x", OptionalType.create(SimpleType.INT))
96+
.addVar("optional_double_x", OptionalType.create(SimpleType.DOUBLE))
97+
.addVar("nested_list_int_x", ListType.create(ListType.create(SimpleType.INT)))
98+
.addVar("nested_list_double_x", ListType.create(ListType.create(SimpleType.DOUBLE)))
99+
.addVar(
100+
"nested_map_list_int_x",
101+
MapType.create(SimpleType.STRING, ListType.create(SimpleType.INT)))
102+
.addVar(
103+
"nested_map_list_double_x",
104+
MapType.create(SimpleType.STRING, ListType.create(SimpleType.DOUBLE)))
105+
.addVar("nullable_int_x", NullableType.create(SimpleType.INT))
106+
.addVar("nullable_double_x", NullableType.create(SimpleType.DOUBLE))
83107
.addVar("bool_var", SimpleType.BOOL)
84108
.addVar("list_var", ListType.create(SimpleType.STRING))
85109
.addVar("map_var", MapType.create(SimpleType.STRING, SimpleType.STRING))
@@ -151,7 +175,48 @@ private static Cel setupEnv(CelBuilder celBuilder) {
151175
@TestParameters("{source: '5 in [1, 1 + 2, 1 + (2 + 3)]', expected: 'false'}")
152176
@TestParameters("{source: '5 in [1, x, y, 5]', expected: 'true'}")
153177
@TestParameters("{source: '!(5 in [1, x, y, 5])', expected: 'false'}")
154-
@TestParameters("{source: 'x in [1, x, y, 5]', expected: 'true'}")
178+
@TestParameters("{source: 'x in [1, x, y, 5]', expected: 'x in [1, x, y, 5]'}")
179+
@TestParameters("{source: 'dyn_x in [1, 2, dyn_x]', expected: 'dyn_x in [1, 2, dyn_x]'}")
180+
@TestParameters("{source: 'int_x in [1, 2, int_x]', expected: 'true'}")
181+
@TestParameters("{source: 'bool_x in [true, false, bool_x]', expected: 'true'}")
182+
@TestParameters("{source: 'string_x in [\"a\", \"b\", string_x]', expected: 'true'}")
183+
@TestParameters(
184+
"{source: 'double_x in [1.0, 2.0, double_x]', expected: 'double_x in [1.0, 2.0, double_x]'}")
185+
@TestParameters("{source: 'int_list_x in [[1], [2], int_list_x]', expected: 'true'}")
186+
@TestParameters(
187+
"{source: 'double_list_x in [[1.0], double_list_x]', expected: 'double_list_x in [[1.0],"
188+
+ " double_list_x]'}")
189+
@TestParameters(
190+
"{source: 'dyn_list_x in [[1], dyn_list_x]', expected: 'dyn_list_x in [[1], dyn_list_x]'}")
191+
@TestParameters(
192+
"{source: 'map_string_int_x in [{\"a\": 1}, map_string_int_x]', expected: 'true'}")
193+
@TestParameters(
194+
"{source: 'map_string_double_x in [{\"a\": 1.0}, map_string_double_x]', expected:"
195+
+ " 'map_string_double_x in [{\"a\": 1.0}, map_string_double_x]'}")
196+
@TestParameters(
197+
"{source: 'optional_int_x in [optional.of(1), optional_int_x]', expected: 'true'}")
198+
@TestParameters(
199+
"{source: 'optional_double_x in [optional.of(1.0), optional_double_x]', expected:"
200+
+ " 'optional_double_x in [optional.of(1.0), optional_double_x]'}")
201+
@TestParameters("{source: 'nullable_int_x in [1, 2, nullable_int_x]', expected: 'true'}")
202+
@TestParameters(
203+
"{source: 'nullable_double_x in [1.0, 2.0, nullable_double_x]', expected:"
204+
+ " 'nullable_double_x in [1.0, 2.0, nullable_double_x]'}")
205+
@TestParameters(
206+
"{source: 'double(\"NaN\") in [double(\"NaN\"), double_x]', expected: 'NaN in"
207+
+ " [NaN, double_x]'}")
208+
@TestParameters(
209+
"{source: 'nested_list_int_x in [[[1]], [[2]], nested_list_int_x]', expected: 'true'}")
210+
@TestParameters(
211+
"{source: 'nested_list_double_x in [[[1.0]], [[2.0]], nested_list_double_x]',"
212+
+ " expected: 'nested_list_double_x in [[[1.0]], [[2.0]], nested_list_double_x]'}")
213+
@TestParameters(
214+
"{source: 'nested_map_list_int_x in [{\"a\": [1]}, nested_map_list_int_x]', expected:"
215+
+ " 'true'}")
216+
@TestParameters(
217+
"{source: 'nested_map_list_double_x in [{\"a\": [1.0]}, nested_map_list_double_x]',"
218+
+ " expected: 'nested_map_list_double_x in [{\"a\": [1.0]},"
219+
+ " nested_map_list_double_x]'}")
155220
@TestParameters("{source: 'x in [1, 1 + 2, 1 + (2 + 3)]', expected: 'x in [1, 3, 6]'}")
156221
@TestParameters("{source: 'duration(string(7 * 24) + ''h'')', expected: 'duration(\"168h\")'}")
157222
@TestParameters("{source: '[1, ?optional.of(3)]', expected: '[1, 3]'}")
@@ -395,7 +460,8 @@ public void constantFold_protoMessageLiteral_success(String source, String expec
395460
@TestParameters(
396461
"{source: 'cel.bind(myMap, {\"foo\": \"bar\"}, myMap[?\"foo\"].optMap(x, x + \"baz\"))', "
397462
+ "expected: 'optional.of(\"barbaz\")'}")
398-
@TestParameters("{source: '(1 + 2 + 3 == x) && (x in [1, 2, x])', expected: '6 == x'}")
463+
@TestParameters(
464+
"{source: '(1 + 2 + 3 == x) && (x in [1, 2, x])', expected: '6 == x && x in [1, 2, x]'}")
399465
public void constantFold_macros_macroCallMetadataPopulated(String source, String expected)
400466
throws Exception {
401467
Cel cel =
@@ -862,4 +928,20 @@ public void iterationLimitReached_throws() throws Exception {
862928
assertThrows(CelOptimizationException.class, () -> optimizer.optimize(ast));
863929
assertThat(e).hasMessageThat().contains("Optimization failure: Max iteration count reached.");
864930
}
931+
932+
@Test
933+
public void constantFold_inOperator_withoutMacros_skipsDoubleNan() throws Exception {
934+
Cel celWithoutMacros =
935+
setupEnv(runtimeFlavor.builder()).toCelBuilder().setStandardMacros().build();
936+
CelOptimizer optimizer =
937+
CelOptimizerFactory.standardCelOptimizerBuilder(celWithoutMacros)
938+
.addAstOptimizers(ConstantFoldingOptimizer.getInstance())
939+
.build();
940+
CelAbstractSyntaxTree ast =
941+
celWithoutMacros.compile("double('NaN') in [double('NaN'), double_x]").getAst();
942+
943+
CelAbstractSyntaxTree optimizedAst = optimizer.optimize(ast);
944+
945+
assertThat(CEL_UNPARSER.unparse(optimizedAst)).isEqualTo("NaN in [NaN, double_x]");
946+
}
865947
}

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)