simplify concat

cucapra · Nov 21, 2024 · 61436ca · 61436ca
1 parent 07c5043
commit 61436ca
Show file tree

Hide file tree

Showing 2 changed files with 57 additions and 0 deletions.
diff --git a/src/expr/simplify.rs b/src/expr/simplify.rs
@@ -46,6 +46,7 @@ pub(crate) fn simplify(ctx: &mut Context, expr: ExprRef, children: &[ExprRef]) -
         (Expr::BVZeroExt { by, .. }, [e]) => simplify_bv_zero_ext(ctx, *e, by),
         (Expr::BVSlice { lo, hi, .. }, [e]) => simplify_bv_slice(ctx, *e, hi, lo),
         (Expr::BVIte { .. }, [cond, tru, fals]) => simplify_ite(ctx, *cond, *tru, *fals),
+        (Expr::BVConcat(..), [a, b]) => simplify_bv_concat(ctx, *a, *b),
         (Expr::BVEqual(..), [a, b]) => simplify_bv_equal(ctx, *a, *b),
         (Expr::BVAnd(..), [a, b]) => simplify_bv_and(ctx, *a, *b),
         (Expr::BVOr(..), [a, b]) => simplify_bv_or(ctx, *a, *b),
@@ -285,6 +286,26 @@ fn simplify_bv_sign_ext(ctx: &mut Context, e: ExprRef, by: WidthInt) -> Option<E
     }
 }
 
+fn simplify_bv_concat(ctx: &mut Context, a: ExprRef, b: ExprRef) -> Option<ExprRef> {
+    match *ctx.get(a) {
+        // normalize concat to be right recursive
+        Expr::BVConcat(a_a, a_b, _) => Some(ctx.build(|c| c.concat(a_a, c.concat(a_b, b)))),
+        Expr::BVLiteral(va) => match *ctx.get(b) {
+            Expr::BVLiteral(vb) => Some(ctx.bv_lit(&va.get(ctx).concat(&vb.get(ctx)))),
+            Expr::BVConcat(b_a, b_b, _) => {
+                if let Expr::BVLiteral(v_b_a) = *ctx.get(b_a) {
+                    let lit = ctx.bv_lit(&va.get(ctx).concat(&v_b_a.get(ctx)));
+                    Some(ctx.concat(lit, b_b))
+                } else {
+                    None
+                }
+            }
+            _ => None,
+        },
+        _ => None,
+    }
+}
+
 fn simplify_bv_slice(ctx: &mut Context, e: ExprRef, hi: WidthInt, lo: WidthInt) -> Option<ExprRef> {
     match ctx.get(e) {
         // combine slices

diff --git a/tests/simplify.rs b/tests/simplify.rs
@@ -218,6 +218,42 @@ fn test_simplify_mul() {
     ts("mul(a : bv<4>, 4'd8)", "concat(a : bv<4>[0],  3'b000)");
 }
 
+#[test]
+fn test_simplify_concat() {
+    // normalize concats
+    ts(
+        "concat(concat(a : bv<1>, b:bv<1>), c:bv<1>)",
+        "concat(a : bv<1>, concat(b:bv<1>, c:bv<1>))",
+    );
+
+    // concat constants
+    ts("concat(3'b110, 2'b01)", "5'b11001");
+    ts("concat(3'b110, concat(false, true))", "5'b11001");
+    ts("concat(concat(3'b110, false), true)", "5'b11001");
+
+    // concat constants in the presence of other expressions
+    ts(
+        "concat(a : bv<3>, concat(false, true))",
+        "concat(a:bv<3>, 2'b01)",
+    );
+    ts(
+        "concat(concat(a : bv<3>, false), true)",
+        "concat(a:bv<3>, 2'b01)",
+    );
+    ts(
+        "concat(true, concat(false, a : bv<3>))",
+        "concat(2'b10, a:bv<3>)",
+    );
+    ts(
+        "concat(concat(true, false), a : bv<3>)",
+        "concat(2'b10, a:bv<3>)",
+    );
+    ts(
+        "concat(3'b101, concat(true, concat(false, a : bv<3>)))",
+        "concat(5'b10110, a:bv<3>)",
+    );
+}
+
 // from maltese-smt:
 // https://github.com/ucb-bar/maltese-smt/blob/main/test/maltese/smt/SMTSimplifierSpec.scala
 #[test]