introduce `mir_keys()`

Each MIR key is a DefId that has MIR associated with it

18 files changed, 132 insertions, 141 deletions

diff --git a/src/librustc/dep_graph/dep_node.rs b/src/librustc/dep_graph/dep_node.rs
index 63a4e6196a2..66505d9a06b 100644
--- a/src/librustc/dep_graph/dep_node.rs
+++ b/src/librustc/dep_graph/dep_node.rs
@@ -76,6 +76,7 @@ pub enum DepNode<D: Clone + Debug> {
     BorrowCheck(D),
     RvalueCheck(D),
     Reachability,
+    MirKeys,
     LateLintCheck,
     TransCrateItem(D),
     TransInlinedItem(D),
@@ -202,6 +203,7 @@ impl<D: Clone + Debug> DepNode<D> {
             Variance => Some(Variance),
             PrivacyAccessLevels(k) => Some(PrivacyAccessLevels(k)),
             Reachability => Some(Reachability),
+            MirKeys => Some(MirKeys),
             LateLintCheck => Some(LateLintCheck),
             TransWriteMetadata => Some(TransWriteMetadata),
 
diff --git a/src/librustc/dep_graph/mod.rs b/src/librustc/dep_graph/mod.rs
index a9f0a44e420..6cb86a30400 100644
--- a/src/librustc/dep_graph/mod.rs
+++ b/src/librustc/dep_graph/mod.rs
@@ -28,6 +28,5 @@ pub use self::graph::WorkProduct;
 pub use self::query::DepGraphQuery;
 pub use self::safe::AssertDepGraphSafe;
 pub use self::safe::DepGraphSafe;
-pub use self::visit::visit_all_bodies_in_krate;
 pub use self::visit::visit_all_item_likes_in_krate;
 pub use self::raii::DepTask;
diff --git a/src/librustc/dep_graph/visit.rs b/src/librustc/dep_graph/visit.rs
index 93f6e3a83a0..bf3748659fe 100644
--- a/src/librustc/dep_graph/visit.rs
+++ b/src/librustc/dep_graph/visit.rs
@@ -75,15 +75,3 @@ pub fn visit_all_item_likes_in_krate<'a, 'tcx, V, F>(tcx: TyCtxt<'a, 'tcx, 'tcx>
     krate.visit_all_item_likes(&mut tracking_visitor)
 }
 
-pub fn visit_all_bodies_in_krate<'a, 'tcx, C>(tcx: TyCtxt<'a, 'tcx, 'tcx>, callback: C)
-    where C: Fn(/* body_owner */
-                DefId,
-                /* body id */
-                hir::BodyId)
-{
-    let krate = tcx.hir.krate();
-    for &body_id in &krate.body_ids {
-        let body_owner_def_id = tcx.hir.body_owner_def_id(body_id);
-        callback(body_owner_def_id, body_id);
-    }
-}
diff --git a/src/librustc/mir/transform.rs b/src/librustc/mir/transform.rs
index 4cbbb67c7e4..78a3c1919db 100644
--- a/src/librustc/mir/transform.rs
+++ b/src/librustc/mir/transform.rs
@@ -10,6 +10,7 @@
 
 use dep_graph::DepNode;
 use hir;
+use hir::def_id::LOCAL_CRATE;
 use hir::map::DefPathData;
 use mir::{Mir, Promoted};
 use ty::TyCtxt;
@@ -114,14 +115,9 @@ impl<'tcx, T: MirPass<'tcx>> MirMapPass<'tcx> for T {
                     tcx: TyCtxt<'a, 'tcx, 'tcx>,
                     hooks: &mut [Box<for<'s> MirPassHook<'s>>])
     {
-        let def_ids = tcx.maps.mir.borrow().keys();
-        for def_id in def_ids {
-            if !def_id.is_local() {
-                continue;
-            }
-
+        for &def_id in tcx.mir_keys(LOCAL_CRATE).iter() {
             let _task = tcx.dep_graph.in_task(DepNode::Mir(def_id));
-            let mir = &mut tcx.maps.mir.borrow()[&def_id].borrow_mut();
+            let mir = &mut tcx.mir(def_id).borrow_mut();
             tcx.dep_graph.write(DepNode::Mir(def_id));
 
             let id = tcx.hir.as_local_node_id(def_id).unwrap();
diff --git a/src/librustc/ty/maps.rs b/src/librustc/ty/maps.rs
index ef5dfab779c..ba475367bc6 100644
--- a/src/librustc/ty/maps.rs
+++ b/src/librustc/ty/maps.rs
@@ -20,7 +20,7 @@ use session::CompileResult;
 use ty::{self, CrateInherentImpls, Ty, TyCtxt};
 use ty::item_path;
 use ty::subst::Substs;
-use util::nodemap::NodeSet;
+use util::nodemap::{DefIdSet, NodeSet};
 
 use rustc_data_structures::indexed_vec::IndexVec;
 use std::cell::{RefCell, RefMut};
@@ -270,8 +270,13 @@ impl<'tcx> QueryDescription for queries::reachable_set<'tcx> {
 
 impl<'tcx> QueryDescription for queries::const_eval<'tcx> {
     fn describe(tcx: TyCtxt, (def_id, _): (DefId, &'tcx Substs<'tcx>)) -> String {
-        format!("const-evaluating `{}`",
-                tcx.item_path_str(def_id))
+        format!("const-evaluating `{}`", tcx.item_path_str(def_id))
+    }
+}
+
+impl<'tcx> QueryDescription for queries::mir_keys<'tcx> {
+    fn describe(_: TyCtxt, _: CrateNum) -> String {
+        format!("getting a list of all mir_keys")
     }
 }
 
@@ -546,6 +551,11 @@ define_maps! { <'tcx>
     /// (in the `RefCell` sense) to prevent accidental mutation.
     [pub] mir: Mir(DefId) -> &'tcx RefCell<mir::Mir<'tcx>>,
 
+    /// Set of all the def-ids in this crate that have MIR associated with
+    /// them. This includes all the body owners, but also things like struct
+    /// constructors.
+    [] mir_keys: mir_keys(CrateNum) -> Rc<DefIdSet>,
+
     /// Maps DefId's that have an associated Mir to the result
     /// of the MIR qualify_consts pass. The actual meaning of
     /// the value isn't known except to the pass itself.
@@ -644,3 +654,7 @@ fn typeck_item_bodies_dep_node(_: CrateNum) -> DepNode<DefId> {
 fn const_eval_dep_node((def_id, _): (DefId, &Substs)) -> DepNode<DefId> {
     DepNode::ConstEval(def_id)
 }
+
+fn mir_keys(_: CrateNum) -> DepNode<DefId> {
+    DepNode::MirKeys
+}
diff --git a/src/librustc/ty/mod.rs b/src/librustc/ty/mod.rs
index c07f41fb223..5e5dbdf20c0 100644
--- a/src/librustc/ty/mod.rs
+++ b/src/librustc/ty/mod.rs
@@ -2049,6 +2049,16 @@ impl<'a, 'gcx, 'tcx> TyCtxt<'a, 'gcx, 'tcx> {
         self.typeck_tables_of(self.hir.body_owner_def_id(body))
     }
 
+    /// Returns an iterator of the def-ids for all body-owners in this
+    /// crate. If you would prefer to iterate over the bodies
+    /// themselves, you can do `self.hir.krate().body_ids.iter()`.
+    pub fn body_owners(self) -> impl Iterator<Item = DefId> + 'a {
+        self.hir.krate()
+                .body_ids
+                .iter()
+                .map(move |&body_id| self.hir.body_owner_def_id(body_id))
+    }
+
     pub fn expr_span(self, id: NodeId) -> Span {
         match self.hir.find(id) {
             Some(hir_map::NodeExpr(e)) => {
@@ -2331,7 +2341,7 @@ impl<'a, 'gcx, 'tcx> TyCtxt<'a, 'gcx, 'tcx> {
     /// Given the DefId of an item, returns its MIR, borrowed immutably.
     /// Returns None if there is no MIR for the DefId
     pub fn maybe_item_mir(self, did: DefId) -> Option<Ref<'gcx, Mir<'gcx>>> {
-        if did.is_local() && !self.maps.mir.borrow().contains_key(&did) {
+        if did.is_local() && !self.mir_keys(LOCAL_CRATE).contains(&did) {
             return None;
         }
 
@@ -2541,17 +2551,6 @@ impl<'a, 'gcx, 'tcx> TyCtxt<'a, 'gcx, 'tcx> {
         dep_graph::visit_all_item_likes_in_krate(self.global_tcx(), dep_node_fn, visitor);
     }
 
-    /// Invokes `callback` for each body in the krate. This will
-    /// create a read edge from `DepNode::Krate` to the current task;
-    /// it is meant to be run in the context of some global task like
-    /// `BorrowckCrate`. The callback would then create a task like
-    /// `BorrowckBody(DefId)` to process each individual item.
-    pub fn visit_all_bodies_in_krate<C>(self, callback: C)
-        where C: Fn(/* body_owner */ DefId, /* body id */ hir::BodyId),
-    {
-        dep_graph::visit_all_bodies_in_krate(self.global_tcx(), callback)
-    }
-
     /// Looks up the span of `impl_did` if the impl is local; otherwise returns `Err`
     /// with the name of the crate containing the impl.
     pub fn span_of_impl(self, impl_did: DefId) -> Result<Span, Symbol> {
diff --git a/src/librustc_borrowck/borrowck/mod.rs b/src/librustc_borrowck/borrowck/mod.rs
index e5e5045bc29..1259816639a 100644
--- a/src/librustc_borrowck/borrowck/mod.rs
+++ b/src/librustc_borrowck/borrowck/mod.rs
@@ -63,9 +63,9 @@ pub struct LoanDataFlowOperator;
 pub type LoanDataFlow<'a, 'tcx> = DataFlowContext<'a, 'tcx, LoanDataFlowOperator>;
 
 pub fn check_crate<'a, 'tcx>(tcx: TyCtxt<'a, 'tcx, 'tcx>) {
-    tcx.visit_all_bodies_in_krate(|body_owner_def_id, _body_id| {
+    for body_owner_def_id in tcx.body_owners() {
         tcx.borrowck(body_owner_def_id);
-    });
+    }
 }
 
 pub fn provide(providers: &mut Providers) {
diff --git a/src/librustc_driver/pretty.rs b/src/librustc_driver/pretty.rs
index 18dc504ca8a..d40a2ab0b53 100644
--- a/src/librustc_driver/pretty.rs
+++ b/src/librustc_driver/pretty.rs
@@ -41,7 +41,6 @@ use graphviz as dot;
 use std::cell::Cell;
 use std::fs::File;
 use std::io::{self, Write};
-use std::iter;
 use std::option;
 use std::path::Path;
 use std::str::FromStr;
@@ -999,22 +998,14 @@ fn print_with_analysis<'tcx, 'a: 'tcx>(sess: &'a Session,
                 if let Some(nodeid) = nodeid {
                     let def_id = tcx.hir.local_def_id(nodeid);
                     match ppm {
-                        PpmMir => write_mir_pretty(tcx, iter::once(def_id), &mut out),
-                        PpmMirCFG => write_mir_graphviz(tcx, iter::once(def_id), &mut out),
+                        PpmMir => write_mir_pretty(tcx, Some(def_id), &mut out),
+                        PpmMirCFG => write_mir_graphviz(tcx, Some(def_id), &mut out),
                         _ => unreachable!(),
                     }?;
                 } else {
                     match ppm {
-                        PpmMir => {
-                            write_mir_pretty(tcx,
-                                             tcx.maps.mir.borrow().keys().into_iter(),
-                                             &mut out)
-                        }
-                        PpmMirCFG => {
-                            write_mir_graphviz(tcx,
-                                               tcx.maps.mir.borrow().keys().into_iter(),
-                                               &mut out)
-                        }
+                        PpmMir => write_mir_pretty(tcx, None, &mut out),
+                        PpmMirCFG => write_mir_graphviz(tcx, None, &mut out),
                         _ => unreachable!(),
                     }?;
                 }
diff --git a/src/librustc_metadata/encoder.rs b/src/librustc_metadata/encoder.rs
index 189b94a1b62..2160a56257b 100644
--- a/src/librustc_metadata/encoder.rs
+++ b/src/librustc_metadata/encoder.rs
@@ -621,7 +621,12 @@ impl<'a, 'b: 'a, 'tcx: 'b> EntryBuilder<'a, 'b, 'tcx> {
 
     fn encode_mir(&mut self, def_id: DefId) -> Option<Lazy<mir::Mir<'tcx>>> {
         debug!("EntryBuilder::encode_mir({:?})", def_id);
-        self.tcx.maps.mir.borrow().get(&def_id).map(|mir| self.lazy(&*mir.borrow()))
+        if self.tcx.mir_keys(LOCAL_CRATE).contains(&def_id) {
+            let mir = self.tcx.item_mir(def_id);
+            Some(self.lazy(&mir))
+        } else {
+            None
+        }
     }
 
     // Encodes the inherent implementations of a structure, enumeration, or trait.
diff --git a/src/librustc_mir/callgraph.rs b/src/librustc_mir/callgraph.rs
index 69416289d8e..961e202ba1c 100644
--- a/src/librustc_mir/callgraph.rs
+++ b/src/librustc_mir/callgraph.rs
@@ -12,7 +12,7 @@
 //!
 //! This only considers direct calls
 
-use rustc::hir::def_id::DefId;
+use rustc::hir::def_id::{DefId, LOCAL_CRATE};
 use rustc_data_structures::graph;
 
 use rustc::mir::*;
@@ -31,16 +31,12 @@ impl CallGraph {
     // FIXME: allow for construction of a callgraph that inspects
     // cross-crate MIRs if available.
     pub fn build<'a, 'tcx>(tcx: ty::TyCtxt<'a, 'tcx, 'tcx>) -> CallGraph {
-        let def_ids = tcx.maps.mir.borrow().keys();
-
         let mut callgraph = CallGraph {
             node_map: DefIdMap(),
             graph: graph::Graph::new()
         };
 
-        for def_id in def_ids {
-            if !def_id.is_local() { continue; }
-
+        for &def_id in tcx.mir_keys(LOCAL_CRATE).iter() {
             let idx = callgraph.add_node(def_id);
 
             let mut call_visitor = CallVisitor {
diff --git a/src/librustc_mir/mir_map.rs b/src/librustc_mir/mir_map.rs
index 1abae515ae6..353ca3bbd09 100644
--- a/src/librustc_mir/mir_map.rs
+++ b/src/librustc_mir/mir_map.rs
@@ -17,7 +17,7 @@
 //! - `#[rustc_mir(pretty="file.mir")]`
 
 use build;
-use rustc::hir::def_id::DefId;
+use rustc::hir::def_id::{CrateNum, DefId, LOCAL_CRATE};
 use rustc::dep_graph::DepNode;
 use rustc::mir::Mir;
 use rustc::mir::transform::MirSource;
@@ -32,50 +32,67 @@ use rustc::ty::maps::Providers;
 use rustc::ty::subst::Substs;
 use rustc::hir;
 use rustc::hir::intravisit::{self, Visitor, NestedVisitorMap};
+use rustc::util::nodemap::DefIdSet;
 use syntax::abi::Abi;
 use syntax::ast;
 use syntax_pos::Span;
 
 use std::cell::RefCell;
 use std::mem;
+use std::rc::Rc;
 
 pub fn build_mir_for_crate<'a, 'tcx>(tcx: TyCtxt<'a, 'tcx, 'tcx>) {
     tcx.dep_graph.with_task(DepNode::MirKrate, tcx, (), build_mir_for_crate_task);
 
     fn build_mir_for_crate_task<'a, 'tcx>(tcx: TyCtxt<'a, 'tcx, 'tcx>, (): ()) {
-        tcx.visit_all_bodies_in_krate(|body_owner_def_id, _body_id| {
+        for &body_owner_def_id in tcx.mir_keys(LOCAL_CRATE).iter() {
             tcx.item_mir(body_owner_def_id);
-        });
-
-        // Tuple struct/variant constructors don't have a BodyId, so we need
-        // to build them separately.
-        struct GatherCtors<'a, 'tcx: 'a> {
-            tcx: TyCtxt<'a, 'tcx, 'tcx>
-        }
-        impl<'a, 'tcx> Visitor<'tcx> for GatherCtors<'a, 'tcx> {
-            fn visit_variant_data(&mut self,
-                                  v: &'tcx hir::VariantData,
-                                  _: ast::Name,
-                                  _: &'tcx hir::Generics,
-                                  _: ast::NodeId,
-                                  _: Span) {
-                if let hir::VariantData::Tuple(_, node_id) = *v {
-                    self.tcx.item_mir(self.tcx.hir.local_def_id(node_id));
-                }
-                intravisit::walk_struct_def(self, v)
-            }
-            fn nested_visit_map<'b>(&'b mut self) -> NestedVisitorMap<'b, 'tcx> {
-                NestedVisitorMap::None
-            }
         }
-        tcx.hir.krate().visit_all_item_likes(&mut GatherCtors {
-            tcx: tcx
-        }.as_deep_visitor());
     }
 }
 
 pub fn provide(providers: &mut Providers) {
     providers.mir = build_mir;
+    providers.mir_keys = mir_keys;
+}
+
+fn mir_keys<'a, 'tcx>(tcx: TyCtxt<'a, 'tcx, 'tcx>, krate: CrateNum)
+                      -> Rc<DefIdSet> {
+    assert_eq!(krate, LOCAL_CRATE);
+
+    let mut set = DefIdSet();
+
+    // All body-owners have MIR associated with them.
+    set.extend(tcx.body_owners());
+
+    // Additionally, tuple struct/variant constructors have MIR, but
+    // they don't have a BodyId, so we need to build them separately.
+    struct GatherCtors<'a, 'tcx: 'a> {
+        tcx: TyCtxt<'a, 'tcx, 'tcx>,
+        set: &'a mut DefIdSet,
+    }
+    impl<'a, 'tcx> Visitor<'tcx> for GatherCtors<'a, 'tcx> {
+        fn visit_variant_data(&mut self,
+                              v: &'tcx hir::VariantData,
+                              _: ast::Name,
+                              _: &'tcx hir::Generics,
+                              _: ast::NodeId,
+                              _: Span) {
+            if let hir::VariantData::Tuple(_, node_id) = *v {
+                self.set.insert(self.tcx.hir.local_def_id(node_id));
+            }
+            intravisit::walk_struct_def(self, v)
+        }
+        fn nested_visit_map<'b>(&'b mut self) -> NestedVisitorMap<'b, 'tcx> {
+            NestedVisitorMap::None
+        }
+    }
+    tcx.hir.krate().visit_all_item_likes(&mut GatherCtors {
+        tcx: tcx,
+        set: &mut set,
+    }.as_deep_visitor());
+
+    Rc::new(set)
 }
 
 fn build_mir<'a, 'tcx>(tcx: TyCtxt<'a, 'tcx, 'tcx>, def_id: DefId)
diff --git a/src/librustc_mir/transform/dump_mir.rs b/src/librustc_mir/transform/dump_mir.rs
index 5b3113f962b..6895facd6af 100644
--- a/src/librustc_mir/transform/dump_mir.rs
+++ b/src/librustc_mir/transform/dump_mir.rs
@@ -81,6 +81,6 @@ pub fn emit_mir<'a, 'tcx>(
 {
     let path = outputs.path(OutputType::Mir);
     let mut f = File::create(&path)?;
-    mir_util::write_mir_pretty(tcx, tcx.maps.mir.borrow().keys().into_iter(), &mut f)?;
+    mir_util::write_mir_pretty(tcx, None, &mut f)?;
     Ok(())
 }
diff --git a/src/librustc_mir/transform/inline.rs b/src/librustc_mir/transform/inline.rs
index 45bdff9195c..fcd5f970882 100644
--- a/src/librustc_mir/transform/inline.rs
+++ b/src/librustc_mir/transform/inline.rs
@@ -10,7 +10,7 @@
 
 //! Inlining pass for MIR functions
 
-use rustc::hir::def_id::DefId;
+use rustc::hir::def_id::{DefId, LOCAL_CRATE};
 
 use rustc_data_structures::bitvec::BitVector;
 use rustc_data_structures::indexed_vec::{Idx, IndexVec};
@@ -58,24 +58,15 @@ impl<'tcx> MirMapPass<'tcx> for Inline {
             tcx: tcx,
         };
 
-        let def_ids = tcx.maps.mir.borrow().keys();
-        for &def_id in &def_ids {
-            if !def_id.is_local() { continue; }
-
+        for &def_id in tcx.mir_keys(LOCAL_CRATE).iter() {
             let _task = tcx.dep_graph.in_task(DepNode::Mir(def_id));
-            let mut mir = if let Some(mir) = tcx.maps.mir.borrow().get(&def_id) {
-                mir.borrow_mut()
-            } else {
-                continue;
-            };
-
-            tcx.dep_graph.write(DepNode::Mir(def_id));
+            let mir = &tcx.item_mir(def_id);
 
             let id = tcx.hir.as_local_node_id(def_id).unwrap();
             let src = MirSource::from_node(tcx, id);
 
             for hook in &mut *hooks {
-                hook.on_mir_pass(tcx, src, &mut mir, self, false);
+                hook.on_mir_pass(tcx, src, mir, self, false);
             }
         }
 
@@ -83,18 +74,15 @@ impl<'tcx> MirMapPass<'tcx> for Inline {
             inliner.inline_scc(&callgraph, &scc);
         }
 
-        for def_id in def_ids {
-            if !def_id.is_local() { continue; }
-
+        for &def_id in tcx.mir_keys(LOCAL_CRATE).iter() {
             let _task = tcx.dep_graph.in_task(DepNode::Mir(def_id));
-            let mut mir = tcx.maps.mir.borrow()[&def_id].borrow_mut();
-            tcx.dep_graph.write(DepNode::Mir(def_id));
+            let mir = &tcx.item_mir(def_id);
 
             let id = tcx.hir.as_local_node_id(def_id).unwrap();
             let src = MirSource::from_node(tcx, id);
 
             for hook in &mut *hooks {
-                hook.on_mir_pass(tcx, src, &mut mir, self, true);
+                hook.on_mir_pass(tcx, src, mir, self, true);
             }
         }
     }
@@ -200,11 +188,7 @@ impl<'a, 'tcx> Inliner<'a, 'tcx> {
 
                 };
 
-                let mut caller_mir = {
-                    let map = self.tcx.maps.mir.borrow();
-                    let mir = map.get(&callsite.caller).unwrap();
-                    mir.borrow_mut()
-                };
+                let mut caller_mir = self.tcx.mir(callsite.caller).borrow_mut();
 
                 let start = caller_mir.basic_blocks().len();
 
@@ -256,11 +240,7 @@ impl<'a, 'tcx> Inliner<'a, 'tcx> {
             let _task = self.tcx.dep_graph.in_task(DepNode::Mir(def_id));
             self.tcx.dep_graph.write(DepNode::Mir(def_id));
 
-            let mut caller_mir = {
-                let map = self.tcx.maps.mir.borrow();
-                let mir = map.get(&def_id).unwrap();
-                mir.borrow_mut()
-            };
+            let mut caller_mir = self.tcx.mir(def_id).borrow_mut();
 
             debug!("Running simplify cfg on {:?}", def_id);
             CfgSimplifier::new(&mut caller_mir).simplify();
diff --git a/src/librustc_mir/transform/qualify_consts.rs b/src/librustc_mir/transform/qualify_consts.rs
index afb775aa01e..f5662295296 100644
--- a/src/librustc_mir/transform/qualify_consts.rs
+++ b/src/librustc_mir/transform/qualify_consts.rs
@@ -19,7 +19,7 @@ use rustc_data_structures::indexed_vec::{IndexVec, Idx};
 use rustc::dep_graph::DepNode;
 use rustc::hir;
 use rustc::hir::map as hir_map;
-use rustc::hir::def_id::DefId;
+use rustc::hir::def_id::{DefId, LOCAL_CRATE};
 use rustc::hir::map::blocks::FnLikeNode;
 use rustc::traits::{self, Reveal};
 use rustc::ty::{self, TyCtxt, Ty, TypeFoldable};
@@ -946,12 +946,7 @@ impl<'tcx> MirMapPass<'tcx> for QualifyAndPromoteConstants {
                     tcx: TyCtxt<'a, 'tcx, 'tcx>,
                     hooks: &mut [Box<for<'s> MirPassHook<'s>>])
     {
-        let def_ids = tcx.maps.mir.borrow().keys();
-        for def_id in def_ids {
-            if !def_id.is_local() {
-                continue;
-            }
-
+        for &def_id in tcx.mir_keys(LOCAL_CRATE).iter() {
             let _task = tcx.dep_graph.in_task(DepNode::Mir(def_id));
             let id = tcx.hir.as_local_node_id(def_id).unwrap();
             let src = MirSource::from_node(tcx, id);
@@ -961,7 +956,7 @@ impl<'tcx> MirMapPass<'tcx> for QualifyAndPromoteConstants {
                 continue;
             }
 
-            let mir = &mut tcx.maps.mir.borrow()[&def_id].borrow_mut();
+            let mir = &mut tcx.mir(def_id).borrow_mut();
             tcx.dep_graph.write(DepNode::Mir(def_id));
 
             for hook in &mut *hooks {
diff --git a/src/librustc_mir/util/graphviz.rs b/src/librustc_mir/util/graphviz.rs
index 91600b947c6..d3d3977a934 100644
--- a/src/librustc_mir/util/graphviz.rs
+++ b/src/librustc_mir/util/graphviz.rs
@@ -18,14 +18,16 @@ use syntax::ast::NodeId;
 
 use rustc_data_structures::indexed_vec::Idx;
 
+use super::pretty::dump_mir_def_ids;
+
 /// Write a graphviz DOT graph of a list of MIRs.
-pub fn write_mir_graphviz<'a, 'b, 'tcx, W, I>(tcx: TyCtxt<'b, 'tcx, 'tcx>,
-                                              iter: I,
-                                              w: &mut W)
-                                              -> io::Result<()>
-    where W: Write, I: Iterator<Item=DefId>
+pub fn write_mir_graphviz<'a, 'tcx, W>(tcx: TyCtxt<'a, 'tcx, 'tcx>,
+                                       single: Option<DefId>,
+                                       w: &mut W)
+                                       -> io::Result<()>
+    where W: Write
 {
-    for def_id in iter {
+    for def_id in dump_mir_def_ids(tcx, single) {
         let nodeid = tcx.hir.as_local_node_id(def_id).unwrap();
         let mir = &tcx.item_mir(def_id);
 
diff --git a/src/librustc_mir/util/pretty.rs b/src/librustc_mir/util/pretty.rs
index b202e149510..c03c9c907a3 100644
--- a/src/librustc_mir/util/pretty.rs
+++ b/src/librustc_mir/util/pretty.rs
@@ -9,7 +9,7 @@
 // except according to those terms.
 
 use rustc::hir;
-use rustc::hir::def_id::DefId;
+use rustc::hir::def_id::{DefId, LOCAL_CRATE};
 use rustc::mir::*;
 use rustc::mir::transform::MirSource;
 use rustc::ty::TyCtxt;
@@ -85,17 +85,16 @@ pub fn dump_mir<'a, 'tcx>(tcx: TyCtxt<'a, 'tcx, 'tcx>,
 }
 
 /// Write out a human-readable textual representation for the given MIR.
-pub fn write_mir_pretty<'a, 'b, 'tcx, I>(tcx: TyCtxt<'b, 'tcx, 'tcx>,
-                                         iter: I,
-                                         w: &mut Write)
-                                         -> io::Result<()>
-    where I: Iterator<Item=DefId>, 'tcx: 'a
+pub fn write_mir_pretty<'a, 'tcx>(tcx: TyCtxt<'a, 'tcx, 'tcx>,
+                                  single: Option<DefId>,
+                                  w: &mut Write)
+                                  -> io::Result<()>
 {
     writeln!(w, "// WARNING: This output format is intended for human consumers only")?;
     writeln!(w, "// and is subject to change without notice. Knock yourself out.")?;
 
     let mut first = true;
-    for def_id in iter.filter(DefId::is_local) {
+    for def_id in dump_mir_def_ids(tcx, single) {
         let mir = &tcx.item_mir(def_id);
 
         if first {
@@ -312,3 +311,11 @@ fn write_temp_decls(mir: &Mir, w: &mut Write) -> io::Result<()> {
 
     Ok(())
 }
+
+pub fn dump_mir_def_ids(tcx: TyCtxt, single: Option<DefId>) -> Vec<DefId> {
+    if let Some(i) = single {
+        vec![i]
+    } else {
+        tcx.mir_keys(LOCAL_CRATE).iter().cloned().collect()
+    }
+}
diff --git a/src/librustc_passes/mir_stats.rs b/src/librustc_passes/mir_stats.rs
index ce02cb0e836..39b3b709af7 100644
--- a/src/librustc_passes/mir_stats.rs
+++ b/src/librustc_passes/mir_stats.rs
@@ -13,6 +13,7 @@
 // completely accurate (some things might be counted twice, others missed).
 
 use rustc_const_math::{ConstUsize};
+use rustc::hir::def_id::LOCAL_CRATE;
 use rustc::middle::const_val::{ConstVal};
 use rustc::mir::{AggregateKind, AssertMessage, BasicBlock, BasicBlockData};
 use rustc::mir::{Constant, Literal, Location, LocalDecl};
@@ -44,10 +45,9 @@ pub fn print_mir_stats<'tcx, 'a>(tcx: TyCtxt<'a, 'tcx, 'tcx>, title: &str) {
     // For debugging instrumentation like this, we don't need to worry
     // about maintaining the dep graph.
     let _ignore = tcx.dep_graph.in_ignore();
-    let mir_map = tcx.maps.mir.borrow();
-    for def_id in mir_map.keys() {
-        let mir = mir_map.get(&def_id).unwrap();
-        collector.visit_mir(&mir.borrow());
+    for &def_id in tcx.mir_keys(LOCAL_CRATE).iter() {
+        let mir = tcx.item_mir(def_id);
+        collector.visit_mir(&mir);
     }
     collector.print(title);
 }
diff --git a/src/librustc_typeck/check/mod.rs b/src/librustc_typeck/check/mod.rs
index c401ed428e4..11095e70f62 100644
--- a/src/librustc_typeck/check/mod.rs
+++ b/src/librustc_typeck/check/mod.rs
@@ -640,9 +640,9 @@ pub fn check_item_bodies<'a, 'tcx>(tcx: TyCtxt<'a, 'tcx, 'tcx>) -> CompileResult
 fn typeck_item_bodies<'a, 'tcx>(tcx: TyCtxt<'a, 'tcx, 'tcx>, crate_num: CrateNum) -> CompileResult {
     debug_assert!(crate_num == LOCAL_CRATE);
     tcx.sess.track_errors(|| {
-        tcx.visit_all_bodies_in_krate(|body_owner_def_id, _body_id| {
+        for body_owner_def_id in tcx.body_owners() {
             tcx.typeck_tables_of(body_owner_def_id);
-        });
+        }
     })
 }