Refactor pretty printing to use more of the driver

3 files changed, 330 insertions, 263 deletions

diff --git a/src/librustc_driver/driver.rs b/src/librustc_driver/driver.rs
index bf1457c576e..2b56366f1c6 100644
--- a/src/librustc_driver/driver.rs
+++ b/src/librustc_driver/driver.rs
@@ -112,7 +112,7 @@ pub fn compile_input(sess: &Session,
             let expanded_crate = phase_2_configure_and_expand(sess,
                                                               &cstore,
                                                               krate,
-                                                              &id[..],
+                                                              &id,
                                                               addl_plugins)?;
 
             (outputs, expanded_crate, id)
@@ -123,8 +123,10 @@ pub fn compile_input(sess: &Session,
                                 CompileState::state_after_expand(input,
                                                                  sess,
                                                                  outdir,
+                                                                 output,
+                                                                 &cstore,
                                                                  &expanded_crate,
-                                                                 &id[..]),
+                                                                 &id),
                                 Ok(()));
 
         let expanded_crate = assign_node_ids(sess, expanded_crate);
@@ -169,10 +171,13 @@ pub fn compile_input(sess: &Session,
                                     CompileState::state_after_write_deps(input,
                                                                          sess,
                                                                          outdir,
+                                                                         output,
+                                                                         &arenas,
+                                                                         &cstore,
                                                                          &hir_map,
                                                                          &expanded_crate,
                                                                          &hir_map.krate(),
-                                                                         &id[..]),
+                                                                         &id),
                                     Ok(()));
         }
 
@@ -202,8 +207,9 @@ pub fn compile_input(sess: &Session,
                 let _ignore = tcx.dep_graph.in_ignore();
 
                 let mut state = CompileState::state_after_analysis(input,
-                                                                   &tcx.sess,
+                                                                   sess,
                                                                    outdir,
+                                                                   output,
                                                                    opt_crate,
                                                                    tcx.map.krate(),
                                                                    &analysis,
@@ -243,7 +249,7 @@ pub fn compile_input(sess: &Session,
 
     controller_entry_point!(after_llvm,
                             sess,
-                            CompileState::state_after_llvm(input, sess, outdir, &trans),
+                            CompileState::state_after_llvm(input, sess, outdir, output, &trans),
                             phase5_result);
     phase5_result?;
 
@@ -334,34 +340,36 @@ impl<'a> PhaseController<'a> {
 /// State that is passed to a callback. What state is available depends on when
 /// during compilation the callback is made. See the various constructor methods
 /// (`state_*`) in the impl to see which data is provided for any given entry point.
-pub struct CompileState<'a, 'ast: 'a, 'tcx: 'a> {
+pub struct CompileState<'a, 'b, 'ast: 'a, 'tcx: 'b> where 'ast: 'tcx {
     pub input: &'a Input,
-    pub session: &'a Session,
+    pub session: &'ast Session,
     pub krate: Option<ast::Crate>,
     pub cstore: Option<&'a CStore>,
     pub crate_name: Option<&'a str>,
     pub output_filenames: Option<&'a OutputFilenames>,
     pub out_dir: Option<&'a Path>,
     pub out_file: Option<&'a Path>,
+    pub arenas: Option<&'ast ty::CtxtArenas<'ast>>,
     pub expanded_crate: Option<&'a ast::Crate>,
     pub hir_crate: Option<&'a hir::Crate>,
     pub ast_map: Option<&'a hir_map::Map<'ast>>,
-    pub mir_map: Option<&'a MirMap<'tcx>>,
+    pub mir_map: Option<&'b MirMap<'tcx>>,
     pub analysis: Option<&'a ty::CrateAnalysis<'a>>,
-    pub tcx: Option<&'a TyCtxt<'tcx>>,
+    pub tcx: Option<&'b TyCtxt<'tcx>>,
     pub trans: Option<&'a trans::CrateTranslation>,
 }
 
-impl<'a, 'ast, 'tcx> CompileState<'a, 'ast, 'tcx> {
+impl<'a, 'b, 'ast, 'tcx> CompileState<'a, 'b, 'ast, 'tcx> {
     fn empty(input: &'a Input,
-             session: &'a Session,
+             session: &'ast Session,
              out_dir: &'a Option<PathBuf>)
-             -> CompileState<'a, 'ast, 'tcx> {
+             -> CompileState<'a, 'b, 'ast, 'tcx> {
         CompileState {
             input: input,
             session: session,
             out_dir: out_dir.as_ref().map(|s| &**s),
             out_file: None,
+            arenas: None,
             krate: None,
             cstore: None,
             crate_name: None,
@@ -377,12 +385,12 @@ impl<'a, 'ast, 'tcx> CompileState<'a, 'ast, 'tcx> {
     }
 
     fn state_after_parse(input: &'a Input,
-                         session: &'a Session,
+                         session: &'ast Session,
                          out_dir: &'a Option<PathBuf>,
                          out_file: &'a Option<PathBuf>,
                          krate: ast::Crate,
                          cstore: &'a CStore)
-                         -> CompileState<'a, 'ast, 'tcx> {
+                         -> CompileState<'a, 'b, 'ast, 'tcx> {
         CompileState {
             krate: Some(krate),
             cstore: Some(cstore),
@@ -392,45 +400,56 @@ impl<'a, 'ast, 'tcx> CompileState<'a, 'ast, 'tcx> {
     }
 
     fn state_after_expand(input: &'a Input,
-                          session: &'a Session,
+                          session: &'ast Session,
                           out_dir: &'a Option<PathBuf>,
+                          out_file: &'a Option<PathBuf>,
+                          cstore: &'a CStore,
                           expanded_crate: &'a ast::Crate,
                           crate_name: &'a str)
-                          -> CompileState<'a, 'ast, 'tcx> {
+                          -> CompileState<'a, 'b, 'ast, 'tcx> {
         CompileState {
             crate_name: Some(crate_name),
+            cstore: Some(cstore),
             expanded_crate: Some(expanded_crate),
+            out_file: out_file.as_ref().map(|s| &**s),
             ..CompileState::empty(input, session, out_dir)
         }
     }
 
     fn state_after_write_deps(input: &'a Input,
-                              session: &'a Session,
+                              session: &'ast Session,
                               out_dir: &'a Option<PathBuf>,
+                              out_file: &'a Option<PathBuf>,
+                              arenas: &'ast ty::CtxtArenas<'ast>,
+                              cstore: &'a CStore,
                               hir_map: &'a hir_map::Map<'ast>,
                               krate: &'a ast::Crate,
                               hir_crate: &'a hir::Crate,
                               crate_name: &'a str)
-                              -> CompileState<'a, 'ast, 'tcx> {
+                              -> CompileState<'a, 'b, 'ast, 'tcx> {
         CompileState {
             crate_name: Some(crate_name),
+            arenas: Some(arenas),
+            cstore: Some(cstore),
             ast_map: Some(hir_map),
             expanded_crate: Some(krate),
             hir_crate: Some(hir_crate),
+            out_file: out_file.as_ref().map(|s| &**s),
             ..CompileState::empty(input, session, out_dir)
         }
     }
 
     fn state_after_analysis(input: &'a Input,
-                            session: &'a Session,
+                            session: &'ast Session,
                             out_dir: &'a Option<PathBuf>,
+                            out_file: &'a Option<PathBuf>,
                             krate: Option<&'a ast::Crate>,
                             hir_crate: &'a hir::Crate,
-                            analysis: &'a ty::CrateAnalysis,
-                            mir_map: Option<&'a MirMap<'tcx>>,
-                            tcx: &'a TyCtxt<'tcx>,
+                            analysis: &'a ty::CrateAnalysis<'a>,
+                            mir_map: Option<&'b MirMap<'tcx>>,
+                            tcx: &'b TyCtxt<'tcx>,
                             crate_name: &'a str)
-                            -> CompileState<'a, 'ast, 'tcx> {
+                            -> CompileState<'a, 'b, 'ast, 'tcx> {
         CompileState {
             analysis: Some(analysis),
             mir_map: mir_map,
@@ -438,17 +457,23 @@ impl<'a, 'ast, 'tcx> CompileState<'a, 'ast, 'tcx> {
             expanded_crate: krate,
             hir_crate: Some(hir_crate),
             crate_name: Some(crate_name),
+            out_file: out_file.as_ref().map(|s| &**s),
             ..CompileState::empty(input, session, out_dir)
         }
     }
 
 
     fn state_after_llvm(input: &'a Input,
-                        session: &'a Session,
+                        session: &'ast Session,
                         out_dir: &'a Option<PathBuf>,
+                        out_file: &'a Option<PathBuf>,
                         trans: &'a trans::CrateTranslation)
-                        -> CompileState<'a, 'ast, 'tcx> {
-        CompileState { trans: Some(trans), ..CompileState::empty(input, session, out_dir) }
+                        -> CompileState<'a, 'b, 'ast, 'tcx> {
+        CompileState {
+            trans: Some(trans),
+            out_file: out_file.as_ref().map(|s| &**s),
+            ..CompileState::empty(input, session, out_dir)
+        }
     }
 }
 
@@ -814,16 +839,16 @@ pub fn phase_3_run_analysis_passes<'tcx, F, R>(sess: &'tcx Session,
     let index = stability::Index::new(&hir_map);
 
     TyCtxt::create_and_enter(sess,
-                               arenas,
-                               def_map,
-                               named_region_map,
-                               hir_map,
-                               freevars,
-                               region_map,
-                               lang_items,
-                               index,
-                               name,
-                               |tcx| {
+                             arenas,
+                             def_map,
+                             named_region_map,
+                             hir_map,
+                             freevars,
+                             region_map,
+                             lang_items,
+                             index,
+                             name,
+                             |tcx| {
         time(time_passes,
              "load_dep_graph",
              || rustc_incremental::load_dep_graph(tcx));
diff --git a/src/librustc_driver/lib.rs b/src/librustc_driver/lib.rs
index 8ffdf6843e7..74f6d3f2515 100644
--- a/src/librustc_driver/lib.rs
+++ b/src/librustc_driver/lib.rs
@@ -450,17 +450,39 @@ impl<'a> CompilerCalls<'a> for RustcDefaultCalls {
     fn build_controller(&mut self, sess: &Session, matches: &getopts::Matches) -> CompileController<'a> {
         let mut control = CompileController::basic();
 
-        if let Some((ppm, opt_uii)) = parse_pretty(&sess, &matches) {
-            control.after_parse.stop = Compilation::Stop;
-            control.after_parse.callback = box move |state| {
-                pretty::pretty_print_input(state.session,
-                                           state.cstore.unwrap(),
-                                           state.input,
-                                           state.krate.take().unwrap(),
-                                           ppm,
-                                           opt_uii.clone(),
-                                           state.out_file);
-            };
+        if let Some((ppm, opt_uii)) = parse_pretty(sess, matches) {
+            if ppm.needs_ast_map(&opt_uii) {
+                control.after_write_deps.stop = Compilation::Stop;
+
+                control.after_parse.callback = box move |state| {
+                    state.krate = Some(pretty::fold_crate(state.krate.take().unwrap(), ppm));
+                };
+                control.after_write_deps.callback = box move |state| {
+                    pretty::print_after_write_deps(state.session,
+                                                   state.cstore.unwrap(),
+                                                   state.ast_map.unwrap(),
+                                                   state.input,
+                                                   &state.expanded_crate.take().unwrap(),
+                                                   state.crate_name.unwrap(),
+                                                   ppm,
+                                                   state.arenas.unwrap(),
+                                                   opt_uii.clone(),
+                                                   state.out_file);
+                };
+            } else {
+                control.after_parse.stop = Compilation::Stop;
+
+                control.after_parse.callback = box move |state| {
+                    let krate = pretty::fold_crate(state.krate.take().unwrap(), ppm);
+                    pretty::print_after_parsing(state.session,
+                                                state.input,
+                                                &krate,
+                                                ppm,
+                                                state.out_file);
+                };
+            }
+
+            return control;
         }
 
         if sess.opts.parse_only || sess.opts.debugging_opts.show_span.is_some() ||
diff --git a/src/librustc_driver/pretty.rs b/src/librustc_driver/pretty.rs
index 80d8a40b4e9..223b6629a88 100644
--- a/src/librustc_driver/pretty.rs
+++ b/src/librustc_driver/pretty.rs
@@ -15,8 +15,6 @@ pub use self::PpSourceMode::*;
 pub use self::PpMode::*;
 use self::NodesMatchingUII::*;
 
-use rustc_trans::back::link;
-
 use {driver, abort_on_err};
 
 use rustc::dep_graph::DepGraph;
@@ -55,7 +53,6 @@ use std::str::FromStr;
 use rustc::hir::map as hir_map;
 use rustc::hir::map::{blocks, NodePrinter};
 use rustc::hir;
-use rustc::hir::lowering::{lower_crate, LoweringContext};
 use rustc::hir::print as pprust_hir;
 
 use rustc::mir::mir_map::MirMap;
@@ -89,6 +86,32 @@ pub enum PpMode {
     PpmMirCFG,
 }
 
+impl PpMode {
+    pub fn needs_ast_map(&self, opt_uii: &Option<UserIdentifiedItem>) -> bool {
+        match *self {
+            PpmSource(PpmNormal) |
+            PpmSource(PpmEveryBodyLoops) |
+            PpmSource(PpmIdentified) => opt_uii.is_some(),
+
+            PpmSource(PpmExpanded) |
+            PpmSource(PpmExpandedIdentified) |
+            PpmSource(PpmExpandedHygiene) |
+            PpmHir(_) |
+            PpmMir |
+            PpmMirCFG |
+            PpmFlowGraph(_) => true,
+            PpmSource(PpmTyped) => panic!("invalid state"),
+        }
+    }
+
+    pub fn needs_analysis(&self) -> bool {
+        match *self {
+             PpmMir | PpmMirCFG | PpmFlowGraph(_) => true,
+             _ => false,
+        }
+    }
+}
+
 pub fn parse_pretty(sess: &Session,
                     name: &str,
                     extended: bool)
@@ -147,7 +170,7 @@ impl PpSourceMode {
     /// Constructs a `PrinterSupport` object and passes it to `f`.
     fn call_with_pp_support<'tcx, A, B, F>(&self,
                                            sess: &'tcx Session,
-                                           ast_map: Option<hir_map::Map<'tcx>>,
+                                           ast_map: Option<&hir_map::Map<'tcx>>,
                                            payload: B,
                                            f: F)
                                            -> A
@@ -157,7 +180,7 @@ impl PpSourceMode {
             PpmNormal | PpmEveryBodyLoops | PpmExpanded => {
                 let annotation = NoAnn {
                     sess: sess,
-                    ast_map: ast_map,
+                    ast_map: ast_map.map(|m| m.clone()),
                 };
                 f(&annotation, payload)
             }
@@ -165,14 +188,14 @@ impl PpSourceMode {
             PpmIdentified | PpmExpandedIdentified => {
                 let annotation = IdentifiedAnnotation {
                     sess: sess,
-                    ast_map: ast_map,
+                    ast_map: ast_map.map(|m| m.clone()),
                 };
                 f(&annotation, payload)
             }
             PpmExpandedHygiene => {
                 let annotation = HygieneAnnotation {
                     sess: sess,
-                    ast_map: ast_map,
+                    ast_map: ast_map.map(|m| m.clone()),
                 };
                 f(&annotation, payload)
             }
@@ -582,40 +605,6 @@ impl UserIdentifiedItem {
     }
 }
 
-fn needs_ast_map(ppm: &PpMode, opt_uii: &Option<UserIdentifiedItem>) -> bool {
-    match *ppm {
-        PpmSource(PpmNormal) |
-        PpmSource(PpmEveryBodyLoops) |
-        PpmSource(PpmIdentified) => opt_uii.is_some(),
-
-        PpmSource(PpmExpanded) |
-        PpmSource(PpmExpandedIdentified) |
-        PpmSource(PpmExpandedHygiene) |
-        PpmHir(_) |
-        PpmMir |
-        PpmMirCFG |
-        PpmFlowGraph(_) => true,
-        PpmSource(PpmTyped) => panic!("invalid state"),
-    }
-}
-
-fn needs_expansion(ppm: &PpMode) -> bool {
-    match *ppm {
-        PpmSource(PpmNormal) |
-        PpmSource(PpmEveryBodyLoops) |
-        PpmSource(PpmIdentified) => false,
-
-        PpmSource(PpmExpanded) |
-        PpmSource(PpmExpandedIdentified) |
-        PpmSource(PpmExpandedHygiene) |
-        PpmHir(_) |
-        PpmMir |
-        PpmMirCFG |
-        PpmFlowGraph(_) => true,
-        PpmSource(PpmTyped) => panic!("invalid state"),
-    }
-}
-
 struct ReplaceBodyWithLoop {
     within_static_or_const: bool,
 }
@@ -700,89 +689,176 @@ impl fold::Folder for ReplaceBodyWithLoop {
     }
 }
 
-pub fn pretty_print_input(sess: &Session,
-                          cstore: &CStore,
-                          input: &Input,
-                          krate: ast::Crate,
-                          ppm: PpMode,
-                          opt_uii: Option<UserIdentifiedItem>,
-                          ofile: Option<&Path>) {
-    let krate = if let PpmSource(PpmEveryBodyLoops) = ppm {
-        let mut fold = ReplaceBodyWithLoop::new();
-        fold.fold_crate(krate)
-    } else {
-        krate
+fn print_flowgraph<'tcx, W: Write>(variants: Vec<borrowck_dot::Variant>,
+                                   tcx: &TyCtxt<'tcx>,
+                                   mir_map: Option<&MirMap<'tcx>>,
+                                   code: blocks::Code,
+                                   mode: PpFlowGraphMode,
+                                   mut out: W)
+                                   -> io::Result<()> {
+    let cfg = match code {
+        blocks::BlockCode(block) => cfg::CFG::new(tcx, &block),
+        blocks::FnLikeCode(fn_like) => cfg::CFG::new(tcx, &fn_like.body()),
+    };
+    let labelled_edges = mode != PpFlowGraphMode::UnlabelledEdges;
+    let lcfg = LabelledCFG {
+        ast_map: &tcx.map,
+        cfg: &cfg,
+        name: format!("node_{}", code.id()),
+        labelled_edges: labelled_edges,
     };
 
-    let id = link::find_crate_name(Some(sess), &krate.attrs, input);
+    match code {
+        _ if variants.is_empty() => {
+            let r = dot::render(&lcfg, &mut out);
+            return expand_err_details(r);
+        }
+        blocks::BlockCode(_) => {
+            tcx.sess.err("--pretty flowgraph with -Z flowgraph-print annotations requires \
+                          fn-like node id.");
+            return Ok(());
+        }
+        blocks::FnLikeCode(fn_like) => {
+            let (bccx, analysis_data) =
+                borrowck::build_borrowck_dataflow_data_for_fn(tcx,
+                                                              mir_map,
+                                                              fn_like.to_fn_parts(),
+                                                              &cfg);
 
-    let is_expanded = needs_expansion(&ppm);
-    let compute_ast_map = needs_ast_map(&ppm, &opt_uii);
-    let krate = if compute_ast_map {
-        match driver::phase_2_configure_and_expand(sess, &cstore, krate, &id, None) {
-            Err(_) => return,
-            Ok(k) => driver::assign_node_ids(sess, k),
+            let lcfg = borrowck_dot::DataflowLabeller {
+                inner: lcfg,
+                variants: variants,
+                borrowck_ctxt: &bccx,
+                analysis_data: &analysis_data,
+            };
+            let r = dot::render(&lcfg, &mut out);
+            return expand_err_details(r);
         }
-    } else {
-        krate
-    };
+    }
 
-    // There is some twisted, god-forsaken tangle of lifetimes here which makes
-    // the ordering of stuff super-finicky.
-    let mut hir_forest;
-    let mut _defs = None;
-    let dep_graph = DepGraph::new(false);
-    let arenas = ty::CtxtArenas::new();
-    let _ignore = dep_graph.in_ignore();
-    let ast_map = if compute_ast_map {
-        _defs = Some(RefCell::new(hir_map::collect_definitions(&krate)));
-        let defs = _defs.as_ref().unwrap();
-        LocalCrateReader::new(sess, &cstore, defs, &krate, &id).read_crates(&dep_graph);
-        let lcx = LoweringContext::new(sess, Some(&krate), defs);
-
-        hir_forest = hir_map::Forest::new(lower_crate(&lcx, &krate), dep_graph.clone());
-        let map = hir_map::map_crate(&mut hir_forest, defs);
-        Some(map)
+    fn expand_err_details(r: io::Result<()>) -> io::Result<()> {
+        r.map_err(|ioerr| {
+            io::Error::new(io::ErrorKind::Other,
+                           &format!("graphviz::render failed: {}", ioerr)[..])
+        })
+    }
+}
+
+pub fn fold_crate(krate: ast::Crate, ppm: PpMode) -> ast::Crate {
+    if let PpmSource(PpmEveryBodyLoops) = ppm {
+        let mut fold = ReplaceBodyWithLoop::new();
+        fold.fold_crate(krate)
     } else {
-        None
-    };
+        krate
+    }
+}
 
+fn get_source(input: &Input, sess: &Session) -> (Vec<u8>, String) {
     let src_name = driver::source_name(input);
     let src = sess.codemap()
-                  .get_filemap(&src_name[..])
+                  .get_filemap(&src_name)
                   .src
                   .as_ref()
                   .unwrap()
                   .as_bytes()
                   .to_vec();
+    (src, src_name)
+}
+
+fn write_output(out: Vec<u8>, ofile: Option<&Path>) {
+    match ofile {
+        None => print!("{}", String::from_utf8(out).unwrap()),
+        Some(p) => {
+            match File::create(p) {
+                Ok(mut w) => w.write_all(&out).unwrap(),
+                Err(e) => panic!("print-print failed to open {} due to {}", p.display(), e),
+            }
+        }
+    }
+}
+
+pub fn print_after_parsing(sess: &Session,
+                           input: &Input,
+                           krate: &ast::Crate,
+                           ppm: PpMode,
+                           ofile: Option<&Path>) {
+    let dep_graph = DepGraph::new(false);
+    let _ignore = dep_graph.in_ignore();
+
+    let (src, src_name) = get_source(input, sess);
+
     let mut rdr = &*src;
+    let mut out = Vec::new();
+
+    if let PpmSource(s) = ppm {
+        // Silently ignores an identified node.
+        let out: &mut Write = &mut out;
+        s.call_with_pp_support(sess, None, box out, |annotation, out| {
+            debug!("pretty printing source code {:?}", s);
+            let sess = annotation.sess();
+            pprust::print_crate(sess.codemap(),
+                                sess.diagnostic(),
+                                krate,
+                                src_name.to_string(),
+                                &mut rdr,
+                                out,
+                                annotation.pp_ann(),
+                                false)
+        }).unwrap()
+    } else {
+        unreachable!();
+    };
+
+    write_output(out, ofile);
+}
+
+pub fn print_after_write_deps<'tcx, 'a: 'tcx>(sess: &'a Session,
+                                              cstore: &CStore,
+                                              ast_map: &hir_map::Map<'tcx>,
+                                              input: &Input,
+                                              krate: &ast::Crate,
+                                              crate_name: &str,
+                                              ppm: PpMode,
+                                              arenas: &'tcx ty::CtxtArenas<'tcx>,
+                                              opt_uii: Option<UserIdentifiedItem>,
+                                              ofile: Option<&Path>) {
+    let dep_graph = DepGraph::new(false);
+    let _ignore = dep_graph.in_ignore();
 
+    if ppm.needs_analysis() {
+        print_with_analysis(sess, cstore, ast_map, crate_name, arenas, ppm, opt_uii, ofile);
+        return;
+    }
+
+    let (src, src_name) = get_source(input, sess);
+
+    let mut rdr = &src[..];
     let mut out = Vec::new();
 
     match (ppm, opt_uii) {
         (PpmSource(s), _) => {
             // Silently ignores an identified node.
             let out: &mut Write = &mut out;
-            s.call_with_pp_support(sess, ast_map, box out, |annotation, out| {
+            s.call_with_pp_support(sess, Some(ast_map), box out, |annotation, out| {
                 debug!("pretty printing source code {:?}", s);
                 let sess = annotation.sess();
                 pprust::print_crate(sess.codemap(),
                                     sess.diagnostic(),
-                                    &krate,
+                                    krate,
                                     src_name.to_string(),
                                     &mut rdr,
                                     out,
                                     annotation.pp_ann(),
-                                    is_expanded)
+                                    true)
             })
         }
 
         (PpmHir(s), None) => {
             let out: &mut Write = &mut out;
             s.call_with_pp_support_hir(sess,
-                                       &ast_map.unwrap(),
-                                       &arenas,
-                                       &id,
+                                       ast_map,
+                                       arenas,
+                                       crate_name,
                                        box out,
                                        |annotation, out, krate| {
                                            debug!("pretty printing source code {:?}", s);
@@ -794,16 +870,16 @@ pub fn pretty_print_input(sess: &Session,
                                                                    &mut rdr,
                                                                    out,
                                                                    annotation.pp_ann(),
-                                                                   is_expanded)
+                                                                   true)
                                        })
         }
 
         (PpmHir(s), Some(uii)) => {
             let out: &mut Write = &mut out;
             s.call_with_pp_support_hir(sess,
-                                       &ast_map.unwrap(),
-                                       &arenas,
-                                       &id,
+                                       ast_map,
+                                       arenas,
+                                       crate_name,
                                        (out,uii),
                                        |annotation, (out,uii), _| {
                 debug!("pretty printing source code {:?}", s);
@@ -829,156 +905,100 @@ pub fn pretty_print_input(sess: &Session,
                 }
                 pp::eof(&mut pp_state.s)
             })
-        }
+       }
+       _ => unreachable!(),
+    }.unwrap();
+
+    write_output(out, ofile);
+}
+
+// In an ideal world, this would be a public function called by the driver after
+// analsysis is performed. However, we want to call `phase_3_run_analysis_passes`
+// with a different callback than the standard driver, so that isn't easy.
+// Instead, we call that function ourselves.
+fn print_with_analysis<'tcx, 'a: 'tcx>(sess: &'a Session,
+                                       cstore: &CStore,
+                                       ast_map: &hir_map::Map<'tcx>,
+                                       crate_name: &str,
+                                       arenas: &'tcx ty::CtxtArenas<'tcx>,
+                                       ppm: PpMode,
+                                       uii: Option<UserIdentifiedItem>,
+                                       ofile: Option<&Path>) {
+    let nodeid = if let Some(uii) = uii {
+        debug!("pretty printing for {:?}", uii);
+        Some(uii.to_one_node_id("--unpretty", sess, &ast_map))
+    } else {
+        debug!("pretty printing for whole crate");
+        None
+    };
 
-        (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,
-                                                             ast_map,
-                                                             &arenas,
-                                                             &id,
-                                                             resolve::MakeGlobMap::No,
-                                                             |tcx, mir_map, _, _| {
+    let mut out = Vec::new();
+
+    abort_on_err(driver::phase_3_run_analysis_passes(sess,
+                                                     ast_map.clone(),
+                                                     arenas,
+                                                     crate_name,
+                                                     resolve::MakeGlobMap::No,
+                                                     |tcx, mir_map, _, _| {
+        match ppm {
+            PpmMir | PpmMirCFG => {
                 if let Some(mir_map) = mir_map {
                     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))
                         });
-                        match pp_type {
+                        match ppm {
                             PpmMir => write_mir_pretty(tcx, iter::once((&nodeid, mir)), &mut out),
-                            _ => write_mir_graphviz(tcx, iter::once((&nodeid, mir)), &mut out)
+                            PpmMirCFG => write_mir_graphviz(tcx, iter::once((&nodeid, mir)), &mut out),
+                            _ => unreachable!(),
                         }?;
                     } else {
-                        match pp_type {
+                        match ppm {
                             PpmMir => write_mir_pretty(tcx, mir_map.map.iter(), &mut out),
-                            _ => write_mir_graphviz(tcx, mir_map.map.iter(), &mut out)
+                            PpmMirCFG => write_mir_graphviz(tcx, mir_map.map.iter(), &mut out),
+                            _ => unreachable!(),
                         }?;
                     }
                 }
                 Ok(())
-            }), sess)
-        }
-
-        (PpmFlowGraph(mode), opt_uii) => {
-            debug!("pretty printing flow graph for {:?}", opt_uii);
-            let uii = opt_uii.unwrap_or_else(|| {
-                sess.fatal(&format!("`pretty flowgraph=..` needs NodeId (int) or
-                                     \
-                                     unique path suffix (b::c::d)"))
+            }
+            PpmFlowGraph(mode) => {
+                let nodeid = nodeid.expect("`pretty flowgraph=..` needs NodeId (int) or \
+                                            unique path suffix (b::c::d)");
+                let node = tcx.map.find(nodeid).unwrap_or_else(|| {
+                    tcx.sess.fatal(&format!("--pretty flowgraph couldn't find id: {}", nodeid))
+                });
 
-            });
-            let ast_map = ast_map.expect("--pretty flowgraph missing ast_map");
-            let nodeid = uii.to_one_node_id("--pretty", sess, &ast_map);
+                let code = blocks::Code::from_node(node);
+                match code {
+                    Some(code) => {
+                        let variants = gather_flowgraph_variants(tcx.sess);
 
-            let node = ast_map.find(nodeid).unwrap_or_else(|| {
-                sess.fatal(&format!("--pretty flowgraph couldn't find id: {}", nodeid))
-            });
+                        let out: &mut Write = &mut out;
 
-            let code = blocks::Code::from_node(node);
-            let out: &mut Write = &mut out;
-            match code {
-                Some(code) => {
-                    let variants = gather_flowgraph_variants(sess);
-                    abort_on_err(driver::phase_3_run_analysis_passes(sess,
-                                                                     ast_map,
-                                                                     &arenas,
-                                                                     &id,
-                                                                     resolve::MakeGlobMap::No,
-                                                                     |tcx, mir_map, _, _| {
                         print_flowgraph(variants,
                                         tcx,
                                         mir_map.as_ref(),
                                         code,
                                         mode,
                                         out)
-                    }), sess)
-                }
-                None => {
-                    let message = format!("--pretty=flowgraph needs block, fn, or method; got \
-                                           {:?}",
-                                          node);
-
-                    // point to what was found, if there's an
-                    // accessible span.
-                    match ast_map.opt_span(nodeid) {
-                        Some(sp) => sess.span_fatal(sp, &message[..]),
-                        None => sess.fatal(&message[..]),
+                    }
+                    None => {
+                        let message = format!("--pretty=flowgraph needs block, fn, or method; got \
+                                               {:?}",
+                                              node);
+
+                        // Point to what was found, if there's an accessible span.
+                        match tcx.map.opt_span(nodeid) {
+                            Some(sp) => tcx.sess.span_fatal(sp, &message),
+                            None => tcx.sess.fatal(&message),
+                        }
                     }
                 }
             }
+            _ => unreachable!(),
         }
-    }
-    .unwrap();
-
-    match ofile {
-        None => print!("{}", String::from_utf8(out).unwrap()),
-        Some(p) => {
-            match File::create(p) {
-                Ok(mut w) => w.write_all(&out).unwrap(),
-                Err(e) => panic!("print-print failed to open {} due to {}", p.display(), e),
-            }
-        }
-    }
-}
+    }), sess).unwrap();
 
-fn print_flowgraph<'tcx, W: Write>(variants: Vec<borrowck_dot::Variant>,
-                                   tcx: &TyCtxt<'tcx>,
-                                   mir_map: Option<&MirMap<'tcx>>,
-                                   code: blocks::Code,
-                                   mode: PpFlowGraphMode,
-                                   mut out: W)
-                                   -> io::Result<()> {
-    let cfg = match code {
-        blocks::BlockCode(block) => cfg::CFG::new(tcx, &block),
-        blocks::FnLikeCode(fn_like) => cfg::CFG::new(tcx, &fn_like.body()),
-    };
-    let labelled_edges = mode != PpFlowGraphMode::UnlabelledEdges;
-    let lcfg = LabelledCFG {
-        ast_map: &tcx.map,
-        cfg: &cfg,
-        name: format!("node_{}", code.id()),
-        labelled_edges: labelled_edges,
-    };
-
-    match code {
-        _ if variants.is_empty() => {
-            let r = dot::render(&lcfg, &mut out);
-            return expand_err_details(r);
-        }
-        blocks::BlockCode(_) => {
-            tcx.sess.err("--pretty flowgraph with -Z flowgraph-print annotations requires \
-                          fn-like node id.");
-            return Ok(());
-        }
-        blocks::FnLikeCode(fn_like) => {
-            let (bccx, analysis_data) =
-                borrowck::build_borrowck_dataflow_data_for_fn(tcx,
-                                                              mir_map,
-                                                              fn_like.to_fn_parts(),
-                                                              &cfg);
-
-            let lcfg = borrowck_dot::DataflowLabeller {
-                inner: lcfg,
-                variants: variants,
-                borrowck_ctxt: &bccx,
-                analysis_data: &analysis_data,
-            };
-            let r = dot::render(&lcfg, &mut out);
-            return expand_err_details(r);
-        }
-    }
-
-    fn expand_err_details(r: io::Result<()>) -> io::Result<()> {
-        r.map_err(|ioerr| {
-            io::Error::new(io::ErrorKind::Other,
-                           &format!("graphviz::render failed: {}", ioerr)[..])
-        })
-    }
+    write_output(out, ofile);
 }