[CFGPrinter] Add node id formater

[Bertik23]

This PR is part of the LLVM IR LSP server project (RFC)

Sometimes it is nice to be able to specify IDs of nodes in the printed CFG. For better manipulation of the outputed CFG.
In our case we will use it for navigation between IR and CFG views.

This adds an argument to DOTFuncInfo - a function that takes a BasicBlock and returns a node ID, to be printed in the result dot.

[llvmbot]

@llvm/pr-subscribers-llvm-analysis

Author: Bertik23 (Bertik23)

Changes

This PR is part of the LLVM IR LSP server project (RFC)

Sometimes it is nice to be able to specify IDs of nodes in the printed CFG. For better manipulation of the outputed CFG.
In our case we will use it for navigation between IR and CFG views.

This adds an argument to DOTFuncInfo - a function that takes a BasicBlock and returns a node ID, to be printed in the result dot.

---

Full diff: https://github.com/llvm/llvm-project/pull/164623.diff

2 Files Affected:

• (modified) llvm/include/llvm/Analysis/CFGPrinter.h (+33-15)
• (modified) llvm/lib/Analysis/CFGPrinter.cpp (+4-2)

diff --git a/llvm/include/llvm/Analysis/CFGPrinter.h b/llvm/include/llvm/Analysis/CFGPrinter.h
index ec26da87eb916..aa711642a3a6d 100644
--- a/llvm/include/llvm/Analysis/CFGPrinter.h
+++ b/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);
diff --git a/llvm/lib/Analysis/CFGPrinter.cpp b/llvm/lib/Analysis/CFGPrinter.cpp
index 38aad849755be..39108a906f081 100644
--- a/llvm/lib/Analysis/CFGPrinter.cpp
+++ b/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.