Add ability to annotate (color) nodes in a viewGraph.

llvm-svn: 30686

5 files changed, 94 insertions, 5 deletions

diff --git a/llvm/include/llvm/CodeGen/SelectionDAG.h b/llvm/include/llvm/CodeGen/SelectionDAG.h
index 0475147b915..c173eb122c3 100644
--- a/llvm/include/llvm/CodeGen/SelectionDAG.h
+++ b/llvm/include/llvm/CodeGen/SelectionDAG.h
@@ -69,10 +69,29 @@ public:
   TargetLowering &getTargetLoweringInfo() const { return TLI; }
   MachineDebugInfo *getMachineDebugInfo() const { return DI; }
 
-  /// viewGraph - Pop up a ghostview window with the DAG rendered using 'dot'.
+  /// viewGraph - Pop up a GraphViz/gv window with the DAG rendered using 'dot'.
   ///
   void viewGraph();
+  
+#ifndef NDEBUG
+  std::map<const SDNode *, std::string> NodeGraphAttrs;
+#endif
 
+  /// clearGraphAttrs - Clear all previously defined node graph attributes.
+  /// Intended to be used from a debugging tool (eg. gdb).
+  void clearGraphAttrs();
+  
+  /// setGraphAttrs - Set graph attributes for a node. (eg. "color=red".)
+  ///
+  void setGraphAttrs(const SDNode *N, const char *Attrs);
+  
+  /// getGraphAttrs - Get graph attributes for a node. (eg. "color=red".)
+  /// Used from getNodeAttributes.
+  const std::string getGraphAttrs(const SDNode *N) const;
+  
+  /// setGraphColor - Convenience for setting node color attribute.
+  ///
+  void setGraphColor(const SDNode *N, const char *Color);
 
   typedef ilist<SDNode>::const_iterator allnodes_const_iterator;
   allnodes_const_iterator allnodes_begin() const { return AllNodes.begin(); }
diff --git a/llvm/include/llvm/Support/DOTGraphTraits.h b/llvm/include/llvm/Support/DOTGraphTraits.h
index 9300ea7a8af..62eac0cea90 100644
--- a/llvm/include/llvm/Support/DOTGraphTraits.h
+++ b/llvm/include/llvm/Support/DOTGraphTraits.h
@@ -60,7 +60,9 @@ struct DefaultDOTGraphTraits {
 
   /// If you want to specify custom node attributes, this is the place to do so
   ///
-  static std::string getNodeAttributes(const void *Node) { return ""; }
+  static std::string getNodeAttributes(const void *Node, const void *Graph) {
+    return "";
+  }
 
   /// If you want to override the dot attributes printed for a particular edge,
   /// override this method.
diff --git a/llvm/include/llvm/Support/GraphWriter.h b/llvm/include/llvm/Support/GraphWriter.h
index 6bf43d29c7c..b4e6d845e61 100644
--- a/llvm/include/llvm/Support/GraphWriter.h
+++ b/llvm/include/llvm/Support/GraphWriter.h
@@ -109,7 +109,7 @@ public:
   }
 
   void writeNode(NodeType *Node) {
-    std::string NodeAttributes = DOTTraits::getNodeAttributes(Node);
+    std::string NodeAttributes = DOTTraits::getNodeAttributes(Node, G);
 
     O << "\tNode" << reinterpret_cast<const void*>(Node) << " [shape=record,";
     if (!NodeAttributes.empty()) O << NodeAttributes << ",";
diff --git a/llvm/lib/Analysis/DataStructure/Printer.cpp b/llvm/lib/Analysis/DataStructure/Printer.cpp
index 7cb79247e4b..ce999615811 100644
--- a/llvm/lib/Analysis/DataStructure/Printer.cpp
+++ b/llvm/lib/Analysis/DataStructure/Printer.cpp
@@ -112,7 +112,7 @@ struct DOTGraphTraits<const DSGraph*> : public DefaultDOTGraphTraits {
     return getCaption(Node, Graph);
   }
 
-  static std::string getNodeAttributes(const DSNode *N) {
+  static std::string getNodeAttributes(const DSNode *N, const DSGraph *Graph) {
     return "shape=Mrecord";
   }
 
diff --git a/llvm/lib/CodeGen/SelectionDAG/SelectionDAGPrinter.cpp b/llvm/lib/CodeGen/SelectionDAG/SelectionDAGPrinter.cpp
index 1669c08a402..430c326ba56 100644
--- a/llvm/lib/CodeGen/SelectionDAG/SelectionDAGPrinter.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/SelectionDAGPrinter.cpp
@@ -27,6 +27,9 @@
 using namespace llvm;
 
 namespace llvm {
+#ifndef NDEBUG
+  std::map<const SDNode *, std::string> DagNodeColor;
+#endif
   template<>
   struct DOTGraphTraits<SelectionDAG*> : public DefaultDOTGraphTraits {
     static std::string getGraphName(const SelectionDAG *G) {
@@ -44,7 +47,17 @@ namespace llvm {
 
     static std::string getNodeLabel(const SDNode *Node,
                                     const SelectionDAG *Graph);
-    static std::string getNodeAttributes(const SDNode *N) {
+    static std::string getNodeAttributes(const SDNode *N,
+                                         const SelectionDAG *Graph) {
+#ifndef NDEBUG
+      const std::string &Attrs = Graph->getGraphAttrs(N);
+      if (!Attrs.empty()) {
+        if (Attrs.find("shape=") == std::string::npos)
+          return std::string("shape=Mrecord,") + Attrs;
+        else
+          return Attrs;
+      }
+#endif
       return "shape=Mrecord";
     }
 
@@ -138,3 +151,58 @@ void SelectionDAG::viewGraph() {
             << "systems with Graphviz or gv!\n";
 #endif  // NDEBUG
 }
+
+
+/// clearGraphAttrs - Clear all previously defined node graph attributes.
+/// Intended to be used from a debugging tool (eg. gdb).
+void SelectionDAG::clearGraphAttrs() {
+#ifndef NDEBUG
+  NodeGraphAttrs.clear();
+#else
+  std::cerr << "SelectionDAG::clearGraphAttrs is only available in debug builds"
+            << " on systems with Graphviz or gv!\n";
+#endif
+}
+
+
+/// setGraphAttrs - Set graph attributes for a node. (eg. "color=red".)
+///
+void SelectionDAG::setGraphAttrs(const SDNode *N, const char *Attrs) {
+#ifndef NDEBUG
+  NodeGraphAttrs[N] = Attrs;
+#else
+  std::cerr << "SelectionDAG::setGraphAttrs is only available in debug builds"
+            << " on systems with Graphviz or gv!\n";
+#endif
+}
+
+
+/// getGraphAttrs - Get graph attributes for a node. (eg. "color=red".)
+/// Used from getNodeAttributes.
+const std::string SelectionDAG::getGraphAttrs(const SDNode *N) const {
+#ifndef NDEBUG
+  std::map<const SDNode *, std::string>::const_iterator I =
+    NodeGraphAttrs.find(N);
+    
+  if (I != NodeGraphAttrs.end())
+    return I->second;
+  else
+    return "";
+#else
+  std::cerr << "SelectionDAG::getGraphAttrs is only available in debug builds"
+            << " on systems with Graphviz or gv!\n";
+  return std::string("");
+#endif
+}
+
+/// setGraphColor - Convenience for setting node color attribute.
+///
+void SelectionDAG::setGraphColor(const SDNode *N, const char *Color) {
+#ifndef NDEBUG
+  NodeGraphAttrs[N] = std::string("color=") + Color;
+#else
+  std::cerr << "SelectionDAG::setGraphColor is only available in debug builds"
+            << " on systems with Graphviz or gv!\n";
+#endif
+}
+