Skip to content

Commit e8f7ded

Browse files
l46kokcopybara-github
l46kok
authored and
copybara-github
committed
Add support for optional field pruning in verifier
PiperOrigin-RevId: 951556122
1 parent c11a0f5 commit e8f7ded

5 files changed

Lines changed: 175 additions & 51 deletions

File tree

‎verifier/src/main/java/dev/cel/verifier/CelAstAlphaHasher.java‎

Lines changed: 14 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,9 @@
2424
import dev.cel.common.ast.CelConstant;
2525
import dev.cel.common.ast.CelExpr;
2626
import java.util.ArrayList;
27+
import java.util.HashMap;
2728
import java.util.List;
29+
import java.util.Map;
2830
import org.jspecify.annotations.Nullable;
2931

3032
/**
@@ -83,16 +85,11 @@ private static void hashAst(CelExpr expr, @Nullable Scope scope, HasherContext c
8385
context.hasher.putByte((byte) 0); // 0 = bound
8486
context.hasher.putInt(bIdx);
8587
} else {
86-
int fIdx = -1;
87-
for (int i = 0; i < context.freeVars.size(); i++) {
88-
if (context.freeVars.get(i).ident().name().equals(name)) {
89-
fIdx = i;
90-
break;
91-
}
92-
}
93-
if (fIdx == -1) {
88+
Integer fIdx = context.freeVarIndices.get(name);
89+
if (fIdx == null) {
90+
fIdx = context.freeVars.size();
9491
context.freeVars.add(expr);
95-
fIdx = context.freeVars.size() - 1;
92+
context.freeVarIndices.put(name, fIdx);
9693
}
9794
context.hasher.putByte((byte) 1); // 1 = free
9895
context.hasher.putInt(fIdx);
@@ -121,6 +118,10 @@ private static void hashAst(CelExpr expr, @Nullable Scope scope, HasherContext c
121118
for (CelExpr elem : expr.list().elements()) {
122119
hashAst(elem, scope, context);
123120
}
121+
context.hasher.putInt(expr.list().optionalIndices().size());
122+
for (int optIndex : expr.list().optionalIndices()) {
123+
context.hasher.putInt(optIndex);
124+
}
124125
break;
125126
case STRUCT:
126127
context.hasher.putString(expr.struct().messageName(), UTF_8);
@@ -208,10 +209,13 @@ private static void hashConstant(CelConstant constant, HasherContext context) {
208209

209210
private static final class HasherContext {
210211
final Hasher hasher;
211-
final List<CelExpr> freeVars = new ArrayList<>();
212+
final List<CelExpr> freeVars;
213+
final Map<String, Integer> freeVarIndices;
212214

213215
HasherContext(HashFunction hashFunction) {
214216
this.hasher = hashFunction.newHasher();
217+
this.freeVars = new ArrayList<>();
218+
this.freeVarIndices = new HashMap<>();
215219
}
216220
}
217221

‎verifier/src/main/java/dev/cel/verifier/CelAstToZ3Translator.java‎

Lines changed: 129 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -88,7 +88,8 @@ final class CelAstToZ3Translator {
8888
private final ImmutableSet<String> unknownIdentifiers;
8989
private final int comprehensionUnrollLimit;
9090
private final CelTypeProvider typeProvider;
91-
private final List<BoolExpr> truncationConditions;
91+
private final Set<BoolExpr> truncationConditions;
92+
private final Set<SeqExpr<?>> boundedMapBijections;
9293
private final Map<String, Expr<?>> emptyMessageCache;
9394
private final Map<Object, Expr<?>> listLiteralCache;
9495
private Expr<?> emptyListCache;
@@ -102,7 +103,7 @@ Set<BoolExpr> getTypeConstraints() {
102103
}
103104

104105
BoolExpr hasTruncation() {
105-
return CelZ3TypeSystem.mkOrFlattened(ctx, truncationConditions);
106+
return CelZ3TypeSystem.mkOrFlattened(ctx, new ArrayList<>(truncationConditions));
106107
}
107108

108109
CelZ3TypeSystem getTypeSystem() {
@@ -284,11 +285,23 @@ private TranslatedValue translateList(CelExpr celExpr, CelAbstractSyntaxTree ast
284285
// check to a trivial identity check (e.g., `list_ref_0 == list_ref_0`).
285286
if (listRef == null) {
286287
SeqExpr seq = ctx.mkEmptySeq(ctx.mkSeqSort(typeSystem.celValueSort()));
287-
for (CelExpr element : createList.elements()) {
288+
ImmutableSet<Integer> optionalIndices = ImmutableSet.copyOf(createList.optionalIndices());
289+
for (int i = 0; i < createList.elements().size(); i++) {
290+
CelExpr element = createList.elements().get(i);
288291
TranslatedValue elem = translateExpr(element, ast);
289292
elementsTv.add(elem);
290293

291-
seq = typeSystem.mkConcatSafe(seq, ctx.mkUnit(elem.z3Expr()));
294+
boolean isOptional = optionalIndices.contains(i);
295+
OptionalUnwrap opt = unwrapOptional(elem.z3Expr(), isOptional);
296+
SeqExpr optSeq =
297+
isOptional
298+
? (SeqExpr)
299+
ctx.mkITE(
300+
opt.isPresent,
301+
ctx.mkUnit(opt.effectiveValue),
302+
ctx.mkEmptySeq(ctx.mkSeqSort(typeSystem.celValueSort())))
303+
: ctx.mkUnit(opt.effectiveValue);
304+
seq = typeSystem.mkConcatSafe(seq, optSeq);
292305
}
293306
listRef = typeSystem.mkListRefConst(LIST_REF_PREFIX);
294307
typeConstraints.add(ctx.mkEq(typeSystem.getSeq(listRef), seq));
@@ -297,7 +310,9 @@ private TranslatedValue translateList(CelExpr celExpr, CelAbstractSyntaxTree ast
297310
}
298311

299312
Expr<?> result = typeSystem.wrapList(listRef);
300-
return TranslatedValue.propagateStrict(ctx, typeSystem, result, celExpr, elementsTv);
313+
BoolExpr baseTaint = ctx.mkFalse();
314+
return TranslatedValue.propagateStrict(
315+
ctx, typeSystem, result, Optional.of(celExpr), baseTaint, elementsTv);
301316
}
302317

303318
private TranslatedValue translateMap(CelExpr celExpr, CelAbstractSyntaxTree ast) {
@@ -318,20 +333,31 @@ private TranslatedValue translateMap(CelExpr celExpr, CelAbstractSyntaxTree ast)
318333
Expr<?> value = valueTv.z3Expr();
319334
elementsTv.add(valueTv);
320335

321-
BoolExpr keyAlreadyPresent = (BoolExpr) ctx.mkSelect(mapPresence, key);
322-
keysSeq =
323-
ctx.mkITE(keyAlreadyPresent, keysSeq, typeSystem.mkConcatSafe(keysSeq, ctx.mkUnit(key)));
336+
OptionalUnwrap opt = unwrapOptional(value, entryAst.optionalEntry());
324337

325-
mapValues = ctx.mkStore(mapValues, key, value);
326-
mapPresence = ctx.mkStore(mapPresence, key, ctx.mkTrue());
338+
BoolExpr keyAlreadyPresent = (BoolExpr) ctx.mkSelect(mapPresence, key);
339+
BoolExpr shouldAddKey = ctx.mkAnd(opt.isPresent, ctx.mkNot(keyAlreadyPresent));
340+
341+
SeqExpr keyOptSeq =
342+
(SeqExpr)
343+
ctx.mkITE(
344+
shouldAddKey,
345+
ctx.mkUnit(key),
346+
ctx.mkEmptySeq(ctx.mkSeqSort(typeSystem.celValueSort())));
347+
348+
keysSeq = typeSystem.mkConcatSafe(keysSeq, keyOptSeq);
349+
mapValues = storeIf(opt.isPresent, mapValues, key, opt.effectiveValue);
350+
mapPresence = storeIf(opt.isPresent, mapPresence, key, ctx.mkTrue());
327351
}
328352

329353
typeConstraints.add(ctx.mkEq(typeSystem.getMapValues(mapRef), mapValues));
330354
typeConstraints.add(ctx.mkEq(typeSystem.getMapPresence(mapRef), mapPresence));
331355
typeConstraints.add(ctx.mkEq(typeSystem.getMapKeys(mapRef), keysSeq));
332356

333357
Expr<?> result = typeSystem.wrapMap(mapRef);
334-
return TranslatedValue.propagateStrict(ctx, typeSystem, result, celExpr, elementsTv);
358+
BoolExpr baseTaint = ctx.mkFalse();
359+
return TranslatedValue.propagateStrict(
360+
ctx, typeSystem, result, Optional.of(celExpr), baseTaint, elementsTv);
335361
}
336362

337363
private TranslatedValue translateStruct(CelExpr celExpr, CelAbstractSyntaxTree ast) {
@@ -379,15 +405,15 @@ private TranslatedValue translateStruct(CelExpr celExpr, CelAbstractSyntaxTree a
379405
// (`msg1 == msg2`) to work without using quantifiers (which avoids MBQI loops).
380406
// Because proto3 singular primitives do not have field presence, we also skip setting
381407
// `msgPresence`.
408+
OptionalUnwrap opt = unwrapOptional(value, entryAst.optionalEntry());
409+
382410
BoolExpr shouldBypass =
383-
fieldType.kind().isPrimitive() ? ctx.mkEq(value, defaultVal) : ctx.mkFalse();
411+
fieldType.kind().isPrimitive() ? ctx.mkEq(opt.effectiveValue, defaultVal) : ctx.mkFalse();
384412

385-
msgValues =
386-
(ArrayExpr) ctx.mkITE(shouldBypass, msgValues, ctx.mkStore(msgValues, key, value));
413+
BoolExpr shouldStore = ctx.mkAnd(opt.isPresent, ctx.mkNot(shouldBypass));
387414

388-
msgPresence =
389-
(ArrayExpr)
390-
ctx.mkITE(shouldBypass, msgPresence, ctx.mkStore(msgPresence, key, ctx.mkTrue()));
415+
msgValues = storeIf(shouldStore, msgValues, key, opt.effectiveValue);
416+
msgPresence = storeIf(shouldStore, msgPresence, key, ctx.mkTrue());
391417
}
392418

393419
typeConstraints.add(
@@ -396,7 +422,9 @@ private TranslatedValue translateStruct(CelExpr celExpr, CelAbstractSyntaxTree a
396422
typeConstraints.add(ctx.mkEq(typeSystem.getMsgPresence(msgRef), msgPresence));
397423

398424
Expr<?> result = typeSystem.wrapMessage(msgRef);
399-
return TranslatedValue.propagateStrict(ctx, typeSystem, result, celExpr, elementsTv);
425+
BoolExpr baseTaint = ctx.mkFalse();
426+
return TranslatedValue.propagateStrict(
427+
ctx, typeSystem, result, Optional.of(celExpr), baseTaint, elementsTv);
400428
}
401429

402430
private Expr<?> getDefaultValueForType(CelType type) {
@@ -472,6 +500,28 @@ private Expr<?> getDefaultValueForType(CelType type) {
472500
return typeSystem.mkUnknown();
473501
}
474502

503+
private static final class OptionalUnwrap {
504+
final BoolExpr isPresent;
505+
final Expr<?> effectiveValue;
506+
507+
OptionalUnwrap(BoolExpr isPresent, Expr<?> effectiveValue) {
508+
this.isPresent = isPresent;
509+
this.effectiveValue = effectiveValue;
510+
}
511+
}
512+
513+
private OptionalUnwrap unwrapOptional(Expr<?> value, boolean isOptional) {
514+
if (!isOptional) {
515+
return new OptionalUnwrap(ctx.mkTrue(), value);
516+
}
517+
Expr<?> optRef = typeSystem.getOptionalRef(value);
518+
return new OptionalUnwrap(typeSystem.optHasValue(optRef), typeSystem.getOptionalValue(optRef));
519+
}
520+
521+
private ArrayExpr storeIf(BoolExpr condition, ArrayExpr array, Expr<?> key, Expr<?> value) {
522+
return (ArrayExpr) ctx.mkITE(condition, ctx.mkStore(array, key, value), array);
523+
}
524+
475525
private static final class FieldAccess {
476526
final Expr<?> presence;
477527
final Expr<?> value;
@@ -657,24 +707,53 @@ private TranslatedValue translateComprehension(CelExpr celExpr, CelAbstractSynta
657707
// For statically known list/map literals, unroll them exactly.
658708
if (iterRangeExpr.exprKind().getKind() == ExprKind.Kind.LIST) {
659709
ImmutableList<CelExpr> elements = iterRangeExpr.list().elements();
710+
ImmutableList<Integer> optionalIndices = iterRangeExpr.list().optionalIndices();
660711
for (int i = 0; i < elements.size(); i++) {
661712
TranslatedValue valueTv = translateExpr(elements.get(i), ast);
662-
Expr<?> value = valueTv.z3Expr();
663-
taints.add(valueTv.isApproximate());
664-
iterationElements.add(new IterationElement(typeSystem.mkInt(i), value));
665-
allRangeElems.add(value);
713+
boolean isOptional = optionalIndices.contains(i);
714+
OptionalUnwrap opt = unwrapOptional(valueTv.z3Expr(), isOptional);
715+
if (isOptional && ctx.mkFalse().equals(opt.isPresent)) {
716+
continue;
717+
}
718+
taints.add(
719+
isOptional
720+
? CelZ3TypeSystem.mkAndFlattened(ctx, opt.isPresent, valueTv.isApproximate())
721+
: valueTv.isApproximate());
722+
iterationElements.add(
723+
new IterationElement(
724+
typeSystem.mkInt(i),
725+
opt.effectiveValue,
726+
isOptional ? Optional.of(opt.isPresent) : Optional.empty()));
727+
allRangeElems.add(
728+
isOptional
729+
? ctx.mkITE(opt.isPresent, opt.effectiveValue, typeSystem.mkInt(0))
730+
: opt.effectiveValue);
666731
}
667732
} else if (iterRangeExpr.exprKind().getKind() == ExprKind.Kind.MAP) {
668733
for (CelExpr.CelMap.Entry entry : iterRangeExpr.map().entries()) {
669734
TranslatedValue keyTv = translateExpr(entry.key(), ast);
670735
Expr<?> key = keyTv.z3Expr();
671736
taints.add(keyTv.isApproximate());
672737
TranslatedValue valueTv = translateExpr(entry.value(), ast);
673-
Expr<?> value = valueTv.z3Expr();
674-
taints.add(valueTv.isApproximate());
675-
iterationElements.add(new IterationElement(key, value));
738+
boolean isOptional = entry.optionalEntry();
739+
OptionalUnwrap opt = unwrapOptional(valueTv.z3Expr(), isOptional);
740+
if (isOptional && ctx.mkFalse().equals(opt.isPresent)) {
741+
continue;
742+
}
743+
taints.add(
744+
isOptional
745+
? CelZ3TypeSystem.mkAndFlattened(ctx, opt.isPresent, valueTv.isApproximate())
746+
: valueTv.isApproximate());
747+
iterationElements.add(
748+
new IterationElement(
749+
key,
750+
opt.effectiveValue,
751+
isOptional ? Optional.of(opt.isPresent) : Optional.empty()));
676752
allRangeElems.add(key);
677-
allRangeElems.add(value);
753+
allRangeElems.add(
754+
isOptional
755+
? ctx.mkITE(opt.isPresent, opt.effectiveValue, typeSystem.mkInt(0))
756+
: opt.effectiveValue);
678757
}
679758
} else {
680759
return translateDynamicComprehension(celExpr, ast);
@@ -693,13 +772,23 @@ private TranslatedValue translateComprehension(CelExpr celExpr, CelAbstractSynta
693772
comp, ast, iterElem.keyOrIndex, iterElem.value, currentAccu, isMap, isTwoVar);
694773
Expr<?> condition = condAndStep[0].z3Expr();
695774
Expr<?> step = condAndStep[1].z3Expr();
696-
taints.add(condAndStep[1].isApproximate());
775+
if (iterElem.hasValue.isPresent()) {
776+
taints.add(
777+
CelZ3TypeSystem.mkAndFlattened(
778+
ctx, iterElem.hasValue.get(), condAndStep[1].isApproximate()));
779+
} else {
780+
taints.add(condAndStep[1].isApproximate());
781+
}
697782

698783
Expr<?> stepVal = ctx.mkITE((BoolExpr) typeSystem.unwrapBool(condition), step, currentAccu);
699784
Expr<?> typeErrorOrStep =
700785
typeSystem.withRuntimeError(stepVal, ctx.mkNot(typeSystem.isBool(condition)));
701786

702-
accu = typeSystem.propagateErrorAndUnknown(typeErrorOrStep, condition);
787+
Expr<?> updatedAccu = typeSystem.propagateErrorAndUnknown(typeErrorOrStep, condition);
788+
accu =
789+
iterElem.hasValue.isPresent()
790+
? ctx.mkITE(iterElem.hasValue.get(), updatedAccu, currentAccu)
791+
: updatedAccu;
703792
}
704793

705794
TranslatedValue resultTv =
@@ -773,6 +862,9 @@ private TranslatedValue translateDynamicComprehension(
773862

774863
private void applyBoundedMapBijection(
775864
ArrayExpr mapPresence, SeqExpr<?> seq, ArithExpr lengthExpr) {
865+
if (!boundedMapBijections.add(seq)) {
866+
return;
867+
}
776868
for (int i = 0; i < comprehensionUnrollLimit; i++) {
777869
for (int j = i + 1; j < comprehensionUnrollLimit; j++) {
778870
BoolExpr validPair = ctx.mkLt(ctx.mkInt(j), lengthExpr);
@@ -1223,16 +1315,23 @@ private BoolExpr createTypeConstraintForType(Expr<?> val, CelType type) {
12231315
this.emptyMessageCache = new HashMap<>();
12241316
this.listLiteralCache = new HashMap<>();
12251317
this.typeProvider = typeProvider;
1226-
this.truncationConditions = new ArrayList<>();
1318+
this.truncationConditions = new LinkedHashSet<>();
1319+
this.boundedMapBijections = new LinkedHashSet<>();
12271320
}
12281321

12291322
private static class IterationElement {
12301323
final Expr<?> keyOrIndex;
12311324
final Expr<?> value;
1325+
final Optional<BoolExpr> hasValue;
12321326

12331327
IterationElement(Expr<?> keyOrIndex, Expr<?> value) {
1328+
this(keyOrIndex, value, Optional.empty());
1329+
}
1330+
1331+
IterationElement(Expr<?> keyOrIndex, Expr<?> value, Optional<BoolExpr> hasValue) {
12341332
this.keyOrIndex = keyOrIndex;
12351333
this.value = value;
1334+
this.hasValue = hasValue;
12361335
}
12371336
}
12381337

@@ -1249,6 +1348,7 @@ private Optional<Object> toCacheKey(CelExpr expr) {
12491348
}
12501349
builder.add(elemKey.get());
12511350
}
1351+
builder.add(expr.list().optionalIndices());
12521352
return Optional.of(builder.build());
12531353
default:
12541354
return Optional.empty();

‎verifier/src/main/java/dev/cel/verifier/CelZ3OperatorTranslator.java‎

Lines changed: 4 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -33,7 +33,6 @@
3333
import dev.cel.common.Operator;
3434
import dev.cel.common.ast.CelConstant;
3535
import dev.cel.common.ast.CelExpr;
36-
import dev.cel.common.ast.CelExpr.ExprKind;
3736
import dev.cel.common.ast.CelReference;
3837
import dev.cel.common.types.CelKind;
3938
import dev.cel.common.types.CelType;
@@ -500,7 +499,7 @@ private BoolExpr getDynamicNumericEquality(Expr<?> z3Expr0, Expr<?> z3Expr1) {
500499
private BoolExpr unrollListEquality(
501500
TranslatedValue listA, TranslatedValue listB, CelAbstractSyntaxTree ast) {
502501
CelExpr literalListAst =
503-
listA.isLiteral(ExprKind.Kind.LIST) ? listA.celExpr().get() : listB.celExpr().get();
502+
listA.isUnrollableList() ? listA.celExpr().get() : listB.celExpr().get();
504503

505504
SeqExpr<?> seq0 = typeSystem.getSeq(typeSystem.getListRef(listA.z3Expr()));
506505
SeqExpr<?> seq1 = typeSystem.getSeq(typeSystem.getListRef(listB.z3Expr()));
@@ -538,13 +537,12 @@ private TranslatedValue translateEquality(
538537
CelType type0 = extractAstTypeOrDefault(arg0, ast);
539538
CelType type1 = extractAstTypeOrDefault(arg1, ast);
540539

540+
boolean canUnrollList = arg0.isUnrollableList() || arg1.isUnrollableList();
541541
BoolExpr equality;
542542

543543
if (isNumericType(type0) && isNumericType(type1)) {
544544
equality = getNumericEquality(arg0, arg1, ast);
545-
} else if (type0.kind() == CelKind.LIST
546-
&& type1.kind() == CelKind.LIST
547-
&& (arg0.isLiteral(ExprKind.Kind.LIST) || arg1.isLiteral(ExprKind.Kind.LIST))) {
545+
} else if (type0.kind() == CelKind.LIST && type1.kind() == CelKind.LIST && canUnrollList) {
548546
equality = unrollListEquality(arg0, arg1, ast);
549547
} else if (isStaticallyKnown(type0) && isStaticallyKnown(type1)) {
550548
equality = typeSystem.getStructuralEquality(z3Arg0, z3Arg1);
@@ -554,7 +552,7 @@ private TranslatedValue translateEquality(
554552

555553
// Check if one side is an explicit LIST that we can unroll
556554
BoolExpr structuralEq = typeSystem.getStructuralEquality(z3Arg0, z3Arg1);
557-
if (arg0.isLiteral(ExprKind.Kind.LIST) || arg1.isLiteral(ExprKind.Kind.LIST)) {
555+
if (canUnrollList) {
558556
structuralEq =
559557
(BoolExpr)
560558
ctx.mkITE(

0 commit comments

Comments
 (0)