Implemented incomplete merging

11 files changed, 248 insertions, 8 deletions

diff --git a/include/klee/ExecutionState.h b/include/klee/ExecutionState.h
index 6aec61ea..ceb07864 100644
--- a/include/klee/ExecutionState.h
+++ b/include/klee/ExecutionState.h
@@ -149,6 +149,9 @@ public:
   // The objects handling the klee_open_merge calls this state ran through
   std::vector<ref<MergeHandler> > openMergeStack;
 
+  // The numbers of times this state has run through Executor::stepInstruction
+  std::uint64_t steppedInstructions;
+
 private:
   ExecutionState() : ptreeNode(0) {}
 
diff --git a/include/klee/MergeHandler.h b/include/klee/MergeHandler.h
index 0c596825..79357d7c 100644
--- a/include/klee/MergeHandler.h
+++ b/include/klee/MergeHandler.h
@@ -67,6 +67,10 @@ extern llvm::cl::opt<bool> UseMerge;
 
 extern llvm::cl::opt<bool> DebugLogMerge;
 
+extern llvm::cl::opt<bool> UseIncompleteMerge;
+
+extern llvm::cl::opt<bool> DebugLogIncompleteMerge;
+
 class Executor;
 class ExecutionState;
 
@@ -76,6 +80,24 @@ class MergeHandler {
 private:
   Executor *executor;
 
+  /// @brief The instruction count when the state ran into the klee_open_merge
+  uint64_t openInstruction;
+
+  /// @brief The average number of instructions between the open and close merge of each
+  /// state that has finished so far
+  double closeMean;
+
+  /// @brief Number of states that are tracked by this MergeHandler, that ran
+  /// into a relevant klee_close_merge
+  unsigned closedStateCount;
+
+  /// @brief Get distance of state from the openInstruction
+  unsigned getInstrDistance(ExecutionState *es);
+
+  /// @brief States that ran through the klee_open_merge, but not yet into a
+  /// corresponding klee_close_merge
+  std::vector<ExecutionState *> openStates;
+
   /// @brief Mapping the different 'klee_close_merge' calls to the states that ran into
   /// them
   std::map<llvm::Instruction *, std::vector<ExecutionState *> >
@@ -86,6 +108,18 @@ public:
   /// @brief Called when a state runs into a 'klee_close_merge()' call
   void addClosedState(ExecutionState *es, llvm::Instruction *mp);
 
+  /// @brief Return state that should be prioritized to complete this merge
+  ExecutionState *getPrioritizeState();
+
+  /// @brief Add state to the 'openStates' vector
+  void addOpenState(ExecutionState *es);
+
+  /// @brief Remove state from the 'openStates' vector
+  void removeOpenState(ExecutionState *es);
+
+  /// @brief Remove state from the 'inCloseMerge' set in the executor
+  void removeFromCloseMergeSet(ExecutionState *es);
+
   /// @brief True, if any states have run into 'klee_close_merge()' and have
   /// not been released yet
   bool hasMergedStates();
@@ -94,11 +128,15 @@ public:
   /// 'klee_merge_close()'
   void releaseStates();
 
+  // Return the mean time it takes for a state to get from klee_open_merge to
+  // klee_close_merge
+  double getMean();
+
   /// @brief Required by klee::ref objects
   unsigned refCount;
 
 
-  MergeHandler(Executor *_executor);
+  MergeHandler(Executor *_executor, ExecutionState *es);
   ~MergeHandler();
 };
 }
diff --git a/lib/Core/ExecutionState.cpp b/lib/Core/ExecutionState.cpp
index 522bdd15..00e372e4 100644
--- a/lib/Core/ExecutionState.cpp
+++ b/lib/Core/ExecutionState.cpp
@@ -73,7 +73,8 @@ ExecutionState::ExecutionState(KFunction *kf) :
     instsSinceCovNew(0),
     coveredNew(false),
     forkDisabled(false),
-    ptreeNode(0) {
+    ptreeNode(0),
+    steppedInstructions(0){
   pushFrame(0, kf);
 }
 
@@ -90,6 +91,11 @@ ExecutionState::~ExecutionState() {
       delete mo;
   }
 
+  for (auto cur_mergehandler: openMergeStack){
+    cur_mergehandler->removeOpenState(this);
+  }
+
+
   while (!stack.empty()) popFrame();
 }
 
@@ -117,10 +123,14 @@ ExecutionState::ExecutionState(const ExecutionState& state):
     ptreeNode(state.ptreeNode),
     symbolics(state.symbolics),
     arrayNames(state.arrayNames),
-    openMergeStack(state.openMergeStack)
+    openMergeStack(state.openMergeStack),
+    steppedInstructions(state.steppedInstructions)
 {
   for (unsigned int i=0; i<symbolics.size(); i++)
     symbolics[i].first->refCount++;
+
+  for (auto cur_mergehandler: openMergeStack)
+    cur_mergehandler->addOpenState(this);
 }
 
 ExecutionState *ExecutionState::branch() {
diff --git a/lib/Core/Executor.cpp b/lib/Core/Executor.cpp
index 7b054e5b..4fd94dd4 100644
--- a/lib/Core/Executor.cpp
+++ b/lib/Core/Executor.cpp
@@ -1177,6 +1177,7 @@ void Executor::stepInstruction(ExecutionState &state) {
     statsTracker->stepInstruction(state);
 
   ++stats::instructions;
+  ++state.steppedInstructions;
   state.prevPC = state.pc;
   ++state.pc;
 
diff --git a/lib/Core/Executor.h b/lib/Core/Executor.h
index 1bc91be0..852aaf3c 100644
--- a/lib/Core/Executor.h
+++ b/lib/Core/Executor.h
@@ -85,6 +85,7 @@ class Executor : public Interpreter {
   friend class SpecialFunctionHandler;
   friend class StatsTracker;
   friend class MergeHandler;
+  friend class MergingSearcher;
 
 public:
   class Timer {
@@ -133,6 +134,12 @@ private:
   std::vector<TimerInfo*> timers;
   PTree *processTree;
 
+  std::vector<MergeHandler *> mergeGroups;
+
+  // Set of vectors that are currently paused from scheduling because they are
+  // waiting to be merged in a klee_close_merge instruction
+  std::set<ExecutionState *> inCloseMerge;
+
   /// Used to track states that have been added during the current
   /// instructions step. 
   /// \invariant \ref addedStates is a subset of \ref states. 
diff --git a/lib/Core/MergeHandler.cpp b/lib/Core/MergeHandler.cpp
index ad8ffd62..45373002 100644
--- a/lib/Core/MergeHandler.cpp
+++ b/lib/Core/MergeHandler.cpp
@@ -24,8 +24,65 @@ llvm::cl::opt<bool>
         llvm::cl::init(false),
         llvm::cl::desc("Enhanced verbosity for region based merge operations"));
 
+llvm::cl::opt<bool>
+    UseIncompleteMerge("use-incomplete-merge",
+        llvm::cl::init(false),
+        llvm::cl::desc("Heuristic based merging"));
+
+llvm::cl::opt<bool>
+    DebugLogIncompleteMerge("debug-log-incomplete-merge",
+        llvm::cl::init(false),
+        llvm::cl::desc("Debug info about incomplete merging"));
+
+double MergeHandler::getMean() {
+  if (closedStateCount == 0)
+    return 0;
+
+  return closeMean;
+}
+
+unsigned MergeHandler::getInstrDistance(ExecutionState *es){
+  return es->steppedInstructions - openInstruction;
+}
+
+ExecutionState *MergeHandler::getPrioritizeState(){
+  for (ExecutionState *cur_state : openStates) {
+    bool stateIsClosed = (executor->inCloseMerge.find(cur_state) !=
+                          executor->inCloseMerge.end());
+
+    if (!stateIsClosed && (getInstrDistance(cur_state) < 2 * getMean())) {
+      return cur_state;
+    }
+  }
+  return 0;
+}
+
+
+void MergeHandler::addOpenState(ExecutionState *es){
+  openStates.push_back(es);
+}
+
+void MergeHandler::removeOpenState(ExecutionState *es){
+  auto it = std::find(openStates.begin(), openStates.end(), es);
+  assert(it != openStates.end());
+  std::swap(*it, openStates.back());
+  openStates.pop_back();
+}
+
+void MergeHandler::removeFromCloseMergeSet(ExecutionState *es){
+  executor->inCloseMerge.erase(es);
+}
+
 void MergeHandler::addClosedState(ExecutionState *es,
                                          llvm::Instruction *mp) {
+  // Update stats
+  ++closedStateCount;
+  closeMean += (static_cast<double>(getInstrDistance(es)) - closeMean) /
+               closedStateCount;
+
+  // Remove from openStates
+  removeOpenState(es);
+
   auto closePoint = reachedMergeClose.find(mp);
 
   // If no other state has yet encountered this klee_close_merge instruction,
@@ -66,11 +123,19 @@ bool MergeHandler::hasMergedStates() {
   return (!reachedMergeClose.empty());
 }
 
-MergeHandler::MergeHandler(Executor *_executor)
-    : executor(_executor), refCount(0) {
+MergeHandler::MergeHandler(Executor *_executor, ExecutionState *es)
+    : executor(_executor), openInstruction(es->steppedInstructions),
+      closedStateCount(0), refCount(0) {
+  executor->mergeGroups.push_back(this);
+  addOpenState(es);
 }
 
 MergeHandler::~MergeHandler() {
+  auto it = std::find(executor->mergeGroups.begin(),
+                      executor->mergeGroups.end(), this);
+  std::swap(*it, executor->mergeGroups.back());
+  executor->mergeGroups.pop_back();
+
   releaseStates();
 }
 }
diff --git a/lib/Core/Searcher.cpp b/lib/Core/Searcher.cpp
index 05ebe6a6..dfeb15de 100644
--- a/lib/Core/Searcher.cpp
+++ b/lib/Core/Searcher.cpp
@@ -15,6 +15,7 @@
 #include "StatsTracker.h"
 
 #include "klee/ExecutionState.h"
+#include "klee/MergeHandler.h"
 #include "klee/Statistics.h"
 #include "klee/Internal/Module/InstructionInfoTable.h"
 #include "klee/Internal/Module/KInstruction.h"
@@ -42,6 +43,7 @@
 using namespace klee;
 using namespace llvm;
 
+
 namespace klee {
   extern RNG theRNG;
 }
@@ -257,7 +259,6 @@ bool WeightedRandomSearcher::empty() {
 }
 
 ///
-
 RandomPathSearcher::RandomPathSearcher(Executor &_executor)
   : executor(_executor) {
 }
@@ -268,7 +269,6 @@ RandomPathSearcher::~RandomPathSearcher() {
 ExecutionState &RandomPathSearcher::selectState() {
   unsigned flips=0, bits=0;
   PTree::Node *n = executor.processTree->root;
-  
   while (!n->data) {
     if (!n->left) {
       n = n->right;
@@ -299,6 +299,42 @@ bool RandomPathSearcher::empty() {
 
 ///
 
+MergingSearcher::MergingSearcher(Executor &_executor, Searcher *_baseSearcher)
+  : executor(_executor),
+  baseSearcher(_baseSearcher){}
+
+MergingSearcher::~MergingSearcher() {
+  delete baseSearcher;
+}
+
+ExecutionState& MergingSearcher::selectState() {
+  assert(!baseSearcher->empty() && "base searcher is empty");
+
+  // Iterate through all MergeHandlers
+  for (auto cur_mergehandler: executor.mergeGroups) {
+    // Find one that has states that could be released
+    if (!cur_mergehandler->hasMergedStates()) {
+      continue;
+    }
+    // Find a state that can be prioritized
+    ExecutionState *es = cur_mergehandler->getPrioritizeState();
+    if (es) {
+      return *es;
+    } else {
+      if (DebugLogIncompleteMerge){
+        llvm::errs() << "Preemptively releasing states\n";
+      }
+      // If no state can be prioritized, they all exceeded the amount of time we
+      // are willing to wait for them. Release the states that already arrived at close_merge.
+      cur_mergehandler->releaseStates();
+    }
+  }
+  // If we were not able to prioritize a merging state, just return some state
+  return baseSearcher->selectState();
+}
+
+///
+
 BatchingSearcher::BatchingSearcher(Searcher *_baseSearcher,
                                    double _timeBudget,
                                    unsigned _instructionBudget) 
diff --git a/lib/Core/Searcher.h b/lib/Core/Searcher.h
index 27c8aed1..47026410 100644
--- a/lib/Core/Searcher.h
+++ b/lib/Core/Searcher.h
@@ -180,6 +180,32 @@ namespace klee {
     }
   };
 
+  class MergeHandler;
+  class MergingSearcher : public Searcher {
+    friend class MergeHandler;
+
+    private:
+
+    Executor &executor;
+    Searcher *baseSearcher;
+
+    public:
+    MergingSearcher(Executor &executor, Searcher *baseSearcher);
+    ~MergingSearcher();
+
+    ExecutionState &selectState();
+
+    void update(ExecutionState *current,
+                const std::vector<ExecutionState *> &addedStates,
+                const std::vector<ExecutionState *> &removedStates) {
+      baseSearcher->update(current, addedStates, removedStates);
+    }
+    bool empty() { return baseSearcher->empty(); }
+    void printName(llvm::raw_ostream &os) {
+      os << "MergingSearcher\n";
+    }
+  };
+
   class BatchingSearcher : public Searcher {
     Searcher *baseSearcher;
     double timeBudget;
diff --git a/lib/Core/SpecialFunctionHandler.cpp b/lib/Core/SpecialFunctionHandler.cpp
index 75456856..e927adf0 100644
--- a/lib/Core/SpecialFunctionHandler.cpp
+++ b/lib/Core/SpecialFunctionHandler.cpp
@@ -342,7 +342,7 @@ void SpecialFunctionHandler::handleOpenMerge(ExecutionState &state,
   }
 
   state.openMergeStack.push_back(
-      ref<MergeHandler>(new MergeHandler(&executor)));
+      ref<MergeHandler>(new MergeHandler(&executor, &state)));
 
   if (DebugLogMerge)
     llvm::errs() << "open merge: " << &state << "\n";
@@ -365,6 +365,7 @@ void SpecialFunctionHandler::handleCloseMerge(ExecutionState &state,
     warning << &state << " ran into a close at " << i << " without a preceding open";
     klee_warning("%s", warning.str().c_str());
   } else {
+    executor.inCloseMerge.insert(&state);
     state.openMergeStack.back()->addClosedState(&state, i);
     state.openMergeStack.pop_back();
   }
diff --git a/lib/Core/UserSearcher.cpp b/lib/Core/UserSearcher.cpp
index f0b2f85a..5e88df7c 100644
--- a/lib/Core/UserSearcher.cpp
+++ b/lib/Core/UserSearcher.cpp
@@ -14,6 +14,8 @@
 
 #include "klee/Internal/Support/ErrorHandling.h"
 #include "klee/CommandLine.h"
+#include "klee/MergeHandler.h"
+
 
 #include "llvm/Support/CommandLine.h"
 
@@ -121,6 +123,10 @@ Searcher *klee::constructUserSearcher(Executor &executor) {
     searcher = new BatchingSearcher(searcher, BatchTime, BatchInstructions);
   }
 
+  if (UseMerge && UseIncompleteMerge) {
+    searcher = new MergingSearcher(executor, searcher);
+  }
+
   if (UseIterativeDeepeningTimeSearch) {
     searcher = new IterativeDeepeningTimeSearcher(searcher);
   }
diff --git a/test/Merging/incomplete_merge.c b/test/Merging/incomplete_merge.c
new file mode 100644
index 00000000..1962cfaa
--- /dev/null
+++ b/test/Merging/incomplete_merge.c
@@ -0,0 +1,47 @@
+// RUN: %llvmgcc -emit-llvm -g -c -o %t.bc %s
+// RUN: rm -rf %t.klee-out
+// RUN: klee --output-dir=%t.klee-out --use-merge --debug-log-merge --use-incomplete-merge --debug-log-incomplete-merge --search=nurs:covnew --use-batching-search %t.bc 2>&1 | FileCheck %s
+// RUN: rm -rf %t.klee-out
+// RUN: klee --output-dir=%t.klee-out --use-merge --debug-log-merge --use-incomplete-merge --debug-log-incomplete-merge --search=bfs --use-batching-search %t.bc 2>&1 | FileCheck %s
+// RUN: rm -rf %t.klee-out
+// RUN: klee --output-dir=%t.klee-out --use-merge --debug-log-merge --use-incomplete-merge --debug-log-incomplete-merge --search=dfs --use-batching-search %t.bc 2>&1 | FileCheck %s
+// RUN: rm -rf %t.klee-out
+// RUN: klee --output-dir=%t.klee-out --use-merge --debug-log-merge --use-incomplete-merge --debug-log-incomplete-merge --search=nurs:covnew %t.bc 2>&1 | FileCheck %s
+
+// CHECK: open merge:
+// 5 close merges
+// CHECK: close merge:
+// CHECK: close merge:
+// CHECK: close merge:
+// CHECK: close merge:
+// CHECK: close merge:
+
+#include <klee/klee.h>
+
+int main(int argc, char **args) {
+
+  int x;
+  int a;
+  int foo = 0;
+
+  klee_make_symbolic(&x, sizeof(x), "x");
+  klee_make_symbolic(&a, sizeof(a), "a");
+
+  klee_open_merge();
+  if (a == 0) {
+    klee_open_merge();
+
+    if (x == 1) {
+      foo = 5;
+    } else if (x == 2) {
+      foo = 6;
+    } else {
+      foo = 7;
+    }
+
+    klee_close_merge();
+  }
+  klee_close_merge();
+
+  return foo;
+}