Find the shortest distance between source and target basic blocks in Binary Ninja

binaryninja-dijkstra.cpp

namespace BN = BinaryNinja;

using std::map;
using std::pair;
using std::priority_queue;
using std::vector;

int32_t dijkstra(vector<BN::Ref<BN::BasicBlock>> blocks, BN::Ref<BN::BasicBlock> source, BN::Ref<BN::BasicBlock> target) {
  using path = pair<int32_t, uint64_t>;
  
  map<uint64_t, vector<uint64_t>> adj;
  priority_queue<int32_t, vector<path>, std::greater<path>> pq;
  map<uint64_t, int32_t> dist;
  
  for (auto&& block : blocks) {
    auto&& edges = block->GetOutgoingEdges();
    
    vector<uint64_t> vec;
    vec.reserve(edges.size());
    for (auto&& edge : edges) {
      vec.emplace_back(edge.target);
    }
    
    adj[block->GetStart()] = vec;
    dist[block->GetStart()] = std::numeric_limits<int32_t>::max();
  }
  
  dist[source->GetStart()] = 0;
  pq.push(make_pair(0, source->GetStart()));
  
  while (!pq.empty()) {
    uint64_t u = pq.top().second;
    pq.pop();
    
    if (u == target->GetStart())
      break;
    
    for (int64_t v : adj.find(u)->second) {
      if (dist[v] > dist[u] + 1) {
        dist[v] = dist[u] + 1;
        pq.push(make_pair(dist[v], v));
      }
    }
  }
  
  return dist[target->GetStart()];
}