[chore] Optimized getDistance

Columpio · misonijnik · commit febf79db0e7c · 2023-10-21T09:24:15.000+04:00
diff --git a/include/klee/Module/CodeGraphInfo.h b/include/klee/Module/CodeGraphInfo.h
@@ -16,31 +16,35 @@
 
 namespace klee {
 
+using Block = llvm::BasicBlock;
+using BlockDistanceMap = std::unordered_map<Block *, unsigned>;
+using FunctionDistanceMap = std::unordered_map<llvm::Function *, unsigned>;
+using SortedBlockDistances = std::vector<std::pair<Block *, unsigned>>;
+using SortedFunctionDistances =
+    std::vector<std::pair<llvm::Function *, unsigned>>;
+
 class CodeGraphInfo {
 
-  using blockDistanceMap =
-      std::unordered_map<KBlock *, std::unordered_map<KBlock *, unsigned>>;
-  using blockDistanceList =
-      std::unordered_map<KBlock *, std::vector<std::pair<KBlock *, unsigned>>>;
+  using blockToDistanceMap = std::unordered_map<Block *, BlockDistanceMap>;
+  using blockDistanceList = std::unordered_map<Block *, SortedBlockDistances>;
 
-  using functionDistanceMap =
-      std::unordered_map<KFunction *,
-                         std::unordered_map<KFunction *, unsigned>>;
+  using functionToDistanceMap =
+      std::unordered_map<llvm::Function *, FunctionDistanceMap>;
   using functionDistanceList =
-      std::unordered_map<KFunction *,
-                         std::vector<std::pair<KFunction *, unsigned>>>;
+      std::unordered_map<llvm::Function *, SortedFunctionDistances>;
 
   using functionBranchesSet =
       std::unordered_map<KFunction *, std::map<KBlock *, std::set<unsigned>>>;
 
 private:
-  blockDistanceMap blockDistance;
-  blockDistanceMap blockBackwardDistance;
+  blockToDistanceMap blockDistance;
+  blockToDistanceMap blockBackwardDistance;
+  std::set<Block *> blockCycles;
   blockDistanceList blockSortedDistance;
   blockDistanceList blockSortedBackwardDistance;
 
-  functionDistanceMap functionDistance;
-  functionDistanceMap functionBackwardDistance;
+  functionToDistanceMap functionDistance;
+  functionToDistanceMap functionBackwardDistance;
   functionDistanceList functionSortedDistance;
   functionDistanceList functionSortedBackwardDistance;
 
@@ -49,33 +53,24 @@ class CodeGraphInfo {
   functionBranchesSet functionBlocks;
 
 private:
-  void calculateDistance(KBlock *bb);
-  void calculateBackwardDistance(KBlock *bb);
+  void calculateDistance(Block *bb);
+  void calculateBackwardDistance(Block *bb);
 
-  void calculateDistance(KFunction *kf);
-  void calculateBackwardDistance(KFunction *kf);
+  void calculateDistance(KFunction *f);
+  void calculateBackwardDistance(KFunction *f);
 
   void calculateFunctionBranches(KFunction *kf);
   void calculateFunctionConditionalBranches(KFunction *kf);
   void calculateFunctionBlocks(KFunction *kf);
 
 public:
-  const std::unordered_map<KBlock *, unsigned int> &getDistance(KBlock *kb);
-  const std::unordered_map<KBlock *, unsigned int> &
-  getBackwardDistance(KBlock *kb);
-  const std::vector<std::pair<KBlock *, unsigned int>> &
-  getSortedDistance(KBlock *kb);
-  const std::vector<std::pair<KBlock *, unsigned int>> &
-  getSortedBackwardDistance(KBlock *kb);
-
-  const std::unordered_map<KFunction *, unsigned int> &
-  getDistance(KFunction *kf);
-  const std::unordered_map<KFunction *, unsigned int> &
-  getBackwardDistance(KFunction *kf);
-  const std::vector<std::pair<KFunction *, unsigned int>> &
-  getSortedDistance(KFunction *kf);
-  const std::vector<std::pair<KFunction *, unsigned int>> &
-  getSortedBackwardDistance(KFunction *kf);
+  const BlockDistanceMap &getDistance(Block *b);
+  const BlockDistanceMap &getDistance(KBlock *kb);
+  const BlockDistanceMap &getBackwardDistance(KBlock *kb);
+  bool hasCycle(KBlock *kb);
+
+  const FunctionDistanceMap &getDistance(KFunction *kf);
+  const FunctionDistanceMap &getBackwardDistance(KFunction *kf);
 
   void getNearestPredicateSatisfying(KBlock *from, KBlockPredicate predicate,
                                      std::set<KBlock *> &result);
diff --git a/lib/Core/DistanceCalculator.cpp b/lib/Core/DistanceCalculator.cpp
@@ -56,18 +56,23 @@ DistanceResult DistanceCalculator::getDistance(const ExecutionState &state,
 DistanceResult DistanceCalculator::getDistance(KBlock *kb, TargetKind kind,
                                                KBlock *target) {
   SpeculativeState specState(kb, kind);
-  if (distanceResultCache.count(target) == 0 ||
-      distanceResultCache.at(target).count(specState) == 0) {
+  auto it1 = distanceResultCache.find(target);
+  if (it1 == distanceResultCache.end()) {
+    SpeculativeStateToDistanceResultMap m;
+    it1 = distanceResultCache.emplace(target, std::move(m)).first;
+  }
+  auto &m = it1->second;
+  auto it2 = m.find(specState);
+  if (it2 == m.end()) {
     auto result = computeDistance(kb, kind, target);
-    distanceResultCache[target][specState] = result;
+    m.emplace(specState, result);
+    return result;
   }
-  return distanceResultCache.at(target).at(specState);
+  return it2->second;
 }
 
 DistanceResult DistanceCalculator::computeDistance(KBlock *kb, TargetKind kind,
                                                    KBlock *target) const {
-  const auto &distanceToTargetFunction =
-      codeGraphInfo.getBackwardDistance(target->parent);
   weight_type weight = 0;
   WeightResult res = Miss;
   bool isInsideFunction = true;
@@ -77,7 +82,7 @@ DistanceResult DistanceCalculator::computeDistance(KBlock *kb, TargetKind kind,
     break;
 
   case PreTarget:
-    res = tryGetPreTargetWeight(kb, weight, distanceToTargetFunction, target);
+    res = tryGetPreTargetWeight(kb, weight, target);
     isInsideFunction = false;
     break;
 
@@ -102,26 +107,21 @@ DistanceResult DistanceCalculator::getDistance(
   auto sfi = frames.rbegin(), sfe = frames.rend();
   bool strictlyAfterKB =
       sfi != sfe && sfi->kf->parent->inMainModule(*sfi->kf->function);
-  for (; sfi != sfe; sfi++) {
+  for (; sfi != sfe; sfi++, sfNum++) {
     unsigned callWeight;
     if (distanceInCallGraph(sfi->kf, kb, callWeight, distanceToTargetFunction,
                             target, strictlyAfterKB && sfNum != 0)) {
-      callWeight *= 2;
-      callWeight += sfNum;
+      callWeight = 2 * callWeight + sfNum;
 
       if (callWeight < UINT_MAX) {
         minCallWeight = callWeight;
         minSfNum = sfNum;
+        break;
       }
     }
 
-    if (sfi->caller) {
+    if (sfi->caller)
       kb = sfi->caller->parent;
-    }
-    sfNum++;
-
-    if (minCallWeight < UINT_MAX)
-      break;
   }
 
   TargetKind kind = NoneTarget;
@@ -138,71 +138,38 @@ DistanceResult DistanceCalculator::getDistance(
 
 bool DistanceCalculator::distanceInCallGraph(
     KFunction *kf, KBlock *origKB, unsigned int &distance,
-    const std::unordered_map<KFunction *, unsigned int>
-        &distanceToTargetFunction,
-    KBlock *target, bool strictlyAfterKB) const {
-  distance = UINT_MAX;
-  const std::unordered_map<KBlock *, unsigned> &dist =
-      codeGraphInfo.getDistance(origKB);
-  KBlock *targetBB = target;
-  KFunction *targetF = targetBB->parent;
-
-  if (kf == targetF && dist.count(targetBB) != 0) {
+    const FunctionDistanceMap &distanceToTargetFunction, KBlock *targetKB,
+    bool strictlyAfterKB) const {
+  auto &dist = codeGraphInfo.getDistance(origKB->basicBlock);
+  if (kf == targetKB->parent && dist.count(targetKB->basicBlock)) {
     distance = 0;
     return true;
   }
 
-  if (!strictlyAfterKB)
-    return distanceInCallGraph(kf, origKB, distance, distanceToTargetFunction,
-                               target);
-  auto min_distance = UINT_MAX;
   distance = UINT_MAX;
-  for (auto bb : successors(origKB->basicBlock)) {
-    auto kb = kf->blockMap[bb];
-    distanceInCallGraph(kf, kb, distance, distanceToTargetFunction, target);
-    if (distance < min_distance)
-      min_distance = distance;
-  }
-  distance = min_distance;
-  return distance != UINT_MAX;
-}
-
-bool DistanceCalculator::distanceInCallGraph(
-    KFunction *kf, KBlock *kb, unsigned int &distance,
-    const std::unordered_map<KFunction *, unsigned int>
-        &distanceToTargetFunction,
-    KBlock *target) const {
-  distance = UINT_MAX;
-  const std::unordered_map<KBlock *, unsigned> &dist =
-      codeGraphInfo.getDistance(kb);
-
-  for (auto &kCallBlock : kf->kCallBlocks) {
-    if (dist.count(kCallBlock) == 0)
+  bool cannotReachItself = strictlyAfterKB && !codeGraphInfo.hasCycle(origKB);
+  for (auto kCallBlock : kf->kCallBlocks) {
+    if (!dist.count(kCallBlock->basicBlock) ||
+        (cannotReachItself && origKB == kCallBlock))
       continue;
-    for (auto &calledFunction : kCallBlock->calledFunctions) {
-      KFunction *calledKFunction = kf->parent->functionMap[calledFunction];
-      if (distanceToTargetFunction.count(calledKFunction) != 0 &&
-          distance > distanceToTargetFunction.at(calledKFunction) + 1) {
-        distance = distanceToTargetFunction.at(calledKFunction) + 1;
-      }
+    for (auto calledFunction : kCallBlock->calledFunctions) {
+      auto it = distanceToTargetFunction.find(calledFunction);
+      if (it != distanceToTargetFunction.end() && distance > it->second + 1)
+        distance = it->second + 1;
     }
   }
   return distance != UINT_MAX;
 }
 
-WeightResult
-DistanceCalculator::tryGetLocalWeight(KBlock *kb, weight_type &weight,
-                                      const std::vector<KBlock *> &localTargets,
-                                      KBlock *target) const {
-  KBlock *currentKB = kb;
-  const std::unordered_map<KBlock *, unsigned> &dist =
-      codeGraphInfo.getDistance(currentKB);
+WeightResult DistanceCalculator::tryGetLocalWeight(
+    KBlock *kb, weight_type &weight,
+    const std::vector<KBlock *> &localTargets) const {
+  const auto &dist = codeGraphInfo.getDistance(kb);
   weight = UINT_MAX;
-  for (auto &end : localTargets) {
-    if (dist.count(end) > 0) {
-      unsigned int w = dist.at(end);
-      weight = std::min(w, weight);
-    }
+  for (auto end : localTargets) {
+    auto it = dist.find(end->basicBlock);
+    if (it != dist.end())
+      weight = std::min(it->second, weight);
   }
 
   if (weight == UINT_MAX)
@@ -214,27 +181,26 @@ DistanceCalculator::tryGetLocalWeight(KBlock *kb, weight_type &weight,
   return Continue;
 }
 
-WeightResult DistanceCalculator::tryGetPreTargetWeight(
-    KBlock *kb, weight_type &weight,
-    const std::unordered_map<KFunction *, unsigned int>
-        &distanceToTargetFunction,
-    KBlock *target) const {
+WeightResult DistanceCalculator::tryGetPreTargetWeight(KBlock *kb,
+                                                       weight_type &weight,
+                                                       KBlock *target) const {
+  auto &distanceToTargetFunction =
+      codeGraphInfo.getBackwardDistance(target->parent);
   KFunction *currentKF = kb->parent;
   std::vector<KBlock *> localTargets;
-  for (auto &kCallBlock : currentKF->kCallBlocks) {
-    for (auto &calledFunction : kCallBlock->calledFunctions) {
-      KFunction *calledKFunction =
-          currentKF->parent->functionMap[calledFunction];
-      if (distanceToTargetFunction.count(calledKFunction) > 0) {
+  for (auto kCallBlock : currentKF->kCallBlocks) {
+    for (auto calledFunction : kCallBlock->calledFunctions) {
+      if (distanceToTargetFunction.count(calledFunction)) {
         localTargets.push_back(kCallBlock);
+        break;
       }
     }
   }
 
   if (localTargets.empty())
     return Miss;
 
-  WeightResult res = tryGetLocalWeight(kb, weight, localTargets, target);
+  WeightResult res = tryGetLocalWeight(kb, weight, localTargets);
   return res == Done ? Continue : res;
 }
 
@@ -247,14 +213,14 @@ WeightResult DistanceCalculator::tryGetPostTargetWeight(KBlock *kb,
   if (localTargets.empty())
     return Miss;
 
-  WeightResult res = tryGetLocalWeight(kb, weight, localTargets, target);
+  WeightResult res = tryGetLocalWeight(kb, weight, localTargets);
   return res == Done ? Continue : res;
 }
 
 WeightResult DistanceCalculator::tryGetTargetWeight(KBlock *kb,
                                                     weight_type &weight,
                                                     KBlock *target) const {
   std::vector<KBlock *> localTargets = {target};
-  WeightResult res = tryGetLocalWeight(kb, weight, localTargets, target);
+  WeightResult res = tryGetLocalWeight(kb, weight, localTargets);
   return res;
 }
diff --git a/lib/Core/DistanceCalculator.h b/lib/Core/DistanceCalculator.h
@@ -11,6 +11,7 @@
 #define KLEE_DISTANCE_CALCULATOR_H
 
 #include "ExecutionState.h"
+#include "klee/Module/CodeGraphInfo.h"
 
 namespace llvm {
 class BasicBlock;
@@ -107,22 +108,14 @@ class DistanceCalculator {
                                  KBlock *target) const;
 
   bool distanceInCallGraph(KFunction *kf, KBlock *kb, unsigned int &distance,
-                           const std::unordered_map<KFunction *, unsigned int>
-                               &distanceToTargetFunction,
-                           KBlock *target) const;
-  bool distanceInCallGraph(KFunction *kf, KBlock *kb, unsigned int &distance,
-                           const std::unordered_map<KFunction *, unsigned int>
-                               &distanceToTargetFunction,
+                           const FunctionDistanceMap &distanceToTargetFunction,
                            KBlock *target, bool strictlyAfterKB) const;
 
-  WeightResult tryGetLocalWeight(KBlock *kb, weight_type &weight,
-                                 const std::vector<KBlock *> &localTargets,
-                                 KBlock *target) const;
   WeightResult
-  tryGetPreTargetWeight(KBlock *kb, weight_type &weight,
-                        const std::unordered_map<KFunction *, unsigned int>
-                            &distanceToTargetFunction,
-                        KBlock *target) const;
+  tryGetLocalWeight(KBlock *kb, weight_type &weight,
+                    const std::vector<KBlock *> &localTargets) const;
+  WeightResult tryGetPreTargetWeight(KBlock *kb, weight_type &weight,
+                                     KBlock *target) const;
   WeightResult tryGetTargetWeight(KBlock *kb, weight_type &weight,
                                   KBlock *target) const;
   WeightResult tryGetPostTargetWeight(KBlock *kb, weight_type &weight,
diff --git a/lib/Core/TargetedExecutionManager.cpp b/lib/Core/TargetedExecutionManager.cpp
@@ -379,18 +379,18 @@ bool TargetedExecutionManager::canReach(const ref<Location> &from,
         }
 
         const auto &blockDist = codeGraphInfo.getDistance(fromBlock);
-        if (blockDist.count(toBlock) != 0) {
+        if (blockDist.count(toBlock->basicBlock) != 0) {
           return true;
         }
       } else {
         const auto &funcDist = codeGraphInfo.getDistance(fromKf);
-        if (funcDist.count(toKf) != 0) {
+        if (funcDist.count(toKf->function) != 0) {
           return true;
         }
 
         const auto &backwardFuncDist =
             codeGraphInfo.getBackwardDistance(fromKf);
-        if (backwardFuncDist.count(toKf) != 0) {
+        if (backwardFuncDist.count(toKf->function) != 0) {
           return true;
         }
       }
@@ -464,9 +464,9 @@ KFunction *TargetedExecutionManager::tryResolveEntryFunction(
           KFunction *curKf = nullptr;
           for (size_t m = 0; m < currKFs.size() && !curKf; ++m) {
             curKf = currKFs.at(m);
-            if (funcDist.count(curKf) == 0) {
+            if (funcDist.count(curKf->function) == 0) {
               const auto &curFuncDist = codeGraphInfo.getDistance(curKf);
-              if (curFuncDist.count(resKf) == 0) {
+              if (curFuncDist.count(resKf->function) == 0) {
                 curKf = nullptr;
               } else {
                 i = j;
diff --git a/lib/Module/CodeGraphInfo.cpp b/lib/Module/CodeGraphInfo.cpp
diff --git a/lib/Module/KModule.cpp b/lib/Module/KModule.cpp
