Improve handling of constant array in Z3

4 files changed, 72 insertions, 28 deletions

diff --git a/lib/Expr/ExprUtil.cpp b/lib/Expr/ExprUtil.cpp
index ed60a4a9..2106b226 100644
--- a/lib/Expr/ExprUtil.cpp
+++ b/lib/Expr/ExprUtil.cpp
@@ -104,6 +104,21 @@ public:
     : objects(_objects) {}
 };
 
+ExprVisitor::Action ConstantArrayFinder::visitRead(const ReadExpr &re) {
+  const UpdateList &ul = re.updates;
+
+  // FIXME should we memo better than what ExprVisitor is doing for us?
+  for (const UpdateNode *un = ul.head; un; un = un->next) {
+    visit(un->index);
+    visit(un->value);
+  }
+
+  if (ul.root->isConstantArray()) {
+    results.insert(ul.root);
+  }
+
+  return Action::doChildren();
+}
 }
 
 template<typename InputIterator>
diff --git a/lib/Solver/Z3Builder.cpp b/lib/Solver/Z3Builder.cpp
index 6c0653e8..d03c4c89 100644
--- a/lib/Solver/Z3Builder.cpp
+++ b/lib/Solver/Z3Builder.cpp
@@ -101,6 +101,7 @@ Z3Builder::~Z3Builder() {
   // they aren associated with.
   clearConstructCache();
   _arr_hash.clear();
+  constant_array_assertions.clear();
   Z3_del_context(ctx);
   if (z3LogInteractionFile.length() > 0) {
     Z3_close_log();
@@ -400,16 +401,21 @@ Z3ASTHandle Z3Builder::getInitialArray(const Array *root) {
     array_expr = buildArray(unique_name.c_str(), root->getDomain(),
                             root->getRange());
 
-    if (root->isConstantArray()) {
-      // FIXME: Flush the concrete values into Z3. Ideally we would do this
-      // using assertions, which might be faster, but we need to fix the caching
-      // to work correctly in that case.
+    if (root->isConstantArray() && constant_array_assertions.count(root) == 0) {
+      std::vector<Z3ASTHandle> array_assertions;
       for (unsigned i = 0, e = root->size; i != e; ++i) {
-        Z3ASTHandle prev = array_expr;
-        array_expr = writeExpr(
-            prev, construct(ConstantExpr::alloc(i, root->getDomain()), 0),
-            construct(root->constantValues[i], 0));
+        // construct(= (select i root) root->value[i]) to be asserted in
+        // Z3Solver.cpp
+        int width_out;
+        Z3ASTHandle array_value =
+            construct(root->constantValues[i], &width_out);
+        assert(width_out == (int)root->getRange() &&
+               "Value doesn't match root range");
+        array_assertions.push_back(
+            eqExpr(readExpr(array_expr, bvConst32(root->getDomain(), i)),
+                   array_value));
       }
+      constant_array_assertions[root] = std::move(array_assertions);
     }
 
     _arr_hash.hashArrayExpr(root, array_expr);
diff --git a/lib/Solver/Z3Builder.h b/lib/Solver/Z3Builder.h
index a3473f82..7cdb2b5e 100644
--- a/lib/Solver/Z3Builder.h
+++ b/lib/Solver/Z3Builder.h
@@ -10,9 +10,10 @@
 #ifndef __UTIL_Z3BUILDER_H__
 #define __UTIL_Z3BUILDER_H__
 
-#include "klee/util/ExprHashMap.h"
-#include "klee/util/ArrayExprHash.h"
 #include "klee/Config/config.h"
+#include "klee/util/ArrayExprHash.h"
+#include "klee/util/ExprHashMap.h"
+#include <unordered_map>
 #include <z3.h>
 
 namespace klee {
@@ -72,7 +73,7 @@ public:
   // To be specialised
   void dump();
 
-  operator T() { return node; }
+  operator T() const { return node; }
 };
 
 // Specialise for Z3_sort
@@ -169,6 +170,8 @@ private:
 
 public:
   Z3_context ctx;
+  std::unordered_map<const Array *, std::vector<Z3ASTHandle> >
+      constant_array_assertions;
   Z3Builder(bool autoClearConstructCache, const char *z3LogInteractionFile);
   ~Z3Builder();
 
diff --git a/lib/Solver/Z3Solver.cpp b/lib/Solver/Z3Solver.cpp
index 985c143d..e89c9463 100644
--- a/lib/Solver/Z3Solver.cpp
+++ b/lib/Solver/Z3Solver.cpp
@@ -153,19 +153,10 @@ char *Z3SolverImpl::getConstraintLog(const Query &query) {
   // with whatever the solver's builder is set to do.
   Z3Builder temp_builder(/*autoClearConstructCache=*/false,
                          /*z3LogInteractionFile=*/NULL);
-
-  for (std::vector<ref<Expr> >::const_iterator it = query.constraints.begin(),
-                                               ie = query.constraints.end();
-       it != ie; ++it) {
-    assumptions.push_back(temp_builder.construct(*it));
-  }
-  ::Z3_ast *assumptionsArray = NULL;
-  int numAssumptions = query.constraints.size();
-  if (numAssumptions) {
-    assumptionsArray = (::Z3_ast *)malloc(sizeof(::Z3_ast) * numAssumptions);
-    for (int index = 0; index < numAssumptions; ++index) {
-      assumptionsArray[index] = (::Z3_ast)assumptions[index];
-    }
+  ConstantArrayFinder constant_arrays_in_query;
+  for (auto const &constraint : query.constraints) {
+    assumptions.push_back(temp_builder.construct(constraint));
+    constant_arrays_in_query.visit(constraint);
   }
 
   // KLEE Queries are validity queries i.e.
@@ -176,6 +167,25 @@ char *Z3SolverImpl::getConstraintLog(const Query &query) {
   Z3ASTHandle formula = Z3ASTHandle(
       Z3_mk_not(temp_builder.ctx, temp_builder.construct(query.expr)),
       temp_builder.ctx);
+  constant_arrays_in_query.visit(query.expr);
+
+  for (auto const &constant_array : constant_arrays_in_query.results) {
+    assert(builder->constant_array_assertions.count(constant_array) == 1 &&
+           "Constant array found in query, but not handled by Z3Builder");
+    for (auto const &arrayIndexValueExpr :
+         builder->constant_array_assertions[constant_array]) {
+      assumptions.push_back(arrayIndexValueExpr);
+    }
+  }
+
+  ::Z3_ast *assumptionsArray = NULL;
+  int numAssumptions = assumptions.size();
+  if (numAssumptions) {
+    assumptionsArray = (::Z3_ast *)malloc(sizeof(::Z3_ast) * numAssumptions);
+    for (int index = 0; index < numAssumptions; ++index) {
+      assumptionsArray[index] = (::Z3_ast)assumptions[index];
+    }
+  }
 
   ::Z3_string result = Z3_benchmark_to_smtlib_string(
       temp_builder.ctx,
@@ -249,10 +259,10 @@ bool Z3SolverImpl::internalRunSolver(
 
   runStatusCode = SOLVER_RUN_STATUS_FAILURE;
 
-  for (ConstraintManager::const_iterator it = query.constraints.begin(),
-                                         ie = query.constraints.end();
-       it != ie; ++it) {
-    Z3_solver_assert(builder->ctx, theSolver, builder->construct(*it));
+  ConstantArrayFinder constant_arrays_in_query;
+  for (auto const &constraint : query.constraints) {
+    Z3_solver_assert(builder->ctx, theSolver, builder->construct(constraint));
+    constant_arrays_in_query.visit(constraint);
   }
   ++stats::queries;
   if (objects)
@@ -260,6 +270,16 @@ bool Z3SolverImpl::internalRunSolver(
 
   Z3ASTHandle z3QueryExpr =
       Z3ASTHandle(builder->construct(query.expr), builder->ctx);
+  constant_arrays_in_query.visit(query.expr);
+
+  for (auto const &constant_array : constant_arrays_in_query.results) {
+    assert(builder->constant_array_assertions.count(constant_array) == 1 &&
+           "Constant array found in query, but not handled by Z3Builder");
+    for (auto const &arrayIndexValueExpr :
+         builder->constant_array_assertions[constant_array]) {
+      Z3_solver_assert(builder->ctx, theSolver, arrayIndexValueExpr);
+    }
+  }
 
   // KLEE Queries are validity queries i.e.
   // ∀ X Constraints(X) → query(X)