4 files changed, 95 insertions, 2 deletions

diff --git a/include/klee/util/ExprSMTLIBPrinter.h b/include/klee/util/ExprSMTLIBPrinter.h
index 6b0d260a..e90a49f1 100644
--- a/include/klee/util/ExprSMTLIBPrinter.h
+++ b/include/klee/util/ExprSMTLIBPrinter.h
@@ -319,6 +319,7 @@ protected:
   void printNotEqualExpr(const ref<NeExpr> &e);
   void printSelectExpr(const ref<SelectExpr> &e,
                                ExprSMTLIBPrinter::SMTLIB_SORT s);
+  void printAShrExpr(const ref<AShrExpr> &e);
 
   // For the set of operators that take sort "s" arguments
   void printSortArgsExpr(const ref<Expr> &e,
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 "=";
diff --git a/test/Solver/AShr_to_smtlib.kquery b/test/Solver/AShr_to_smtlib.kquery
new file mode 100644
index 00000000..774c46d9
--- /dev/null
+++ b/test/Solver/AShr_to_smtlib.kquery
@@ -0,0 +1,16 @@
+# RUN: %kleaver -print-smtlib -smtlib-abbreviation-mode=none %s > %t
+# RUN: diff -u %t %s.good.smt2
+
+# This test tries to check that the SMT-LIBv2 we generate when a AShrExpr is
+# used is correct.
+#
+# FIXME: We should really pass the generated query to an SMT solver that supports
+# SMT-LIBv2 and check it gives the correct answer ``unsat``. An older version of
+# KLEE where AShrExpr wasn't handled correctly would give ``sat`` for this query.
+#
+# We could fix this if we required STP to be in the user's PATH and made available
+# as a substitution in llvm-lit
+
+array value[1] : w32 -> w8 = symbolic
+array shift[1] : w32 -> w8 = symbolic
+(query [(Ule 8 (Read w8 0 shift))] (Eq 0 (AShr w8 (Read w8 0 value) (Read w8 0 shift))) )
diff --git a/test/Solver/AShr_to_smtlib.kquery.good.smt2 b/test/Solver/AShr_to_smtlib.kquery.good.smt2
new file mode 100644
index 00000000..e9478da6
--- /dev/null
+++ b/test/Solver/AShr_to_smtlib.kquery.good.smt2
@@ -0,0 +1,7 @@
+;SMTLIBv2 Query 0
+(set-logic QF_AUFBV )
+(declare-fun shift () (Array (_ BitVec 32) (_ BitVec 8) ) )
+(declare-fun value () (Array (_ BitVec 32) (_ BitVec 8) ) )
+(assert (and  (=  false (=  (_ bv0 8) (ite (bvuge (select  shift (_ bv0 32) ) (_ bv8 8) ) (_ bv0 8) (bvashr (select  value (_ bv0 32) ) (select  shift (_ bv0 32) ) ) ) ) ) (bvule  (_ bv8 8) (select  shift (_ bv0 32) ) ) ) )
+(check-sat)
+(exit)