Added support for KLEE index-based array optimization

1 files changed, 327 insertions, 0 deletions

diff --git a/lib/Expr/AssignmentGenerator.cpp b/lib/Expr/AssignmentGenerator.cpp
new file mode 100644
index 00000000..755263f0
--- /dev/null
+++ b/lib/Expr/AssignmentGenerator.cpp
@@ -0,0 +1,327 @@
+#include "klee/AssignmentGenerator.h"
+
+#include "llvm/Support/raw_ostream.h"
+
+using namespace klee;
+
+bool AssignmentGenerator::generatePartialAssignment(const ref<Expr> &e,
+                                                    ref<Expr> &val,
+                                                    Assignment *&a) {
+  return helperGenerateAssignment(e, val, a, e->getWidth(), false);
+}
+
+bool AssignmentGenerator::helperGenerateAssignment(const ref<Expr> &e,
+                                                   ref<Expr> &val,
+                                                   Assignment *&a,
+                                                   Expr::Width width,
+                                                   bool sign) {
+  Expr &ep = *e.get();
+  switch (ep.getKind()) {
+
+  // ARITHMETIC
+  case Expr::Add: {
+    // val = val - kid
+    ref<Expr> kid_val;
+    ref<Expr> kid_expr;
+    if (isa<ConstantExpr>(ep.getKid(0))) {
+      kid_val = ep.getKid(0);
+      kid_expr = ep.getKid(1);
+    } else if (isa<ConstantExpr>(ep.getKid(1))) {
+      kid_val = ep.getKid(1);
+      kid_expr = ep.getKid(0);
+    } else {
+      return false;
+    }
+    if (!ExtractExpr::create(kid_val, kid_val.get()->getWidth() - 1, 1).get()->isZero()) {
+      // FIXME: really bad hack to support those cases in which KLEE creates
+      // Add expressions with negative values
+      val = createAddExpr(kid_val, val);
+    } else {
+      val = createSubExpr(kid_val, val);
+    }
+    return helperGenerateAssignment(kid_expr, val, a, width, sign);
+  }
+  case Expr::Sub: {
+    // val = val + kid
+    ref<Expr> kid_val;
+    ref<Expr> kid_expr;
+    if (isa<ConstantExpr>(ep.getKid(0))) {
+      kid_val = ep.getKid(0);
+      kid_expr = ep.getKid(1);
+    } else if (isa<ConstantExpr>(ep.getKid(1))) {
+      kid_val = ep.getKid(1);
+      kid_expr = ep.getKid(0);
+    } else {
+      return false;
+    }
+    val = createAddExpr(kid_val, val);
+    return helperGenerateAssignment(kid_expr, val, a, width, sign);
+  }
+  case Expr::Mul: {
+    // val = val / kid (check for sign)
+    ref<Expr> kid_val;
+    ref<Expr> kid_expr;
+    if (isa<ConstantExpr>(ep.getKid(0))) {
+      kid_val = ep.getKid(0);
+      kid_expr = ep.getKid(1);
+    } else if (isa<ConstantExpr>(ep.getKid(1))) {
+      kid_val = ep.getKid(1);
+      kid_expr = ep.getKid(0);
+    } else {
+      return false;
+    }
+
+    if (kid_val.get()->isZero()) {
+      return false;
+    } else if (!createDivRem(kid_val, val, sign)->isZero()) {
+      return false;
+    }
+    val = createDivExpr(kid_val, val, sign);
+    return helperGenerateAssignment(kid_expr, val, a, width, sign);
+  }
+  case Expr::UDiv:
+  // val = val * kid
+  // no break, falling down to case SDiv
+  case Expr::SDiv: {
+    // val = val * kid
+    ref<Expr> kid_val;
+    ref<Expr> kid_expr;
+    if (isa<ConstantExpr>(ep.getKid(0))) {
+      kid_val = ep.getKid(0);
+      kid_expr = ep.getKid(1);
+    } else if (isa<ConstantExpr>(ep.getKid(1))) {
+      kid_val = ep.getKid(1);
+      kid_expr = ep.getKid(0);
+    } else {
+      return false;
+    }
+    val = createMulExpr(kid_val, val);
+    return helperGenerateAssignment(kid_expr, val, a, width, sign);
+  }
+
+  // LOGICAL
+  case Expr::LShr: {
+    if (!isa<ConstantExpr>(ep.getKid(1))) {
+      return false;
+    }
+    ref<Expr> kid_val = ep.getKid(1);
+    val = createShlExpr(val, kid_val);
+    return helperGenerateAssignment(ep.getKid(0), val, a, width, sign);
+  }
+  case Expr::Shl: {
+    if (!isa<ConstantExpr>(ep.getKid(1))) {
+      return false;
+    }
+    ref<Expr> kid_val = ep.getKid(1);
+    val = createLShrExpr(val, kid_val);
+    return helperGenerateAssignment(ep.getKid(0), val, a, width, sign);
+  }
+  case Expr::Not: {
+    return helperGenerateAssignment(ep.getKid(0), val, a, width, sign);
+  }
+  case Expr::And: {
+    // val = val & kid
+    ref<Expr> kid_val;
+    ref<Expr> kid_expr;
+    if (isa<ConstantExpr>(ep.getKid(0))) {
+      kid_val = ep.getKid(0);
+      kid_expr = ep.getKid(1);
+    } else if (isa<ConstantExpr>(ep.getKid(1))) {
+      kid_val = ep.getKid(1);
+      kid_expr = ep.getKid(0);
+    } else {
+      return false;
+    }
+    val = createAndExpr(kid_val, val);
+    return helperGenerateAssignment(kid_expr, val, a, width, sign);
+  }
+
+  // CASTING
+  case Expr::ZExt: {
+    val = createExtractExpr(ep.getKid(0), val);
+    return helperGenerateAssignment(ep.getKid(0), val, a, width, false);
+  }
+  case Expr::SExt: {
+    val = createExtractExpr(ep.getKid(0), val);
+    return helperGenerateAssignment(ep.getKid(0), val, a, width, true);
+  }
+
+  // SPECIAL
+  case Expr::Concat: {
+    ReadExpr *base = hasOrderedReads(&ep);
+    if (base) {
+      return helperGenerateAssignment(ref<Expr>(base), val, a, ep.getWidth(),
+                                      sign);
+    } else {
+      klee_warning("Not supported");
+      ep.printKind(llvm::errs(), ep.getKind());
+      return false;
+    }
+  }
+  case Expr::Extract: {
+    val = createExtendExpr(ep.getKid(0), val);
+    return helperGenerateAssignment(ep.getKid(0), val, a, width, sign);
+  }
+
+  // READ
+  case Expr::Read: {
+    ReadExpr &re = static_cast<ReadExpr &>(ep);
+    if (isa<ConstantExpr>(re.index)) {
+      if (re.updates.root->isSymbolicArray()) {
+        ConstantExpr &index = static_cast<ConstantExpr &>(*re.index.get());
+        std::vector<unsigned char> c_value =
+            getIndexedValue(getByteValue(val), index, re.updates.root->size);
+        if (c_value.size() == 0) {
+          return false;
+        }
+        if (a->bindings.find(re.updates.root) == a->bindings.end()) {
+          a->bindings.insert(std::make_pair(re.updates.root, c_value));
+        } else {
+          return false;
+        }
+      }
+    } else {
+      return helperGenerateAssignment(re.index, val, a, width, sign);
+    }
+    return true;
+  }
+  default:
+    std::string type_str;
+    llvm::raw_string_ostream rso(type_str);
+    ep.printKind(rso, ep.getKind());
+    klee_warning("%s is not supported", rso.str().c_str());
+    return false;
+  }
+}
+
+bool AssignmentGenerator::isReadExprAtOffset(ref<Expr> e, const ReadExpr *base,
+                                             ref<Expr> offset) {
+  const ReadExpr *re = dyn_cast<ReadExpr>(e.get());
+  if (!re || (re->getWidth() != Expr::Int8))
+    return false;
+  return SubExpr::create(re->index, base->index) == offset;
+}
+
+ReadExpr *AssignmentGenerator::hasOrderedReads(ref<Expr> e) {
+  assert(e->getKind() == Expr::Concat);
+
+  const ReadExpr *base = dyn_cast<ReadExpr>(e->getKid(0));
+
+  // right now, all Reads are byte reads but some
+  // transformations might change this
+  if (!base || base->getWidth() != Expr::Int8)
+    return NULL;
+
+  // Get stride expr in proper index width.
+  Expr::Width idxWidth = base->index->getWidth();
+  ref<Expr> strideExpr = ConstantExpr::alloc(-1, idxWidth);
+  ref<Expr> offset = ConstantExpr::create(0, idxWidth);
+
+  e = e->getKid(1);
+
+  // concat chains are unbalanced to the right
+  while (e->getKind() == Expr::Concat) {
+    offset = AddExpr::create(offset, strideExpr);
+    if (!isReadExprAtOffset(e->getKid(0), base, offset))
+      return NULL;
+    e = e->getKid(1);
+  }
+
+  offset = AddExpr::create(offset, strideExpr);
+  if (!isReadExprAtOffset(e, base, offset))
+    return NULL;
+
+  return cast<ReadExpr>(e.get());
+}
+
+ref<Expr> AssignmentGenerator::createSubExpr(const ref<Expr> &l, ref<Expr> &r) {
+  return SubExpr::create(r, l);
+}
+
+ref<Expr> AssignmentGenerator::createAddExpr(const ref<Expr> &l, ref<Expr> &r) {
+  return AddExpr::create(r, l);
+}
+
+ref<Expr> AssignmentGenerator::createMulExpr(const ref<Expr> &l, ref<Expr> &r) {
+  return MulExpr::create(r, l);
+}
+
+ref<Expr> AssignmentGenerator::createDivRem(const ref<Expr> &l, ref<Expr> &r,
+                                            bool sign) {
+  if (sign)
+    return SRemExpr::create(r, l);
+  else
+    return URemExpr::create(r, l);
+}
+
+ref<Expr> AssignmentGenerator::createDivExpr(const ref<Expr> &l, ref<Expr> &r,
+                                             bool sign) {
+  if (sign)
+    return SDivExpr::create(r, l);
+  else
+    return UDivExpr::create(r, l);
+}
+
+ref<Expr> AssignmentGenerator::createShlExpr(const ref<Expr> &l, ref<Expr> &r) {
+  return ShlExpr::create(r, l);
+}
+
+ref<Expr> AssignmentGenerator::createLShrExpr(const ref<Expr> &l,
+                                              ref<Expr> &r) {
+  return LShrExpr::create(r, l);
+}
+
+ref<Expr> AssignmentGenerator::createAndExpr(const ref<Expr> &l, ref<Expr> &r) {
+  return AndExpr::create(r, l);
+}
+
+ref<Expr> AssignmentGenerator::createExtractExpr(const ref<Expr> &l,
+                                                 ref<Expr> &r) {
+  return ExtractExpr::create(r, 0, l.get()->getWidth());
+}
+
+ref<Expr> AssignmentGenerator::createExtendExpr(const ref<Expr> &l,
+                                                ref<Expr> &r) {
+  if (l.get()->getKind() == Expr::ZExt) {
+    return ZExtExpr::create(r, l.get()->getWidth());
+  } else {
+    return SExtExpr::create(r, l.get()->getWidth());
+  }
+}
+
+std::vector<unsigned char> AssignmentGenerator::getByteValue(ref<Expr> &val) {
+  std::vector<unsigned char> toReturn;
+  if (ConstantExpr *value = dyn_cast<ConstantExpr>(val)) {
+    for (unsigned w = 0; w < val.get()->getWidth() / 8; ++w) {
+      ref<ConstantExpr> byte = value->Extract(w * 8, Expr::Int8);
+      unsigned char mem_val;
+      byte->toMemory(&mem_val);
+      toReturn.push_back(mem_val);
+    }
+  }
+  return toReturn;
+}
+
+std::vector<unsigned char>
+AssignmentGenerator::getIndexedValue(const std::vector<unsigned char> &c_val,
+                                     ConstantExpr &index,
+                                     const unsigned int size) {
+  std::vector<unsigned char> toReturn;
+  const llvm::APInt index_val = index.getAPValue();
+  assert(!index_val.isSignBit() && "Negative index");
+  const uint64_t id = index_val.getZExtValue() / c_val.size();
+  uint64_t arraySize = size / c_val.size();
+  for (uint64_t i = 0; i < arraySize; ++i) {
+    if (i == id) {
+      for (unsigned arr_i = 0; arr_i < c_val.size(); ++arr_i) {
+        toReturn.push_back(c_val[arr_i]);
+      }
+    } else {
+      for (unsigned arr_i = 0; arr_i < c_val.size(); ++arr_i) {
+        toReturn.push_back(0);
+      }
+    }
+  }
+
+  return toReturn;
+}