lib/Transforms/Instrumentation/MaximumSpanningTree.h

   1 //===- llvm/Analysis/MaximumSpanningTree.h - Interface ----------*- C++ -*-===//
   2 //
   3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
   4 // See https://llvm.org/LICENSE.txt for license information.
   5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
   6 //
   7 //===----------------------------------------------------------------------===//
   8 //
   9 // This module provides means for calculating a maximum spanning tree for a
  10 // given set of weighted edges. The type parameter T is the type of a node.
  11 //
  12 //===----------------------------------------------------------------------===//
  13
  14 #ifndef LLVM_LIB_TRANSFORMS_INSTRUMENTATION_MAXIMUMSPANNINGTREE_H
  15 #define LLVM_LIB_TRANSFORMS_INSTRUMENTATION_MAXIMUMSPANNINGTREE_H
  16
  17 #include "llvm/ADT/EquivalenceClasses.h"
  18 #include "llvm/IR/BasicBlock.h"
  19 #include <algorithm>
  20 #include <vector>
  21
  22 namespace llvm {
  23
  24   /// MaximumSpanningTree - A MST implementation.
  25   /// The type parameter T determines the type of the nodes of the graph.
  26   template <typename T>
  27   class MaximumSpanningTree {
  28   public:
  29     typedef std::pair<const T*, const T*> Edge;
  30     typedef std::pair<Edge, double> EdgeWeight;
  31     typedef std::vector<EdgeWeight> EdgeWeights;
  32   protected:
  33     typedef std::vector<Edge> MaxSpanTree;
  34
  35     MaxSpanTree MST;
  36
  37   private:
  38     // A comparing class for comparing weighted edges.
  39     struct EdgeWeightCompare {
  40       static bool getBlockSize(const T *X) {
  41         const BasicBlock *BB = dyn_cast_or_null<BasicBlock>(X);
  42         return BB ? BB->size() : 0;
  43       }
  44
  45       bool operator()(EdgeWeight X, EdgeWeight Y) const {
  46         if (X.second > Y.second) return true;
  47         if (X.second < Y.second) return false;
  48
  49         // Equal edge weights: break ties by comparing block sizes.
  50         size_t XSizeA = getBlockSize(X.first.first);
  51         size_t YSizeA = getBlockSize(Y.first.first);
  52         if (XSizeA > YSizeA) return true;
  53         if (XSizeA < YSizeA) return false;
  54
  55         size_t XSizeB = getBlockSize(X.first.second);
  56         size_t YSizeB = getBlockSize(Y.first.second);
  57         if (XSizeB > YSizeB) return true;
  58         if (XSizeB < YSizeB) return false;
  59
  60         return false;
  61       }
  62     };
  63
  64   public:
  65     static char ID; // Class identification, replacement for typeinfo
  66
  67     /// MaximumSpanningTree() - Takes a vector of weighted edges and returns a
  68     /// spanning tree.
  69     MaximumSpanningTree(EdgeWeights &EdgeVector) {
  70
  71       std::stable_sort(EdgeVector.begin(), EdgeVector.end(), EdgeWeightCompare());
  72
  73       // Create spanning tree, Forest contains a special data structure
  74       // that makes checking if two nodes are already in a common (sub-)tree
  75       // fast and cheap.
  76       EquivalenceClasses<const T*> Forest;
  77       for (typename EdgeWeights::iterator EWi = EdgeVector.begin(),
  78            EWe = EdgeVector.end(); EWi != EWe; ++EWi) {
  79         Edge e = (*EWi).first;
  80
  81         Forest.insert(e.first);
  82         Forest.insert(e.second);
  83       }
  84
  85       // Iterate over the sorted edges, biggest first.
  86       for (typename EdgeWeights::iterator EWi = EdgeVector.begin(),
  87            EWe = EdgeVector.end(); EWi != EWe; ++EWi) {
  88         Edge e = (*EWi).first;
  89
  90         if (Forest.findLeader(e.first) != Forest.findLeader(e.second)) {
  91           Forest.unionSets(e.first, e.second);
  92           // So we know now that the edge is not already in a subtree, so we push
  93           // the edge to the MST.
  94           MST.push_back(e);
  95         }
  96       }
  97     }
  98
  99     typename MaxSpanTree::iterator begin() {
 100       return MST.begin();
 101     }
 102
 103     typename MaxSpanTree::iterator end() {
 104       return MST.end();
 105     }
 106   };
 107
 108 } // End llvm namespace
 109
 110 #endif // LLVM_LIB_TRANSFORMS_INSTRUMENTATION_MAXIMUMSPANNINGTREE_H