From edd39c68fb291f8bb4a811a9ba2a60338adf479e Mon Sep 17 00:00:00 2001
From: Nguyễn Gia Phong <mcsinyx@disroot.org>
Date: Thu, 22 Jun 2023 16:05:59 +0900
Subject: Add search heuristic for patch locations

---
 lib/Core/ExecutionState.cpp | 2 ++
 lib/Core/ExecutionState.h   | 3 +++
 lib/Core/Searcher.cpp       | 3 +++
 lib/Core/Searcher.h         | 2 ++
 lib/Core/UserSearcher.cpp   | 3 +++
 5 files changed, 13 insertions(+)

diff --git a/lib/Core/ExecutionState.cpp b/lib/Core/ExecutionState.cpp
index b614f6e4..9a39d2dc 100644
--- a/lib/Core/ExecutionState.cpp
+++ b/lib/Core/ExecutionState.cpp
@@ -109,6 +109,7 @@ ExecutionState::ExecutionState(const ExecutionState& state):
     cexPreferences(state.cexPreferences),
     arrayNames(state.arrayNames),
     openMergeStack(state.openMergeStack),
+    patchLocs(state.patchLocs),
     patchNo(state.patchNo),
     metaEnvVar(state.metaEnvVar),
     formula(state.formula),
@@ -423,6 +424,7 @@ std::string extractMetaEnvVar(ref<Expr> e) {
 void ExecutionState::addConstraint(ref<Expr> e) {
   auto v = extractMetaEnvVar(e);
   if (!v.empty()) {
+    this->patchLocs++;
     if (v.substr(v.size() - 2) != "=0")
       this->metaEnvVar = v;
     return;
diff --git a/lib/Core/ExecutionState.h b/lib/Core/ExecutionState.h
index 93ed8214..91a5a9f9 100644
--- a/lib/Core/ExecutionState.h
+++ b/lib/Core/ExecutionState.h
@@ -226,6 +226,9 @@ public:
   /// @brief The objects handling the klee_open_merge calls this state ran through
   std::vector<ref<MergeHandler>> openMergeStack;
 
+  /// @ brief The number of patch locations reached
+  std::uint64_t patchLocs = 0;
+
   /// @ brief The patch number, starting from 1; 0 being the original.
   std::uint64_t patchNo = 0;
 
diff --git a/lib/Core/Searcher.cpp b/lib/Core/Searcher.cpp
index e94511ea..d8b76ef2 100644
--- a/lib/Core/Searcher.cpp
+++ b/lib/Core/Searcher.cpp
@@ -159,6 +159,7 @@ WeightedRandomSearcher::WeightedRandomSearcher(WeightType type, RNG &rng)
   switch(type) {
   case Depth:
   case RP:
+  case PatchLocCount:
     updateWeights = false;
     break;
   case InstCount:
@@ -200,6 +201,8 @@ double WeightedRandomSearcher::getWeight(ExecutionState *es) {
       return (es->queryMetaData.queryCost.toSeconds() < .1)
                  ? 1.
                  : 1. / es->queryMetaData.queryCost.toSeconds();
+    case PatchLocCount:
+      return es->patchLocs;
     case CoveringNew:
     case MinDistToUncovered: {
       uint64_t md2u = computeMinDistToUncovered(es->pc,
diff --git a/lib/Core/Searcher.h b/lib/Core/Searcher.h
index 36efe67a..07fee8ca 100644
--- a/lib/Core/Searcher.h
+++ b/lib/Core/Searcher.h
@@ -67,6 +67,7 @@ namespace klee {
       RandomPath,
       NURS_CovNew,
       NURS_MD2U,
+      NURS_Patch,
       NURS_Depth,
       NURS_RP,
       NURS_ICnt,
@@ -131,6 +132,7 @@ namespace klee {
       QueryCost,
       InstCount,
       CPInstCount,
+      PatchLocCount,
       MinDistToUncovered,
       CoveringNew
     };
diff --git a/lib/Core/UserSearcher.cpp b/lib/Core/UserSearcher.cpp
index 7e0fd31f..ca7eac0d 100644
--- a/lib/Core/UserSearcher.cpp
+++ b/lib/Core/UserSearcher.cpp
@@ -41,6 +41,8 @@ cl::list<Searcher::CoreSearchType> CoreSearch(
                    "use Non Uniform Random Search (NURS) with Coverage-New"),
         clEnumValN(Searcher::NURS_MD2U, "nurs:md2u",
                    "use NURS with Min-Dist-to-Uncovered"),
+        clEnumValN(Searcher::NURS_Patch, "nurs:patch",
+                   "use NURS with patch locations count"),
         clEnumValN(Searcher::NURS_Depth, "nurs:depth", "use NURS with depth"),
         clEnumValN(Searcher::NURS_RP, "nurs:rp", "use NURS with 1/2^depth"),
         clEnumValN(Searcher::NURS_ICnt, "nurs:icnt",
@@ -116,6 +118,7 @@ Searcher *getNewSearcher(Searcher::CoreSearchType type, RNG &rng,
     case Searcher::RandomPath: searcher = new RandomPathSearcher(executionTree, rng); break;
     case Searcher::NURS_CovNew: searcher = new WeightedRandomSearcher(WeightedRandomSearcher::CoveringNew, rng); break;
     case Searcher::NURS_MD2U: searcher = new WeightedRandomSearcher(WeightedRandomSearcher::MinDistToUncovered, rng); break;
+    case Searcher::NURS_Patch: searcher = new WeightedRandomSearcher(WeightedRandomSearcher::PatchLocCount, rng); break;
     case Searcher::NURS_Depth: searcher = new WeightedRandomSearcher(WeightedRandomSearcher::Depth, rng); break;
     case Searcher::NURS_RP: searcher = new WeightedRandomSearcher(WeightedRandomSearcher::RP, rng); break;
     case Searcher::NURS_ICnt: searcher = new WeightedRandomSearcher(WeightedRandomSearcher::InstCount, rng); break;
-- 
cgit 1.4.1