2 files changed, 91 insertions, 13 deletions

diff --git a/lib/Expr/Constraints.cpp b/lib/Expr/Constraints.cpp
index ae4563f4..dbdfd999 100644
--- a/lib/Expr/Constraints.cpp
+++ b/lib/Expr/Constraints.cpp
@@ -23,6 +23,14 @@
 
 using namespace klee;
 
+namespace {
+  llvm::cl::opt<bool>
+  RewriteEqualities("rewrite-equalities",
+		    llvm::cl::init(true),
+		    llvm::cl::desc("Rewrite existing constraints when an equality with a constant is added (default=on)"));
+}
+
+
 class ExprReplaceVisitor : public ExprVisitor {
 private:
   ref<Expr> src, dst;
@@ -118,13 +126,7 @@ ref<Expr> ConstraintManager::simplifyExpr(ref<Expr> e) const {
 }
 
 void ConstraintManager::addConstraintInternal(ref<Expr> e) {
-  // rewrite any known equalities 
-
-  // XXX should profile the effects of this and the overhead.
-  // traversing the constraints looking for equalities is hardly the
-  // slowest thing we do, but it is probably nicer to have a
-  // ConstraintSet ADT which efficiently remembers obvious patterns
-  // (byte-constant comparison).
+  // rewrite any known equalities and split Ands into different conjuncts
 
   switch (e->getKind()) {
   case Expr::Constant:
@@ -141,10 +143,17 @@ void ConstraintManager::addConstraintInternal(ref<Expr> e) {
   }
 
   case Expr::Eq: {
-    BinaryExpr *be = cast<BinaryExpr>(e);
-    if (isa<ConstantExpr>(be->left)) {
-      ExprReplaceVisitor visitor(be->right, be->left);
-      rewriteConstraints(visitor);
+    if (RewriteEqualities) {
+      // XXX: should profile the effects of this and the overhead.
+      // traversing the constraints looking for equalities is hardly the
+      // slowest thing we do, but it is probably nicer to have a
+      // ConstraintSet ADT which efficiently remembers obvious patterns
+      // (byte-constant comparison).
+      BinaryExpr *be = cast<BinaryExpr>(e);
+      if (isa<ConstantExpr>(be->left)) {
+	ExprReplaceVisitor visitor(be->right, be->left);
+	rewriteConstraints(visitor);
+      }
     }
     constraints.push_back(e);
     break;
diff --git a/lib/Expr/ExprSMTLIBPrinter.cpp b/lib/Expr/ExprSMTLIBPrinter.cpp
index c780ae90..e8e9a253 100644
--- a/lib/Expr/ExprSMTLIBPrinter.cpp
+++ b/lib/Expr/ExprSMTLIBPrinter.cpp
@@ -255,7 +255,10 @@ void ExprSMTLIBPrinter::printFullExpression(
      * type T.
      */
     printLogicalOrBitVectorExpr(e, expectedSort);
+    return;
 
+  case Expr::AShr:
+    printAShrExpr(cast<AShrExpr>(e));
     return;
 
   default:
@@ -334,6 +337,73 @@ void ExprSMTLIBPrinter::printCastExpr(const ref<CastExpr> &e) {
   *p << ")";
 }
 
+void ExprSMTLIBPrinter::printAShrExpr(const ref<AShrExpr> &e) {
+  // There is a difference between AShr and SMT-LIBv2's
+  // bvashr function when the shift amount is >= the bit width
+  // so we need to treat it specially here.
+  //
+  // Technically this is undefined behaviour for LLVM's ashr operator
+  // but currently llvm::APInt:ashr(...) will emit 0 if the shift amount
+  // is >= the bit width but this does not match how SMT-LIBv2's bvashr
+  // behaves as demonstrates by the following query
+  //
+  // (declare-fun x () (_ BitVec 32))
+  // (declare-fun y () (_ BitVec 32))
+  // (declare-fun result () (_ BitVec 32))
+  // (assert (bvuge y (_ bv32 32)))
+  // (assert (= result (bvashr x y)))
+  // (assert (distinct (_ bv0 32) result))
+  // (check-sat)
+  // sat
+  //
+  // Our solution is to instead emit
+  //
+  // (ite (bvuge shift_amount bit_width)
+  //      (_ bv0 bitwidth)
+  //      (bvashr value_to_shift shift_amount)
+  // )
+  //
+
+  Expr::Width bitWidth = e->getKid(0)->getWidth();
+  assert(bitWidth > 0 && "Invalid bit width");
+  ref<Expr> bitWidthExpr = ConstantExpr::create(bitWidth, bitWidth);
+  ref<Expr> zeroExpr = ConstantExpr::create(0, bitWidth);
+
+  // FIXME: we print e->getKid(1) twice and it might not get
+  // abbreviated
+  *p << "(ite";
+  p->pushIndent();
+  printSeperator();
+
+  *p << "(bvuge";
+  p->pushIndent();
+  printSeperator();
+  printExpression(e->getKid(1), SORT_BITVECTOR);
+  printSeperator();
+  printExpression(bitWidthExpr, SORT_BITVECTOR);
+  p->popIndent();
+  printSeperator();
+  *p << ")";
+
+  printSeperator();
+  printExpression(zeroExpr, SORT_BITVECTOR);
+  printSeperator();
+
+  *p << "(bvashr";
+  p->pushIndent();
+  printSeperator();
+  printExpression(e->getKid(0), SORT_BITVECTOR);
+  printSeperator();
+  printExpression(e->getKid(1), SORT_BITVECTOR);
+  p->popIndent();
+  printSeperator();
+  *p << ")";
+
+  p->popIndent();
+  printSeperator();
+  *p << ")";
+}
+
 const char *ExprSMTLIBPrinter::getSMTLIBKeyword(const ref<Expr> &e) {
 
   switch (e->getKind()) {
@@ -373,8 +443,7 @@ const char *ExprSMTLIBPrinter::getSMTLIBKeyword(const ref<Expr> &e) {
     return "bvshl";
   case Expr::LShr:
     return "bvlshr";
-  case Expr::AShr:
-    return "bvashr";
+  // AShr is not supported here. See printAShrExpr()
 
   case Expr::Eq:
     return "=";