Implemented incomplete merging

8 files changed, 159 insertions, 7 deletions

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);
   }