[CFGPrinter] Add node id formater (#164623)

This PR is part of the LLVM IR LSP server project
([RFC](https://discourse.llvm.org/t/rfc-ir-visualization-with-vs-code-extension-using-an-lsp-server/87773))

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.

Author: Bertik23

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.