Optimize SwitchBuilder to prune dead SMT branches

PiperOrigin-RevId: 951041974

Author: l46kok

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

@@ -654,15 +654,16 @@ private ProbeResult createProbeResult(
     BoolExpr inMap2 = (BoolExpr) ctx.mkSelect(mapPresence, key2);
     Expr<?> val2 = ctx.mkSelect(mapValues, key2);
 
-    BoolExpr altInMap = ctx.mkOr(inMapOrig, ctx.mkAnd(cond1, inMap1), ctx.mkAnd(cond2, inMap2));
+    BoolExpr condMap1 = CelZ3TypeSystem.mkAndFlattened(ctx, cond1, inMap1);
+    BoolExpr condMap2 = CelZ3TypeSystem.mkAndFlattened(ctx, cond2, inMap2);
+
+    BoolExpr altInMap = CelZ3TypeSystem.mkOrFlattened(ctx, inMapOrig, condMap1, condMap2);
     Expr<?> altVal =
-        ctx.mkITE(
-            inMapOrig,
-            valOrig,
-            ctx.mkITE(
-                ctx.mkAnd(cond1, inMap1),
-                val1,
-                ctx.mkITE(ctx.mkAnd(cond2, inMap2), val2, valOrig)));
+        CelZ3TypeSystem.SwitchBuilder.newBuilder(ctx)
+            .addCase(inMapOrig, valOrig)
+            .addCase(condMap1, val1)
+            .addCase(condMap2, val2)
+            .build(valOrig);
 
     return new ProbeResult(altInMap, altVal);
   }

verifier/src/main/java/dev/cel/verifier/CelZ3TypeSystem.java

@@ -723,7 +723,9 @@ public <R extends Sort> SeqExpr<R> mkConcatSafe(Expr<?> arg1, Expr<?> arg2) {
   /**
    * Helper to build a chain of nested ITE (If-Then-Else) conditions.
    *
-   * <p>Conditions are evaluated in the order they are added.
+   * <p>Conditions are evaluated in the order they are added. Branches with {@code isFalse()}
+   * conditions are skipped, and redundant {@code ITE(condition, X, X)} creations are omitted to
+   * avoid allocating dead AST paths in Z3.
    */
   public static final class SwitchBuilder {
 
@@ -746,13 +748,21 @@ public static SwitchBuilder newBuilder(Context ctx) {
 
     @CanIgnoreReturnValue
     public SwitchBuilder addCase(BoolExpr condition, Expr<?> value) {
+      // Skip branches that can never be hit (e.g. `isFalse()` probes).
+      if (condition.isFalse()) {
+        return this;
+      }
       cases.add(new SwitchCase(condition, value));
       return this;
     }
 
     public Expr<?> build(Expr<?> defaultFallback) {
       Expr<?> result = defaultFallback;
       for (SwitchCase c : Lists.reverse(cases)) {
+        // ITE(condition, X, X) simplifies to X; skip calling into native C++ Z3_mk_ite.
+        if (c.value.equals(result)) {
+          continue;
+        }
         result = ctx.mkITE(c.condition, c.value, result);
       }
       return result;
@@ -764,6 +774,15 @@ private SwitchBuilder(Context ctx) {
     }
   }
 
+  /**
+   * Helper to construct a flattened logical OR expression to avoid deep left-leaning ASTs.
+   *
+   * <p>Returns {@code false} if the list is empty.
+   */
+  public static BoolExpr mkOrFlattened(Context ctx, BoolExpr... args) {
+    return mkOrFlattened(ctx, Arrays.asList(args));
+  }
+
   /**
    * Helper to construct a flattened logical OR expression to avoid deep left-leaning ASTs.
    *

verifier/src/main/java/dev/cel/verifier/TranslatedValue.java

@@ -194,12 +194,10 @@ static TranslatedValue propagateStrict(
     BoolExpr isSafe =
         CelZ3TypeSystem.mkOrFlattened(
             ctx,
-            Arrays.asList(
-                hasExactUnknown,
-                CelZ3TypeSystem.mkAndFlattened(
-                    ctx,
-                    Arrays.asList(hasExactError, CelZ3TypeSystem.mkNotFlattened(ctx, hasUnknown))),
-                CelZ3TypeSystem.mkNotFlattened(ctx, anyTaint)));
+            hasExactUnknown,
+            CelZ3TypeSystem.mkAndFlattened(
+                ctx, hasExactError, CelZ3TypeSystem.mkNotFlattened(ctx, hasUnknown)),
+            CelZ3TypeSystem.mkNotFlattened(ctx, anyTaint));
 
     return create(finalResult, celExpr, ts, CelZ3TypeSystem.mkNotFlattened(ctx, isSafe));
   }