1 files changed, 11 insertions, 47 deletions

diff --git a/llvm/lib/CodeGen/MachineBlockFrequencyInfo.cpp b/llvm/lib/CodeGen/MachineBlockFrequencyInfo.cpp
index e994aec4e1e..f02356b92c5 100644
--- a/llvm/lib/CodeGen/MachineBlockFrequencyInfo.cpp
+++ b/llvm/lib/CodeGen/MachineBlockFrequencyInfo.cpp
@@ -29,7 +29,6 @@ using namespace llvm;
 #define DEBUG_TYPE "block-freq"
 
 #ifndef NDEBUG
-enum GVDAGType { GVDT_None, GVDT_Fraction, GVDT_Integer, GVDT_Count };
 
 static cl::opt<GVDAGType> ViewMachineBlockFreqPropagationDAG(
     "view-machine-block-freq-propagation-dags", cl::Hidden,
@@ -79,59 +78,24 @@ template <> struct GraphTraits<MachineBlockFrequencyInfo *> {
   }
 };
 
+typedef BFIDOTGraphTraitsBase<MachineBlockFrequencyInfo,
+                              MachineBranchProbabilityInfo>
+    MBFIDOTGraphTraitsBase;
 template <>
 struct DOTGraphTraits<MachineBlockFrequencyInfo *>
-    : public DefaultDOTGraphTraits {
+    : public MBFIDOTGraphTraitsBase {
   explicit DOTGraphTraits(bool isSimple = false)
-      : DefaultDOTGraphTraits(isSimple) {}
-
-  typedef MachineBasicBlock::const_succ_iterator EdgeIter;
-  static std::string getGraphName(const MachineBlockFrequencyInfo *G) {
-    return G->getFunction()->getName();
-  }
+      : MBFIDOTGraphTraitsBase(isSimple) {}
 
   std::string getNodeLabel(const MachineBasicBlock *Node,
                            const MachineBlockFrequencyInfo *Graph) {
-    std::string Result;
-    raw_string_ostream OS(Result);
-
-    OS << Node->getName().str() << " : ";
-    switch (ViewMachineBlockFreqPropagationDAG) {
-    case GVDT_Fraction:
-      Graph->printBlockFreq(OS, Node);
-      break;
-    case GVDT_Integer:
-      OS << Graph->getBlockFreq(Node).getFrequency();
-      break;
-    case GVDT_Count: {
-      auto Count = Graph->getBlockProfileCount(Node);
-      if (Count)
-        OS << Count.getValue();
-      else
-        OS << "Unknown";
-      break;
-    }
-    case GVDT_None:
-      llvm_unreachable("If we are not supposed to render a graph we should "
-                       "never reach this point.");
-    }
-    return Result;
+    return MBFIDOTGraphTraitsBase::getNodeLabel(
+        Node, Graph, ViewMachineBlockFreqPropagationDAG);
   }
-  static std::string getEdgeAttributes(const MachineBasicBlock *Node,
-                                       EdgeIter EI,
-                                       const MachineBlockFrequencyInfo *MBFI) {
-    std::string Str;
-    const MachineBranchProbabilityInfo *MBPI = MBFI->getMBPI();
-    if (!MBPI)
-      return Str;
-    BranchProbability BP = MBPI->getEdgeProbability(Node, EI);
-    uint32_t N = BP.getNumerator();
-    uint32_t D = BP.getDenominator();
-    double Percent = 100.0 * N / D;
-    raw_string_ostream OS(Str);
-    OS << format("label=\"%.1f%%\"", Percent);
-    OS.flush();
-    return Str;
+
+  std::string getEdgeAttributes(const MachineBasicBlock *Node, EdgeIter EI,
+                                const MachineBlockFrequencyInfo *MBFI) {
+    return MBFIDOTGraphTraitsBase::getEdgeAttributes(Node, EI, MBFI->getMBPI());
   }
 };