5 files changed, 79 insertions, 42 deletions
diff --git a/lib/Expr/ArrayExprOptimizer.cpp b/lib/Expr/ArrayExprOptimizer.cpp
index b55974d0..60f2ca6e 100644
--- a/lib/Expr/ArrayExprOptimizer.cpp
+++ b/lib/Expr/ArrayExprOptimizer.cpp
@@ -313,12 +313,9 @@ ref<Expr> ExprOptimizer::getSelectOptExpr(
         arrayValues.push_back(val);
       }
 
-      ref<Expr> index = read->index;
-      IndexCleanerVisitor ice;
-      ice.visit(index);
-      if (ice.getIndex().get()) {
-        index = ice.getIndex();
-      }
+      ref<Expr> index = UDivExpr::create(
+          read->index,
+          ConstantExpr::create(bytesPerElement, read->index->getWidth()));
 
       ref<Expr> opt =
           buildConstantSelectExpr(index, arrayValues, width, elementsInArray);
@@ -650,12 +647,8 @@ ref<Expr> ExprOptimizer::buildMixedSelectExpr(
       }
     }
 
-    ref<Expr> new_index = re->index;
-    IndexCleanerVisitor ice;
-    ice.visit(new_index);
-    if (ice.getIndex().get()) {
-      new_index = ice.getIndex();
-    }
+    ref<Expr> new_index = UDivExpr::create(
+        re->index, ConstantExpr::create(width / 8, re->index->getWidth()));
 
     int new_index_width = new_index->getWidth();
     // Iterate through all the ranges
diff --git a/lib/Expr/ArrayExprVisitor.cpp b/lib/Expr/ArrayExprVisitor.cpp
index 75604104..cada7867 100644
--- a/lib/Expr/ArrayExprVisitor.cpp
+++ b/lib/Expr/ArrayExprVisitor.cpp
@@ -256,19 +256,3 @@ ExprVisitor::Action ArrayValueOptReplaceVisitor::visitRead(const ReadExpr &re) {
   }
   return Action::doChildren();
 }
-
-ExprVisitor::Action IndexCleanerVisitor::visitMul(const MulExpr &e) {
-  if (mul) {
-    if (!isa<ConstantExpr>(e.getKid(0)))
-      index = e.getKid(0);
-    else if (!isa<ConstantExpr>(e.getKid(1)))
-      index = e.getKid(1);
-    mul = false;
-  }
-  return Action::doChildren();
-}
-
-ExprVisitor::Action IndexCleanerVisitor::visitRead(const ReadExpr &) {
-  mul = false;
-  return Action::doChildren();
-}
diff --git a/lib/Expr/ArrayExprVisitor.h b/lib/Expr/ArrayExprVisitor.h
index 37f14cd1..28f485d9 100644
--- a/lib/Expr/ArrayExprVisitor.h
+++ b/lib/Expr/ArrayExprVisitor.h
@@ -124,20 +124,6 @@ public:
                                        bool recursive = true)
       : ExprVisitor(recursive), optimized(_optimized) {}
 };
-
-class IndexCleanerVisitor : public ExprVisitor {
-private:
-  bool mul{true};
-  ref<Expr> index;
-
-protected:
-  Action visitMul(const MulExpr &) override;
-  Action visitRead(const ReadExpr &) override;
-
-public:
-  IndexCleanerVisitor() : ExprVisitor(true) {}
-  inline ref<Expr> getIndex() { return index; }
-};
 } // namespace klee
 
 #endif /* KLEE_ARRAYEXPRVISITOR_H */
diff --git a/test/ArrayOpt/test_expr_arbitrary.c b/test/ArrayOpt/test_expr_arbitrary.c
new file mode 100644
index 00000000..145efea9
--- /dev/null
+++ b/test/ArrayOpt/test_expr_arbitrary.c
@@ -0,0 +1,35 @@
+// RUN: %clang %s -g -emit-llvm %O0opt -c -o %t.bc
+// RUN: rm -rf %t.klee-out
+// RUN: %klee --write-kqueries --use-query-log=all:kquery --output-dir=%t.klee-out %t.bc > %t.log 2>&1
+// RUN: FileCheck %s -input-file=%t.log -check-prefix=CHECK
+// RUN: rm -rf %t.klee-out
+// RUN: %klee --optimize-array=value --write-kqueries --use-query-log=all:kquery --output-dir=%t.klee-out %t.bc > %t.log 2>&1
+// RUN: FileCheck %s -input-file=%t.log -check-prefix=CHECK
+// RUN: FileCheck %s -input-file=%t.log -check-prefix=CHECK-V
+
+#include "klee/klee.h"
+#include <stdio.h>
+
+short array[10] = {42, 1, 42, 42, 2, 5, 6, 42, 8, 9};
+
+int main() {
+  char k;
+  // CHECK-V: KLEE: WARNING: OPT_V: successful
+
+  klee_make_symbolic(&k, sizeof(k), "k");
+  klee_assume(k < 4);
+  klee_assume(k >= 0);
+
+  short *ptrs[4] = {array + 3, array + 0, array + 7, array + 2};
+
+  // CHECK-DAG: Yes
+  if ((*(ptrs[k])) == 42)
+    printf("Yes\n");
+  else
+    printf("No\n");
+
+  // CHECK-DAG: KLEE: done: completed paths = 1
+  // CHECK-NOT: No
+
+  return 0;
+}
diff --git a/test/ArrayOpt/test_expr_mul.c b/test/ArrayOpt/test_expr_mul.c
new file mode 100644
index 00000000..c2b0b22e
--- /dev/null
+++ b/test/ArrayOpt/test_expr_mul.c
@@ -0,0 +1,39 @@
+// RUN: %clang %s -emit-llvm %O0opt -c -o %t.bc
+// RUN: rm -rf %t.klee-out
+// RUN: %klee --write-kqueries --use-query-log=all:kquery --output-dir=%t.klee-out %t.bc > %t.log 2>&1
+// RUN: FileCheck %s -input-file=%t.log -check-prefix=CHECK
+// RUN: rm -rf %t.klee-out
+// RUN: %klee --optimize-array=value --write-kqueries --use-query-log=all:kquery --output-dir=%t.klee-out %t.bc > %t.log 2>&1
+// RUN: FileCheck %s -input-file=%t.log -check-prefix=CHECK
+// RUN: FileCheck %s -input-file=%t.log -check-prefix=CHECK-V
+
+#include "klee/klee.h"
+#include <stdio.h>
+
+char array[10] = {1, 2, 3, 2, 5, 6, 5, 8, 9, 10};
+char array2[10] = {0, 1, 7, 19, 20, 21, 22, 23, 24, 25};
+
+int main() {
+  char k;
+  // CHECK-V: KLEE: WARNING: OPT_V: successful
+
+  klee_make_symbolic(&k, sizeof(k), "k");
+  klee_assume(k < 4);
+  klee_assume(k >= 1);
+
+  // CHECK-DAG: Yes
+  // CHECK-DAG: Yes
+  char r = (array[k] * 3) & 1;
+  char v = array2[r];
+  if (v == 0)
+    printf("Yes\n");
+  else if (v == 1)
+    printf("Yes\n");
+  else
+    printf("No\n");
+
+  // CHECK-DAG: KLEE: done: completed paths = 2
+  // CHECK-NOT: No
+
+  return 0;
+}