2 files changed, 43 insertions, 21 deletions

diff --git a/lib/Core/Searcher.cpp b/lib/Core/Searcher.cpp
index f5d6543f..bf98ebc7 100644
--- a/lib/Core/Searcher.cpp
+++ b/lib/Core/Searcher.cpp
@@ -412,33 +412,50 @@ void MergingSearcher::printName(llvm::raw_ostream &os) {
 
 ///
 
-BatchingSearcher::BatchingSearcher(Searcher *baseSearcher, time::Span timeBudget, unsigned instructionBudget)
-  : baseSearcher{baseSearcher},
-    timeBudget{timeBudget},
-    instructionBudget{instructionBudget} {};
+BatchingSearcher::BatchingSearcher(Searcher *baseSearcher,
+                                   time::Span timeBudget,
+                                   unsigned instructionBudget)
+    : baseSearcher{baseSearcher}, timeBudgetEnabled{timeBudget},
+      timeBudget{timeBudget}, instructionBudgetEnabled{instructionBudget > 0},
+      instructionBudget{instructionBudget} {};
+
+bool BatchingSearcher::withinTimeBudget() const {
+  return !timeBudgetEnabled ||
+         (time::getWallTime() - lastStartTime) <= timeBudget;
+}
+
+bool BatchingSearcher::withinInstructionBudget() const {
+  return !instructionBudgetEnabled ||
+         (stats::instructions - lastStartInstructions) <= instructionBudget;
+}
 
 ExecutionState &BatchingSearcher::selectState() {
-  if (!lastState ||
-      (((timeBudget.toSeconds() > 0) &&
-        (time::getWallTime() - lastStartTime) > timeBudget)) ||
-      ((instructionBudget > 0) &&
-       (stats::instructions - lastStartInstructions) > instructionBudget)) {
-    if (lastState && timeBudget.toSeconds() > 0) {
-      time::Span delta = time::getWallTime() - lastStartTime;
-      auto t = timeBudget;
-      t *= 1.1;
-      if (delta > t) {
-        klee_message("increased time budget from %f to %f\n", timeBudget.toSeconds(), delta.toSeconds());
-        timeBudget = delta;
-      }
+  if (lastState && withinTimeBudget() && withinInstructionBudget()) {
+    // return same state for as long as possible
+    return *lastState;
+  }
+
+  // ensure time budget is larger than time between two calls (for same state)
+  if (lastState && timeBudgetEnabled) {
+    time::Span delta = time::getWallTime() - lastStartTime;
+    auto t = timeBudget;
+    t *= 1.1;
+    if (delta > t) {
+      klee_message("increased time budget from %f to %f\n",
+                   timeBudget.toSeconds(), delta.toSeconds());
+      timeBudget = delta;
     }
-    lastState = &baseSearcher->selectState();
+  }
+
+  // pick a new state
+  lastState = &baseSearcher->selectState();
+  if (timeBudgetEnabled) {
     lastStartTime = time::getWallTime();
+  }
+  if (instructionBudgetEnabled) {
     lastStartInstructions = stats::instructions;
-    return *lastState;
-  } else {
-    return *lastState;
   }
+  return *lastState;
 }
 
 void BatchingSearcher::update(ExecutionState *current,
diff --git a/lib/Core/Searcher.h b/lib/Core/Searcher.h
index 4eda838d..e399c616 100644
--- a/lib/Core/Searcher.h
+++ b/lib/Core/Searcher.h
@@ -243,13 +243,18 @@ namespace klee {
   /// of instructions.
   class BatchingSearcher final : public Searcher {
     std::unique_ptr<Searcher> baseSearcher;
+    bool timeBudgetEnabled;
     time::Span timeBudget;
+    bool instructionBudgetEnabled;
     unsigned instructionBudget;
 
     ExecutionState *lastState {nullptr};
     time::Point lastStartTime;
     unsigned lastStartInstructions;
 
+    bool withinTimeBudget() const;
+    bool withinInstructionBudget() const;
+
   public:
     /// \param baseSearcher The underlying searcher (takes ownership).
     /// \param timeBudget Time span a state gets selected before choosing a different one.