[CFGPrinter] Add node id formater

llvm/include/llvm/Analysis/CFGPrinter.h

@@ -31,6 +31,9 @@
 #include "llvm/Support/DOTGraphTraits.h"
 #include "llvm/Support/FormatVariadic.h"
 
+#include <functional>
+#include <sstream>
+
 namespace llvm {
 class ModuleSlotTracker;
 
@@ -69,13 +72,18 @@ class DOTFuncInfo {
   bool ShowHeat;
   bool EdgeWeights;
   bool RawWeights;
+  using NodeIdFormatterTy =
+      std::function<std::optional<std::string>(const BasicBlock *)>;
+  std::optional<NodeIdFormatterTy> NodeIdFormatter;
 
 public:
   DOTFuncInfo(const Function *F) : DOTFuncInfo(F, nullptr, nullptr, 0) {}
   LLVM_ABI ~DOTFuncInfo();
 
-  LLVM_ABI DOTFuncInfo(const Function *F, const BlockFrequencyInfo *BFI,
-                       const BranchProbabilityInfo *BPI, uint64_t MaxFreq);
+  LLVM_ABI
+  DOTFuncInfo(const Function *F, const BlockFrequencyInfo *BFI,
+              const BranchProbabilityInfo *BPI, uint64_t MaxFreq,
+              std::optional<NodeIdFormatterTy> NodeIdFormatter = std::nullopt);
 
   const BlockFrequencyInfo *getBFI() const { return BFI; }
 
@@ -102,6 +110,10 @@ class DOTFuncInfo {
   void setEdgeWeights(bool EdgeWeights) { this->EdgeWeights = EdgeWeights; }
 
   bool showEdgeWeights() { return EdgeWeights; }
+
+  std::optional<NodeIdFormatterTy> getNodeIdFormatter() {
+    return NodeIdFormatter;
+  }
 };
 
 template <>
@@ -311,21 +323,27 @@ struct DOTGraphTraits<DOTFuncInfo *> : public DefaultDOTGraphTraits {
   }
 
   std::string getNodeAttributes(const BasicBlock *Node, DOTFuncInfo *CFGInfo) {
+    std::stringstream Attrs;
+
+    if (auto NodeIdFmt = CFGInfo->getNodeIdFormatter())
+      if (auto NodeId = (*NodeIdFmt)(Node))
+        Attrs << "id=\"" << *NodeId << "\"";
+
+    if (CFGInfo->showHeatColors()) {
+      uint64_t Freq = CFGInfo->getFreq(Node);
+      std::string Color = getHeatColor(Freq, CFGInfo->getMaxFreq());
+      std::string EdgeColor = (Freq <= (CFGInfo->getMaxFreq() / 2))
+                                  ? (getHeatColor(0))
+                                  : (getHeatColor(1));
+      if (!Attrs.str().empty())
+        Attrs << ",";
+      Attrs << "color=\"" << EdgeColor << "ff\", style=filled, "
+            << "fillcolor=\"" << Color << "70\", " << "fontname=\"Courier\"";
+    }
 
-    if (!CFGInfo->showHeatColors())
-      return "";
-
-    uint64_t Freq = CFGInfo->getFreq(Node);
-    std::string Color = getHeatColor(Freq, CFGInfo->getMaxFreq());
-    std::string EdgeColor = (Freq <= (CFGInfo->getMaxFreq() / 2))
-                                ? (getHeatColor(0))
-                                : (getHeatColor(1));
-
-    std::string Attrs = "color=\"" + EdgeColor + "ff\", style=filled," +
-                        " fillcolor=\"" + Color + "70\"" +
-                        " fontname=\"Courier\"";
-    return Attrs;
+    return Attrs.str();
   }
+
   LLVM_ABI bool isNodeHidden(const BasicBlock *Node,
                              const DOTFuncInfo *CFGInfo);
   LLVM_ABI void computeDeoptOrUnreachablePaths(const Function *F);

llvm/lib/Analysis/CFGPrinter.cpp

@@ -92,8 +92,10 @@ static void viewCFG(Function &F, const BlockFrequencyInfo *BFI,
 }
 
 DOTFuncInfo::DOTFuncInfo(const Function *F, const BlockFrequencyInfo *BFI,
-                         const BranchProbabilityInfo *BPI, uint64_t MaxFreq)
-    : F(F), BFI(BFI), BPI(BPI), MaxFreq(MaxFreq) {
+                         const BranchProbabilityInfo *BPI, uint64_t MaxFreq,
+                         std::optional<NodeIdFormatterTy> NodeIdFormatter)
+    : F(F), BFI(BFI), BPI(BPI), MaxFreq(MaxFreq),
+      NodeIdFormatter(NodeIdFormatter) {
   ShowHeat = false;
   EdgeWeights = !!BPI; // Print EdgeWeights when BPI is available.
   RawWeights = !!BFI;  // Print RawWeights when BFI is available.