Implement --unpretty mir-cfg for graphviz output

Also change output for --unpretty mir to output function names in a prettier way.

3 files changed, 68 insertions, 59 deletions

diff --git a/src/librustc_driver/pretty.rs b/src/librustc_driver/pretty.rs
index eb4668e6abb..e716c3e82b0 100644
--- a/src/librustc_driver/pretty.rs
+++ b/src/librustc_driver/pretty.rs
@@ -31,6 +31,7 @@ use rustc_resolve as resolve;
 use rustc_metadata::cstore::CStore;
 
 use rustc_mir::pretty::write_mir_pretty;
+use rustc_mir::graphviz::write_mir_graphviz;
 
 use syntax::ast::{self, BlockCheckMode};
 use syntax::codemap;
@@ -44,6 +45,7 @@ use graphviz as dot;
 
 use std::fs::File;
 use std::io::{self, Write};
+use std::iter;
 use std::option;
 use std::path::PathBuf;
 use std::str::FromStr;
@@ -80,6 +82,7 @@ pub enum PpMode {
     PpmHir(PpSourceMode),
     PpmFlowGraph(PpFlowGraphMode),
     PpmMir,
+    PpmMirCFG,
 }
 
 pub fn parse_pretty(sess: &Session,
@@ -100,6 +103,7 @@ pub fn parse_pretty(sess: &Session,
         ("hir,identified", true) => PpmHir(PpmIdentified),
         ("hir,typed", true) => PpmHir(PpmTyped),
         ("mir", true) => PpmMir,
+        ("mir-cfg", true) => PpmMirCFG,
         ("flowgraph", true) => PpmFlowGraph(PpFlowGraphMode::Default),
         ("flowgraph,unlabelled", true) => PpmFlowGraph(PpFlowGraphMode::UnlabelledEdges),
         _ => {
@@ -574,6 +578,7 @@ fn needs_ast_map(ppm: &PpMode, opt_uii: &Option<UserIdentifiedItem>) -> bool {
         PpmSource(PpmExpandedHygiene) |
         PpmHir(_) |
         PpmMir |
+        PpmMirCFG |
         PpmFlowGraph(_) => true,
         PpmSource(PpmTyped) => panic!("invalid state"),
     }
@@ -590,6 +595,7 @@ fn needs_expansion(ppm: &PpMode) -> bool {
         PpmSource(PpmExpandedHygiene) |
         PpmHir(_) |
         PpmMir |
+        PpmMirCFG |
         PpmFlowGraph(_) => true,
         PpmSource(PpmTyped) => panic!("invalid state"),
     }
@@ -807,9 +813,15 @@ pub fn pretty_print_input(sess: Session,
             })
         }
 
-        (PpmMir, None) => {
-            debug!("pretty printing MIR for whole crate");
-            let ast_map = ast_map.expect("--unpretty mir missing ast_map");
+        (pp_type@PpmMir, uii) | (pp_type@PpmMirCFG, uii) => {
+            let ast_map = ast_map.expect("--unpretty missing ast_map");
+            let nodeid = if let Some(uii) = uii {
+                debug!("pretty printing MIR for {:?}", uii);
+                Some(uii.to_one_node_id("--unpretty", &sess, &ast_map))
+            } else {
+                debug!("pretty printing MIR for whole crate");
+                None
+            };
             abort_on_err(driver::phase_3_run_analysis_passes(&sess,
                                                              &cstore,
                                                              ast_map,
@@ -818,38 +830,25 @@ pub fn pretty_print_input(sess: Session,
                                                              resolve::MakeGlobMap::No,
                                                              |tcx, mir_map, _, _| {
                 if let Some(mir_map) = mir_map {
-                    for (nodeid, mir) in &mir_map.map {
-                        try!(writeln!(out, "MIR for {}", tcx.map.node_to_string(*nodeid)));
-                        try!(write_mir_pretty(mir, &mut out));
+                    if let Some(nodeid) = nodeid {
+                        let mir = mir_map.map.get(&nodeid).unwrap_or_else(|| {
+                            sess.fatal(&format!("no MIR map entry for node {}", nodeid))
+                        });
+                        try!(match pp_type {
+                            PpmMir => write_mir_pretty(tcx, iter::once((&nodeid, mir)), &mut out),
+                            _ => write_mir_graphviz(tcx, iter::once((&nodeid, mir)), &mut out)
+                        });
+                    } else {
+                        try!(match pp_type {
+                            PpmMir => write_mir_pretty(tcx, mir_map.map.iter(), &mut out),
+                            _ => write_mir_graphviz(tcx, mir_map.map.iter(), &mut out)
+                        });
                     }
                 }
                 Ok(())
             }), &sess)
         }
 
-        (PpmMir, Some(uii)) => {
-            debug!("pretty printing MIR for {:?}", uii);
-            let ast_map = ast_map.expect("--unpretty mir missing ast_map");
-            let nodeid = uii.to_one_node_id("--unpretty", &sess, &ast_map);
-
-            abort_on_err(driver::phase_3_run_analysis_passes(&sess,
-                                                             &cstore,
-                                                             ast_map,
-                                                             &arenas,
-                                                             &id,
-                                                             resolve::MakeGlobMap::No,
-                                                             |tcx, mir_map, _, _| {
-                if let Some(mir_map) = mir_map {
-                    try!(writeln!(out, "MIR for {}", tcx.map.node_to_string(nodeid)));
-                    let mir = mir_map.map.get(&nodeid).unwrap_or_else(|| {
-                        sess.fatal(&format!("no MIR map entry for node {}", nodeid))
-                    });
-                    try!(write_mir_pretty(mir, &mut out));
-                }
-                Ok(())
-            }), &sess)
-        }
-
         (PpmFlowGraph(mode), opt_uii) => {
             debug!("pretty printing flow graph for {:?}", opt_uii);
             let uii = opt_uii.unwrap_or_else(|| {
diff --git a/src/librustc_mir/graphviz.rs b/src/librustc_mir/graphviz.rs
index 1b8fe650558..f705c0591b5 100644
--- a/src/librustc_mir/graphviz.rs
+++ b/src/librustc_mir/graphviz.rs
@@ -13,30 +13,34 @@ use rustc::mir::repr::*;
 use rustc::middle::ty;
 use std::fmt::Debug;
 use std::io::{self, Write};
+use syntax::ast::NodeId;
 
-/// Write a graphviz DOT graph for the given MIR.
-pub fn write_mir_graphviz<W: Write>(mir: &Mir, w: &mut W) -> io::Result<()> {
-    try!(writeln!(w, "digraph Mir {{"));
+/// Write a graphviz DOT graph of a list of MIRs.
+pub fn write_mir_graphviz<'a, 't, W, I>(tcx: &ty::TyCtxt<'t>, iter: I, w: &mut W) -> io::Result<()>
+where W: Write, I: Iterator<Item=(&'a NodeId, &'a Mir<'a>)> {
+    for (&nodeid, mir) in iter {
+        try!(writeln!(w, "digraph Mir_{} {{", nodeid));
 
-    // Global graph properties
-    try!(writeln!(w, r#"    graph [fontname="monospace"];"#));
-    try!(writeln!(w, r#"    node [fontname="monospace"];"#));
-    try!(writeln!(w, r#"    edge [fontname="monospace"];"#));
+        // Global graph properties
+        try!(writeln!(w, r#"    graph [fontname="monospace"];"#));
+        try!(writeln!(w, r#"    node [fontname="monospace"];"#));
+        try!(writeln!(w, r#"    edge [fontname="monospace"];"#));
 
-    // Graph label
-    try!(write_graph_label(mir, w));
+        // Graph label
+        try!(write_graph_label(tcx, nodeid, mir, w));
 
-    // Nodes
-    for block in mir.all_basic_blocks() {
-        try!(write_node(block, mir, w));
-    }
+        // Nodes
+        for block in mir.all_basic_blocks() {
+            try!(write_node(block, mir, w));
+        }
 
-    // Edges
-    for source in mir.all_basic_blocks() {
-        try!(write_edges(source, mir, w));
+        // Edges
+        for source in mir.all_basic_blocks() {
+            try!(write_edges(source, mir, w));
+        }
+        try!(writeln!(w, "}}"))
     }
-
-    writeln!(w, "}}")
+    Ok(())
 }
 
 /// Write a graphviz DOT node for the given basic block.
@@ -84,8 +88,9 @@ fn write_edges<W: Write>(source: BasicBlock, mir: &Mir, w: &mut W) -> io::Result
 /// Write the graphviz DOT label for the overall graph. This is essentially a block of text that
 /// will appear below the graph, showing the type of the `fn` this MIR represents and the types of
 /// all the variables and temporaries.
-fn write_graph_label<W: Write>(mir: &Mir, w: &mut W) -> io::Result<()> {
-    try!(write!(w, "    label=<fn("));
+fn write_graph_label<W: Write>(tcx: &ty::TyCtxt, nid: NodeId, mir: &Mir, w: &mut W)
+-> io::Result<()> {
+    try!(write!(w, "    label=<fn {}(", dot::escape_html(&tcx.map.path_to_string(nid))));
 
     // fn argument types.
     for (i, arg) in mir.arg_decls.iter().enumerate() {
diff --git a/src/librustc_mir/pretty.rs b/src/librustc_mir/pretty.rs
index ea4036a4d37..c3fe4df40b5 100644
--- a/src/librustc_mir/pretty.rs
+++ b/src/librustc_mir/pretty.rs
@@ -11,19 +11,22 @@
 use rustc::mir::repr::*;
 use rustc::middle::ty;
 use std::io::{self, Write};
+use syntax::ast::NodeId;
 
 const INDENT: &'static str = "    ";
 
 /// Write out a human-readable textual representation for the given MIR.
-pub fn write_mir_pretty<W: Write>(mir: &Mir, w: &mut W) -> io::Result<()> {
-    try!(write_mir_intro(mir, w));
-
-    // Nodes
-    for block in mir.all_basic_blocks() {
-        try!(write_basic_block(block, mir, w));
+pub fn write_mir_pretty<'a, 't, W, I>(tcx: &ty::TyCtxt<'t>, iter: I, w: &mut W) -> io::Result<()>
+where W: Write, I: Iterator<Item=(&'a NodeId, &'a Mir<'a>)> {
+    for (&nodeid, mir) in iter {
+        try!(write_mir_intro(tcx, nodeid, mir, w));
+        // Nodes
+        for block in mir.all_basic_blocks() {
+            try!(write_basic_block(block, mir, w));
+        }
+        try!(writeln!(w, "}}"))
     }
-
-    writeln!(w, "}}")
+    Ok(())
 }
 
 /// Write out a human-readable textual representation for the given basic block.
@@ -46,8 +49,10 @@ fn write_basic_block<W: Write>(block: BasicBlock, mir: &Mir, w: &mut W) -> io::R
 
 /// Write out a human-readable textual representation of the MIR's `fn` type and the types of its
 /// local variables (both user-defined bindings and compiler temporaries).
-fn write_mir_intro<W: Write>(mir: &Mir, w: &mut W) -> io::Result<()> {
-    try!(write!(w, "fn("));
+fn write_mir_intro<W: Write>(tcx: &ty::TyCtxt, nid: NodeId, mir: &Mir, w: &mut W)
+-> io::Result<()> {
+
+    try!(write!(w, "fn {}(", tcx.map.path_to_string(nid)));
 
     // fn argument types.
     for (i, arg) in mir.arg_decls.iter().enumerate() {