1 files changed, 69 insertions, 0 deletions
diff --git a/compiler/rustc_middle/src/mir/generic_graph.rs b/compiler/rustc_middle/src/mir/generic_graph.rs
new file mode 100644
index 000000000..f3621cd99
--- /dev/null
+++ b/compiler/rustc_middle/src/mir/generic_graph.rs
@@ -0,0 +1,69 @@
+use gsgdt::{Edge, Graph, Node, NodeStyle};
+use rustc_hir::def_id::DefId;
+use rustc_middle::mir::*;
+use rustc_middle::ty::TyCtxt;
+
+/// Convert an MIR function into a gsgdt Graph
+pub fn mir_fn_to_generic_graph<'tcx>(tcx: TyCtxt<'tcx>, body: &Body<'_>) -> Graph {
+    let def_id = body.source.def_id();
+    let def_name = graphviz_safe_def_name(def_id);
+    let graph_name = format!("Mir_{}", def_name);
+    let dark_mode = tcx.sess.opts.unstable_opts.graphviz_dark_mode;
+
+    // Nodes
+    let nodes: Vec<Node> = body
+        .basic_blocks()
+        .iter_enumerated()
+        .map(|(block, _)| bb_to_graph_node(block, body, dark_mode))
+        .collect();
+
+    // Edges
+    let mut edges = Vec::new();
+    for (source, _) in body.basic_blocks().iter_enumerated() {
+        let def_id = body.source.def_id();
+        let terminator = body[source].terminator();
+        let labels = terminator.kind.fmt_successor_labels();
+
+        for (target, label) in terminator.successors().zip(labels) {
+            let src = node(def_id, source);
+            let trg = node(def_id, target);
+            edges.push(Edge::new(src, trg, label.to_string()));
+        }
+    }
+
+    Graph::new(graph_name, nodes, edges)
+}
+
+fn bb_to_graph_node(block: BasicBlock, body: &Body<'_>, dark_mode: bool) -> Node {
+    let def_id = body.source.def_id();
+    let data = &body[block];
+    let label = node(def_id, block);
+
+    let (title, bgcolor) = if data.is_cleanup {
+        let color = if dark_mode { "royalblue" } else { "lightblue" };
+        (format!("{} (cleanup)", block.index()), color)
+    } else {
+        let color = if dark_mode { "dimgray" } else { "gray" };
+        (format!("{}", block.index()), color)
+    };
+
+    let style = NodeStyle { title_bg: Some(bgcolor.to_owned()), ..Default::default() };
+    let mut stmts: Vec<String> = data.statements.iter().map(|x| format!("{:?}", x)).collect();
+
+    // add the terminator to the stmts, gsgdt can print it out separately
+    let mut terminator_head = String::new();
+    data.terminator().kind.fmt_head(&mut terminator_head).unwrap();
+    stmts.push(terminator_head);
+
+    Node::new(stmts, label, title, style)
+}
+
+// Must match `[0-9A-Za-z_]*`. This does not appear in the rendered graph, so
+// it does not have to be user friendly.
+pub fn graphviz_safe_def_name(def_id: DefId) -> String {
+    format!("{}_{}", def_id.krate.index(), def_id.index.index(),)
+}
+
+fn node(def_id: DefId, block: BasicBlock) -> String {
+    format!("bb{}__{}", block.index(), graphviz_safe_def_name(def_id))
+}