[feat] Add BlockLevelSearcher

Author: misonijnik

include/klee/Module/KModule.h

@@ -187,6 +187,9 @@ struct KFunction : public KCallable {
   /// Unique index for KFunction and KInstruction inside KModule
   /// from 0 to [KFunction + KInstruction]
   [[nodiscard]] inline unsigned getGlobalIndex() const { return globalIndex; }
+
+  bool operator<(const KFunction &rhs) const { return id < rhs.id; }
+  bool operator<(const KFunction *rhs) const { return id < rhs->id; }
 };
 
 struct KBlockCompare {
@@ -197,6 +200,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

@@ -446,7 +446,10 @@ void ExecutionState::increaseLevel() {
   if (prevPC->inst->isTerminator() && kmodule->inMainModule(*kf->function)) {
     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);
   }
 }
 

lib/Core/ExecutionState.h

@@ -96,6 +96,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

@@ -6763,7 +6763,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

@@ -181,8 +181,7 @@ weight_type TargetedSearcher::getWeight(ExecutionState *es) {
     return weight;
   }
   auto distRes = distanceCalculator.getDistance(*es, target->getBlock());
-  weight = klee::util::ulog2(distRes.weight + es->steppedMemoryInstructions +
-                             1); // [0, 32)
+  weight = klee::util::ulog2(distRes.weight + 1); // [0, 32)
   if (!distRes.isInsideFunction) {
     weight += 32; // [32, 64)
   }
@@ -193,12 +192,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;
@@ -763,3 +762,124 @@ 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;
+          auto si = level.second.begin();
+          index = theRNG.getInt32() % level.second.size();
+          std::advance(si, index);
+          auto &state = *si;
+          return *state;
+        }
+      }
+    }
+  }
+
+  return **(sizesTo.begin()->second.begin()->second.begin()->second.begin());
+}
+
+void BlockLevelSearcher::clear(ExecutionState &state) {
+  KFunction *kf = state.initPC->parent->parent;
+  BlockLevel &bl = stateToBlockLevel[&state];
+  auto &sizeTo = data[kf];
+  auto &sizesTo = sizeTo[bl.sizeOfLevel];
+  auto &levelTo = sizesTo[bl.sizesOfFrameLevels];
+  auto &states = levelTo[bl.maxMultilevel];
+
+  states.erase(&state);
+  if (states.size() == 0) {
+    levelTo.erase(bl.maxMultilevel);
+  }
+  if (levelTo.size() == 0) {
+    sizesTo.erase(bl.sizesOfFrameLevels);
+  }
+  if (sizesTo.size() == 0) {
+    sizeTo.erase(bl.sizeOfLevel);
+  }
+  if (sizeTo.size() == 0) {
+    data.erase(kf);
+  }
+}
+
+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;
+    BlockLevel &bl = stateToBlockLevel[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 != bl.sizesOfFrameLevels ||
+        current->level.size() != bl.sizeOfLevel ||
+        maxMultilevel != bl.maxMultilevel) {
+      clear(*current);
+
+      data[kf][current->level.size()][sizes][maxMultilevel].insert(current);
+
+      stateToBlockLevel[current] =
+          BlockLevel(kf, current->level.size(), sizes, maxMultilevel);
+    }
+  }
+
+  for (const auto state : addedStates) {
+    KFunction *kf = state->initPC->parent->parent;
+
+    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);
+    }
+
+    data[kf][state->level.size()][sizes][maxMultilevel].insert(state);
+
+    stateToBlockLevel[state] =
+        BlockLevel(kf, state->level.size(), sizes, maxMultilevel);
+  }
+
+  // remove states
+  for (const auto state : removedStates) {
+    clear(*state);
+    stateToBlockLevel.erase(state);
+  }
+}
+
+bool BlockLevelSearcher::empty() { return stateToBlockLevel.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,68 @@ class InterleavedSearcher final : public Searcher {
   void printName(llvm::raw_ostream &os) override;
 };
 
+class BlockLevelSearcher final : public Searcher {
+private:
+  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 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;
+    }
+  };
+
+  struct BlockLevel {
+    KFunction *kf;
+    unsigned sizeOfLevel;
+    std::vector<unsigned> sizesOfFrameLevels;
+    unsigned long long maxMultilevel;
+    BlockLevel(KFunction *kf, unsigned sizeOfLevel,
+               const std::vector<unsigned> &sizesOfFrameLevels,
+               unsigned long long maxMultilevel)
+        : kf(kf), sizeOfLevel(sizeOfLevel),
+          sizesOfFrameLevels(sizesOfFrameLevels), maxMultilevel(maxMultilevel) {
+    }
+    BlockLevel() : kf(nullptr), sizeOfLevel(0), maxMultilevel(0){};
+  };
+
+  using states_ty = std::set<ExecutionState *, ExecutionStateIDCompare>;
+  using kblocks_ty = std::set<KBlock *, KBlockCompare>;
+
+  using ThirdLayer =
+      std::map<unsigned long long, states_ty, std::less<unsigned long long>>;
+  using SecondLayer =
+      std::map<std::vector<unsigned>, ThirdLayer, VectorsCompare>;
+  using FirstLayer = std::map<unsigned, SecondLayer, std::greater<unsigned>>;
+  using Data = std::map<KFunction *, FirstLayer, KFunctionCompare>;
+
+  Data data;
+  std::unordered_map<ExecutionState *, BlockLevel> stateToBlockLevel;
+  std::vector<unsigned> sizes;
+  RNG &theRNG;
+
+  void clear(ExecutionState &state);
+
+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