Move merging related code from Executor into MergingSearcher

Co-authored-by: Felix Rath <felix.rath@comsys.rwth-aachen.de>

8 files changed, 95 insertions, 89 deletions

diff --git a/include/klee/MergeHandler.h b/include/klee/MergeHandler.h
index 213264e6..d189e46f 100644
--- a/include/klee/MergeHandler.h
+++ b/include/klee/MergeHandler.h
@@ -83,8 +83,6 @@ 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;
diff --git a/lib/Core/Executor.cpp b/lib/Core/Executor.cpp
index 7ea3aa3b..f90a8909 100644
--- a/lib/Core/Executor.cpp
+++ b/lib/Core/Executor.cpp
@@ -2752,12 +2752,6 @@ void Executor::updateStates(ExecutionState *current) {
     delete es;
   }
   removedStates.clear();
-
-  if (searcher) {
-    searcher->update(nullptr, continuedStates, pausedStates);
-    pausedStates.clear();
-    continuedStates.clear();
-  }
 }
 
 template <typename TypeIt>
@@ -3036,29 +3030,6 @@ std::string Executor::getAddressInfo(ExecutionState &state,
   return info.str();
 }
 
-void Executor::pauseState(ExecutionState &state){
-  auto it = std::find(continuedStates.begin(), continuedStates.end(), &state);
-  // If the state was to be continued, but now gets paused again
-  if (it != continuedStates.end()){
-    // ...just don't continue it
-    std::swap(*it, continuedStates.back());
-    continuedStates.pop_back();
-  } else {
-    pausedStates.push_back(&state);
-  }
-}
-
-void Executor::continueState(ExecutionState &state){
-  auto it = std::find(pausedStates.begin(), pausedStates.end(), &state);
-  // If the state was to be paused, but now gets continued again
-  if (it != pausedStates.end()){
-    // ...don't pause it
-    std::swap(*it, pausedStates.back());
-    pausedStates.pop_back();
-  } else {
-    continuedStates.push_back(&state);
-  }
-}
 
 void Executor::terminateState(ExecutionState &state) {
   if (replayKTest && replayPosition!=replayKTest->numObjects) {
diff --git a/lib/Core/Executor.h b/lib/Core/Executor.h
index f6e58a6b..bf81a1d7 100644
--- a/lib/Core/Executor.h
+++ b/lib/Core/Executor.h
@@ -74,6 +74,7 @@ namespace klee {
   class TimingSolver;
   class TreeStreamWriter;
   class MergeHandler;
+  class MergingSearcher;
   template<class T> class ref;
 
 
@@ -89,7 +90,6 @@ class Executor : public Interpreter {
   friend class SpecialFunctionHandler;
   friend class StatsTracker;
   friend class MergeHandler;
-  friend class MergingSearcher;
 
 public:
   typedef std::pair<ExecutionState*,ExecutionState*> StatePair;
@@ -127,16 +127,6 @@ private:
   TimerGroup timers;
   std::unique_ptr<PTree> processTree;
 
-  /// Keeps track of all currently ongoing merges.
-  /// An ongoing merge is a set of states which branched from a single state
-  /// which ran into a klee_open_merge(), and not all states in the set have
-  /// reached the corresponding klee_close_merge() yet.
-  std::vector<MergeHandler *> mergeGroups;
-
-  /// ExecutionStates 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. 
@@ -148,13 +138,6 @@ private:
   /// \invariant \ref addedStates and \ref removedStates are disjoint.
   std::vector<ExecutionState *> removedStates;
 
-  /// Used to track states that are not terminated, but should not
-  /// be scheduled by the searcher.
-  std::vector<ExecutionState *> pausedStates;
-  /// States that were 'paused' from scheduling, that now may be
-  /// scheduled again
-  std::vector<ExecutionState *> continuedStates;
-
   /// When non-empty the Executor is running in "seed" mode. The
   /// states in this map will be executed in an arbitrary order
   /// (outside the normal search interface) until they terminate. When
@@ -227,6 +210,10 @@ private:
   /// Optimizes expressions
   ExprOptimizer optimizer;
 
+  /// Points to the merging searcher of the searcher chain,
+  /// `nullptr` if merging is disabled
+  MergingSearcher *mergingSearcher = nullptr;
+
   llvm::Function* getTargetFunction(llvm::Value *calledVal,
                                     ExecutionState &state);
   
@@ -410,10 +397,6 @@ private:
 
   bool shouldExitOn(enum TerminateReason termReason);
 
-  // remove state from searcher only
-  void pauseState(ExecutionState& state);
-  // add state to searcher only
-  void continueState(ExecutionState& state);
   // remove state from queue and delete
   void terminateState(ExecutionState &state);
   // call exit handler and terminate state
@@ -526,6 +509,9 @@ public:
 
   /// Returns the errno location in memory of the state
   int *getErrnoLocation(const ExecutionState &state) const;
+
+  MergingSearcher *getMergingSearcher() const { return mergingSearcher; };
+  void setMergingSearcher(MergingSearcher *ms) { mergingSearcher = ms; };
 };
   
 } // End klee namespace
diff --git a/lib/Core/MergeHandler.cpp b/lib/Core/MergeHandler.cpp
index 5832c0d8..5764b5a1 100644
--- a/lib/Core/MergeHandler.cpp
+++ b/lib/Core/MergeHandler.cpp
@@ -11,6 +11,7 @@
 
 #include "CoreStats.h"
 #include "Executor.h"
+#include "Searcher.h"
 #include "klee/ExecutionState.h"
 
 namespace klee {
@@ -54,8 +55,9 @@ unsigned MergeHandler::getInstructionDistance(ExecutionState *es){
 
 ExecutionState *MergeHandler::getPrioritizeState(){
   for (ExecutionState *cur_state : openStates) {
-    bool stateIsClosed = (executor->inCloseMerge.find(cur_state) !=
-                          executor->inCloseMerge.end());
+    bool stateIsClosed =
+        (executor->mergingSearcher->inCloseMerge.find(cur_state) !=
+         executor->mergingSearcher->inCloseMerge.end());
 
     if (!stateIsClosed && (getInstructionDistance(cur_state) < 2 * getMean())) {
       return cur_state;
@@ -95,7 +97,7 @@ void MergeHandler::addClosedState(ExecutionState *es,
   // add a new element to the map
   if (closePoint == reachedCloseMerge.end()) {
     reachedCloseMerge[mp].push_back(es);
-    executor->pauseState(*es);
+    executor->mergingSearcher->pauseState(*es);
   } else {
     // Otherwise try to merge with any state in the map element for this
     // instruction
@@ -105,14 +107,14 @@ void MergeHandler::addClosedState(ExecutionState *es,
     for (auto& mState: cpv) {
       if (mState->merge(*es)) {
         executor->terminateState(*es);
-        executor->inCloseMerge.erase(es);
+        executor->mergingSearcher->inCloseMerge.erase(es);
         mergedSuccessful = true;
         break;
       }
     }
     if (!mergedSuccessful) {
       cpv.push_back(es);
-      executor->pauseState(*es);
+      executor->mergingSearcher->pauseState(*es);
     }
   }
 }
@@ -120,8 +122,8 @@ void MergeHandler::addClosedState(ExecutionState *es,
 void MergeHandler::releaseStates() {
   for (auto& curMergeGroup: reachedCloseMerge) {
     for (auto curState: curMergeGroup.second) {
-      executor->continueState(*curState);
-      executor->inCloseMerge.erase(curState);
+      executor->mergingSearcher->continueState(*curState);
+      executor->mergingSearcher->inCloseMerge.erase(curState);
     }
   }
   reachedCloseMerge.clear();
@@ -134,15 +136,17 @@ bool MergeHandler::hasMergedStates() {
 MergeHandler::MergeHandler(Executor *_executor, ExecutionState *es)
     : executor(_executor), openInstruction(es->steppedInstructions),
       closedMean(0), closedStateCount(0), refCount(0) {
-  executor->mergeGroups.push_back(this);
+  executor->mergingSearcher->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();
+  auto it = std::find(executor->mergingSearcher->mergeGroups.begin(),
+                      executor->mergingSearcher->mergeGroups.end(), this);
+  assert(it != executor->mergingSearcher->mergeGroups.end() &&
+         "All MergeHandlers should be registered in mergeGroups");
+  std::swap(*it, executor->mergingSearcher->mergeGroups.back());
+  executor->mergingSearcher->mergeGroups.pop_back();
 
   releaseStates();
 }
diff --git a/lib/Core/Searcher.cpp b/lib/Core/Searcher.cpp
index 0d5d61e2..e576e8bc 100644
--- a/lib/Core/Searcher.cpp
+++ b/lib/Core/Searcher.cpp
@@ -295,9 +295,8 @@ bool RandomPathSearcher::empty() {
 
 ///
 
-MergingSearcher::MergingSearcher(Executor &_executor, Searcher *_baseSearcher)
-  : executor(_executor),
-  baseSearcher(_baseSearcher){}
+MergingSearcher::MergingSearcher(Searcher *_baseSearcher)
+  : baseSearcher(_baseSearcher){}
 
 MergingSearcher::~MergingSearcher() {
   delete baseSearcher;
@@ -306,8 +305,11 @@ MergingSearcher::~MergingSearcher() {
 ExecutionState& MergingSearcher::selectState() {
   assert(!baseSearcher->empty() && "base searcher is empty");
 
+  if (!UseIncompleteMerge)
+    return baseSearcher->selectState();
+
   // Iterate through all MergeHandlers
-  for (auto cur_mergehandler: executor.mergeGroups) {
+  for (auto cur_mergehandler: mergeGroups) {
     // Find one that has states that could be released
     if (!cur_mergehandler->hasMergedStates()) {
       continue;
diff --git a/lib/Core/Searcher.h b/lib/Core/Searcher.h
index c11c6be5..13941af7 100644
--- a/lib/Core/Searcher.h
+++ b/lib/Core/Searcher.h
@@ -12,6 +12,7 @@
 
 #include "klee/Internal/System/Time.h"
 
+#include "llvm/Support/CommandLine.h"
 #include "llvm/Support/raw_ostream.h"
 
 #include <map>
@@ -183,26 +184,64 @@ namespace klee {
     }
   };
 
+
+  extern llvm::cl::opt<bool> UseIncompleteMerge;
   class MergeHandler;
   class MergingSearcher : public Searcher {
     friend class MergeHandler;
 
     private:
 
-    Executor &executor;
     Searcher *baseSearcher;
 
+    /// States that have been paused by the 'pauseState' function
+    std::vector<ExecutionState*> pausedStates;
+
     public:
-    MergingSearcher(Executor &executor, Searcher *baseSearcher);
+    MergingSearcher(Searcher *baseSearcher);
     ~MergingSearcher();
 
+    /// ExecutionStates currently paused from scheduling because they are
+    /// waiting to be merged in a klee_close_merge instruction
+    std::set<ExecutionState *> inCloseMerge;
+
+    /// Keeps track of all currently ongoing merges.
+    /// An ongoing merge is a set of states (stored in a MergeHandler object)
+    /// which branched from a single state which ran into a klee_open_merge(),
+    /// and not all states in the set have reached the corresponding
+    /// klee_close_merge() yet.
+    std::vector<MergeHandler *> mergeGroups;
+
+    /// Remove state from the searcher chain, while keeping it in the executor.
+    /// This is used here to 'freeze' a state while it is waiting for other
+    /// states in its merge group to reach the same instruction.
+    void pauseState(ExecutionState &state){
+      assert(std::find(pausedStates.begin(), pausedStates.end(), &state) == pausedStates.end());
+      pausedStates.push_back(&state);
+      baseSearcher->removeState(&state);
+    }
+
+    /// Continue a paused state
+    void continueState(ExecutionState &state){
+      auto it = std::find(pausedStates.begin(), pausedStates.end(), &state);
+      assert( it != pausedStates.end());
+      pausedStates.erase(it);
+      baseSearcher->addState(&state);
+    }
+
     ExecutionState &selectState();
 
     void update(ExecutionState *current,
                 const std::vector<ExecutionState *> &addedStates,
                 const std::vector<ExecutionState *> &removedStates) {
-      baseSearcher->update(current, addedStates, removedStates);
+      // We have to check if the current execution state was just deleted, as to
+      // not confuse the nurs searchers
+      if (std::find(pausedStates.begin(), pausedStates.end(), current) ==
+          pausedStates.end()) {
+        baseSearcher->update(current, addedStates, removedStates);
+      }
     }
+
     bool empty() { return baseSearcher->empty(); }
     void printName(llvm::raw_ostream &os) {
       os << "MergingSearcher\n";
diff --git a/lib/Core/SpecialFunctionHandler.cpp b/lib/Core/SpecialFunctionHandler.cpp
index 4acf2e36..bae8085d 100644
--- a/lib/Core/SpecialFunctionHandler.cpp
+++ b/lib/Core/SpecialFunctionHandler.cpp
@@ -12,6 +12,7 @@
 #include "Executor.h"
 #include "Memory.h"
 #include "MemoryManager.h"
+#include "Searcher.h"
 #include "TimingSolver.h"
 
 #include "klee/ExecutionState.h"
@@ -353,16 +354,17 @@ void SpecialFunctionHandler::handleCloseMerge(ExecutionState &state,
   Instruction *i = target->inst;
 
   if (DebugLogMerge)
-    llvm::errs() << "close merge: " << &state << " at " << i << '\n';
+    llvm::errs() << "close merge: " << &state << " at [" << *i << "]\n";
 
   if (state.openMergeStack.empty()) {
     std::ostringstream warning;
     warning << &state << " ran into a close at " << i << " without a preceding open";
     klee_warning("%s", warning.str().c_str());
   } else {
-    assert(executor.inCloseMerge.find(&state) == executor.inCloseMerge.end() &&
+    assert(executor.mergingSearcher->inCloseMerge.find(&state) ==
+               executor.mergingSearcher->inCloseMerge.end() &&
            "State cannot run into close_merge while being closed");
-    executor.inCloseMerge.insert(&state);
+    executor.mergingSearcher->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 c02c349e..3aa07e73 100644
--- a/lib/Core/UserSearcher.cpp
+++ b/lib/Core/UserSearcher.cpp
@@ -125,35 +125,39 @@ Searcher *getNewSearcher(Searcher::CoreSearchType type, Executor &executor) {
 Searcher *klee::constructUserSearcher(Executor &executor) {
 
   Searcher *searcher = getNewSearcher(CoreSearch[0], executor);
-  
+
   if (CoreSearch.size() > 1) {
     std::vector<Searcher *> s;
     s.push_back(searcher);
 
-    for (unsigned i=1; i<CoreSearch.size(); i++)
+    for (unsigned i = 1; i < CoreSearch.size(); i++)
       s.push_back(getNewSearcher(CoreSearch[i], executor));
-    
-    searcher = new InterleavedSearcher(s);
-  }
 
-  if (UseMerge) {
-    if (std::find(CoreSearch.begin(), CoreSearch.end(), Searcher::RandomPath) != CoreSearch.end()){
-      klee_error("use-merge currently does not support random-path, please use another search strategy");
-    }
+    searcher = new InterleavedSearcher(s);
   }
 
   if (UseBatchingSearch) {
-    searcher = new BatchingSearcher(searcher, time::Span(BatchTime), BatchInstructions);
-  }
-
-  if (UseMerge && UseIncompleteMerge) {
-    searcher = new MergingSearcher(executor, searcher);
+    searcher = new BatchingSearcher(searcher, time::Span(BatchTime),
+                                    BatchInstructions);
   }
 
   if (UseIterativeDeepeningTimeSearch) {
     searcher = new IterativeDeepeningTimeSearcher(searcher);
   }
 
+  if (UseMerge) {
+    if (std::find(CoreSearch.begin(), CoreSearch.end(), Searcher::RandomPath) !=
+        CoreSearch.end()) {
+      klee_error("use-merge currently does not support random-path, please use "
+                 "another search strategy");
+    }
+
+    auto *ms = new MergingSearcher(searcher);
+    executor.setMergingSearcher(ms);
+
+    searcher = ms;
+  }
+
   llvm::raw_ostream &os = executor.getHandler().getInfoStream();
 
   os << "BEGIN searcher description\n";