[feat] Add BlockLevelSearcher

Author: misonijnik

include/klee/Module/KModule.h

@@ -197,6 +197,12 @@ struct KBlockCompare {
   }
 };
 
+struct KFunctionCompare {
+  bool operator()(const KFunction *a, const KFunction *b) const {
+    return a->id < b->id;
+  }
+};
+
 class KConstant {
 public:
   /// Actual LLVM constant this represents.

lib/Core/ExecutionState.cpp

@@ -443,10 +443,13 @@ void ExecutionState::increaseLevel() {
   KFunction *kf = prevPC->parent->parent;
   KModule *kmodule = kf->parent;
 
-  if (prevPC->inst->isTerminator() && kmodule->inMainModule(*kf->function)) {
+  if (prevPC->inst->isTerminator()) {
     auto srcLevel = stack.infoStack().back().multilevel[srcbb].second;
     stack.infoStack().back().multilevel.replace({srcbb, srcLevel + 1});
+    stack.infoStack().back().maxMultilevel =
+        std::max(stack.infoStack().back().maxMultilevel, srcLevel + 1);
     level.insert(prevPC->parent);
+    stack.infoStack().back().level.insert(prevPC->parent);
   }
   if (srcbb != dstbb) {
     transitionLevel.insert(std::make_pair(srcbb, dstbb));

lib/Core/ExecutionState.h

@@ -95,6 +95,8 @@ struct InfoStackFrame {
   KFunction *kf;
   CallPathNode *callPathNode = nullptr;
   PersistentMap<llvm::BasicBlock *, unsigned long long> multilevel;
+  unsigned long long maxMultilevel = 0;
+  std::set<KBlock *, KBlockCompare> level;
 
   /// Minimum distance to an uncovered instruction once the function
   /// returns. This is not a good place for this but is used to

lib/Core/Executor.cpp

@@ -6663,7 +6663,7 @@ void Executor::runFunctionAsMain(Function *f, int argc, char **argv,
 
   if (guidanceKind == GuidanceKind::ErrorGuidance) {
     std::map<klee::KFunction *, klee::ref<klee::TargetForest>,
-             klee::TargetedExecutionManager::KFunctionLess>
+             klee::KFunctionCompare>
         prepTargets;
     if (FunctionCallReproduce == "") {
       auto &paths = interpreterOpts.Paths.value();

lib/Core/Searcher.cpp

@@ -195,7 +195,7 @@ weight_type TargetedSearcher::getWeight(ExecutionState *es) {
     return weight;
   }
   auto distRes = distanceCalculator.getDistance(*es, target->getBlock());
-  weight = ulog2(distRes.weight + es->steppedMemoryInstructions + 1); // [0, 32)
+  weight = ulog2(distRes.weight + 1); // [0, 32)
   if (!distRes.isInsideFunction) {
     weight += 32; // [32, 64)
   }
@@ -206,12 +206,12 @@ weight_type TargetedSearcher::getWeight(ExecutionState *es) {
 
 ExecutionState &GuidedSearcher::selectState() {
   unsigned size = historiesAndTargets.size();
-  index = theRNG.getInt32() % (size + 1);
+  interleave ^= 1;
   ExecutionState *state = nullptr;
-  if (index == size) {
+  if (interleave || !size) {
     state = &baseSearcher->selectState();
   } else {
-    index = index % size;
+    index = theRNG.getInt32() % size;
     auto &historyTargetPair = historiesAndTargets[index];
     ref<const TargetsHistory> history = historyTargetPair.first;
     ref<Target> target = historyTargetPair.second;
@@ -776,3 +776,159 @@ void InterleavedSearcher::printName(llvm::raw_ostream &os) {
     searcher->printName(os);
   os << "</InterleavedSearcher>\n";
 }
+
+///
+
+BlockLevelSearcher::BlockLevelSearcher(RNG &rng) : theRNG{rng} {}
+
+ExecutionState &BlockLevelSearcher::selectState() {
+  unsigned rnd = 0;
+  unsigned index = 0;
+  unsigned mod = 10;
+  unsigned border = 9;
+
+  auto kfi = data.begin();
+  index = theRNG.getInt32() % data.size();
+  std::advance(kfi, index);
+  auto &sizesTo = kfi->second;
+
+  for (auto &sizesSize : sizesTo) {
+    rnd = theRNG.getInt32();
+    if (rnd % mod < border) {
+      for (auto &size : sizesSize.second) {
+        rnd = theRNG.getInt32();
+        if (rnd % mod < border) {
+          auto lbi = size.second.begin();
+          index = theRNG.getInt32() % size.second.size();
+          std::advance(lbi, index);
+          auto &level = *lbi;
+          for (auto &levelMultilevels : level.second) {
+            rnd = theRNG.getInt32();
+            if (rnd % mod < border) {
+              auto si = levelMultilevels.second.begin();
+              index = theRNG.getInt32() % levelMultilevels.second.size();
+              std::advance(si, index);
+              auto &state = *si;
+              return *state;
+            }
+          }
+        }
+      }
+    }
+  }
+
+  return **(sizesTo.begin()
+                ->second.begin()
+                ->second.begin()
+                ->second.begin()
+                ->second.begin());
+}
+
+void BlockLevelSearcher::update(ExecutionState *current,
+                                const StateIterable &addedStates,
+                                const StateIterable &removedStates) {
+  if (current && std::find(removedStates.begin(), removedStates.end(),
+                           current) == removedStates.end()) {
+    KFunction *kf = current->initPC->parent->parent;
+    auto &sizesTo = data[kf];
+    auto &sizeTo = sizesTo[stateToSize[current]];
+    auto &levelTo = sizeTo[stateToSizes[current]];
+    auto &multilevelTo = levelTo[stateToLevel[current]];
+    auto &states = multilevelTo[stateToMultilevel[current]];
+    sizes.clear();
+    unsigned long long maxMultilevel = 0u;
+    for (auto &infoFrame : current->stack.infoStack()) {
+      sizes.push_back(infoFrame.level.size());
+      maxMultilevel = std::max(maxMultilevel, infoFrame.maxMultilevel);
+    }
+    for (auto &kfLevel : current->stack.multilevel) {
+      maxMultilevel = std::max(maxMultilevel, kfLevel.second);
+    }
+    if (sizes != stateToSizes[current] ||
+        current->level.size() != stateToLevel[current].size() ||
+        maxMultilevel != stateToMultilevel[current]) {
+      states.erase(current);
+      if (states.size() == 0) {
+        multilevelTo.erase(stateToMultilevel[current]);
+      }
+      if (multilevelTo.size() == 0) {
+        levelTo.erase(stateToLevel[current]);
+      }
+      if (levelTo.size() == 0) {
+        sizeTo.erase(stateToSizes[current]);
+      }
+      if (sizeTo.size() == 0) {
+        sizesTo.erase(stateToSize[current]);
+      }
+
+      sizesTo[current->level.size()][sizes][current->level][maxMultilevel]
+          .insert(current);
+
+      stateToLevel[current] = current->level;
+      stateToMultilevel[current] = maxMultilevel;
+      stateToSizes[current] = sizes;
+      stateToSize[current] = current->level.size();
+    }
+  }
+
+  for (const auto state : addedStates) {
+    KFunction *kf = state->initPC->parent->parent;
+    auto &sizesTo = data[kf];
+
+    sizes.clear();
+    unsigned long long maxMultilevel = 0u;
+    for (auto &infoFrame : state->stack.infoStack()) {
+      sizes.push_back(infoFrame.level.size());
+      maxMultilevel = std::max(maxMultilevel, infoFrame.maxMultilevel);
+    }
+    for (auto &kfLevel : state->stack.multilevel) {
+      maxMultilevel = std::max(maxMultilevel, kfLevel.second);
+    }
+
+    sizesTo[state->level.size()][sizes][state->level][maxMultilevel].insert(
+        state);
+
+    stateToLevel[state] = state->level;
+    stateToMultilevel[state] = maxMultilevel;
+    stateToSize[state] = state->level.size();
+    stateToSizes[state] = sizes;
+  }
+
+  // remove states
+  for (const auto state : removedStates) {
+    KFunction *kf = state->initPC->parent->parent;
+    auto &sizesTo = data[kf];
+    auto &sizeTo = sizesTo[stateToSize[state]];
+    auto &levelTo = sizeTo[stateToSizes[state]];
+    auto &multilevelTo = levelTo[stateToLevel[state]];
+    auto &states = multilevelTo[stateToMultilevel[state]];
+
+    states.erase(state);
+    if (states.size() == 0) {
+      multilevelTo.erase(stateToMultilevel[state]);
+    }
+    if (multilevelTo.size() == 0) {
+      levelTo.erase(stateToLevel[state]);
+    }
+    if (levelTo.size() == 0) {
+      sizeTo.erase(stateToSizes[state]);
+    }
+    if (sizeTo.size() == 0) {
+      sizesTo.erase(stateToSize[state]);
+    }
+    if (sizesTo.size() == 0) {
+      data.erase(kf);
+    }
+
+    stateToLevel.erase(state);
+    stateToMultilevel.erase(state);
+    stateToSizes.erase(state);
+    stateToSize.erase(state);
+  }
+}
+
+bool BlockLevelSearcher::empty() { return stateToLevel.empty(); }
+
+void BlockLevelSearcher::printName(llvm::raw_ostream &os) {
+  os << "BlockLevelSearcher\n";
+}

lib/Core/Searcher.h

@@ -76,6 +76,7 @@ class Searcher {
   enum CoreSearchType : std::uint8_t {
     DFS,
     BFS,
+    BlockLevelSearch,
     RandomState,
     RandomPath,
     NURS_CovNew,
@@ -169,6 +170,7 @@ class GuidedSearcher final : public Searcher, public TargetManagerSubscriber {
   DistanceCalculator &distanceCalculator;
   RNG &theRNG;
   unsigned index{1};
+  bool interleave = true;
 
   TargetForestHistoryTargetSet localHistoryTargets;
   std::vector<ExecutionState *> baseAddedStates;
@@ -363,6 +365,73 @@ class InterleavedSearcher final : public Searcher {
   void printName(llvm::raw_ostream &os) override;
 };
 
+class BlockLevelSearcher final : public Searcher {
+  struct KBlockSetsCompare {
+    bool operator()(const std::set<KBlock *, KBlockCompare> &a,
+                    const std::set<KBlock *, KBlockCompare> &b) const {
+      return a.size() > b.size() || (a.size() == b.size() && a > b);
+    }
+  };
+
+  struct MultilevelCompare {
+    bool operator()(
+        const std::map<KBlock *, unsigned long long, KBlockCompare> &a,
+        const std::map<KBlock *, unsigned long long, KBlockCompare> &b) const {
+      return a < b;
+    }
+  };
+
+  struct VectorsCompare {
+    bool operator()(const std::vector<unsigned> &a,
+                    const std::vector<unsigned> &b) const {
+      auto ai = a.begin(), ae = a.end(), bi = b.begin(), be = b.end();
+      for (; ai != ae && bi != be; ++ai, ++bi) {
+        if (*ai != *bi) {
+          return *ai > *bi;
+        }
+      }
+      return bi != be;
+    }
+  };
+
+  using MultilevelsToStates =
+      std::map<unsigned long long,
+               std::set<ExecutionState *, ExecutionStateIDCompare>,
+               std::less<unsigned long long>>;
+
+  using LevelToMultilevelsToStates =
+      std::map<std::set<KBlock *, KBlockCompare>, MultilevelsToStates,
+               KBlockSetsCompare>;
+
+  using SizesToLevelToMultilevelsToStates =
+      std::map<std::vector<unsigned>, LevelToMultilevelsToStates,
+               VectorsCompare>;
+  using SizeToSizesToLevelToMultilevelsToStates =
+      std::map<unsigned, SizesToLevelToMultilevelsToStates,
+               std::greater<unsigned>>;
+
+  using FunctionToSizeToSizesToLevelToMultilevelsToStates =
+      std::map<KFunction *, SizeToSizesToLevelToMultilevelsToStates,
+               KFunctionCompare>;
+
+  FunctionToSizeToSizesToLevelToMultilevelsToStates data;
+  std::unordered_map<ExecutionState *, std::set<KBlock *, KBlockCompare>>
+      stateToLevel;
+  std::unordered_map<ExecutionState *, unsigned long long> stateToMultilevel;
+  std::unordered_map<ExecutionState *, std::vector<unsigned>> stateToSizes;
+  std::unordered_map<ExecutionState *, unsigned> stateToSize;
+  std::vector<unsigned> sizes;
+  RNG &theRNG;
+
+public:
+  explicit BlockLevelSearcher(RNG &rng);
+  ExecutionState &selectState() override;
+  void update(ExecutionState *current, const StateIterable &addedStates,
+              const StateIterable &removedStates) override;
+  bool empty() override;
+  void printName(llvm::raw_ostream &os) override;
+};
+
 } // namespace klee
 
 #endif /* KLEE_SEARCHER_H */

lib/Core/TargetedExecutionManager.cpp

@@ -493,13 +493,12 @@ KFunction *TargetedExecutionManager::tryResolveEntryFunction(
   return resKf;
 }
 
-std::map<KFunction *, ref<TargetForest>,
-         TargetedExecutionManager::KFunctionLess>
+std::map<KFunction *, ref<TargetForest>, KFunctionCompare>
 TargetedExecutionManager::prepareTargets(KModule *kmodule, SarifReport paths) {
   Locations locations = collectAllLocations(paths);
   LocationToBlocks locToBlocks = prepareAllLocations(kmodule, locations);
 
-  std::map<KFunction *, ref<TargetForest>, KFunctionLess> whitelists;
+  std::map<KFunction *, ref<TargetForest>, KFunctionCompare> whitelists;
 
   for (auto &result : paths.results) {
     bool isFullyResolved = tryResolveLocations(result, locToBlocks);

lib/Core/TargetedExecutionManager.h

@@ -115,17 +115,11 @@ class TargetedExecutionManager {
   StatesSet localStates;
 
 public:
-  struct KFunctionLess {
-    bool operator()(const KFunction *a, const KFunction *b) const {
-      return a->id < b->id;
-    }
-  };
-
   explicit TargetedExecutionManager(CodeGraphInfo &codeGraphInfo_,
                                     TargetManager &targetManager_)
       : codeGraphInfo(codeGraphInfo_), targetManager(targetManager_) {}
   ~TargetedExecutionManager() = default;
-  std::map<KFunction *, ref<TargetForest>, KFunctionLess>
+  std::map<KFunction *, ref<TargetForest>, KFunctionCompare>
   prepareTargets(KModule *kmodule, SarifReport paths);
 
   void reportFalseNegative(ExecutionState &state, ReachWithError error);

lib/Core/UserSearcher.cpp

@@ -38,6 +38,9 @@ cl::list<Searcher::CoreSearchType> CoreSearch(
         clEnumValN(Searcher::BFS, "bfs",
                    "use Breadth First Search (BFS), where scheduling decisions "
                    "are taken at the level of (2-way) forks"),
+        clEnumValN(Searcher::BlockLevelSearch, "bls",
+                   "use Block Level Search (BLS), where selection "
+                   "between states with same level is is carried out randomly"),
         clEnumValN(Searcher::RandomState, "random-state",
                    "randomly select a state to explore"),
         clEnumValN(Searcher::RandomPath, "random-path",
@@ -92,7 +95,7 @@ void klee::initializeSearchOptions() {
   // default values
   if (CoreSearch.empty()) {
     CoreSearch.push_back(Searcher::RandomPath);
-    CoreSearch.push_back(Searcher::NURS_CovNew);
+    CoreSearch.push_back(Searcher::BlockLevelSearch);
   }
 }
 
@@ -119,6 +122,9 @@ Searcher *getNewSearcher(Searcher::CoreSearchType type, RNG &rng,
   case Searcher::BFS:
     searcher = new BFSSearcher();
     break;
+  case Searcher::BlockLevelSearch:
+    searcher = new BlockLevelSearcher(rng);
+    break;
   case Searcher::RandomState:
     searcher = new RandomSearcher(rng);
     break;

test/Feature/BlockPath.ll

@@ -1,7 +1,7 @@
 ; REQUIRES: geq-llvm-12.0
 ; RUN: %llvmas %s -o %t1.bc
 ; RUN: rm -rf %t.klee-out
-; RUN: %klee --output-dir=%t.klee-out --libc=klee --write-kpaths %t1.bc
+; RUN: %klee --output-dir=%t.klee-out --libc=klee --search=dfs --write-kpaths %t1.bc
 ; RUN: grep "(path: 0 (main: %0 %5 %6 %8 (abs: %1 %8 %11 %13) %8 %10 %16 %17 %19" %t.klee-out/test000001.kpath
 ; RUN: grep "(path: 0 (main: %0 %5 %6 %8 (abs: %1 %6 %11 %13) %8 %10 %16 %17 %19" %t.klee-out/test000002.kpath