rustdoc: pretty-print nested bodies in inlined constants.

6 files changed, 79 insertions, 5 deletions

diff --git a/src/librustc/hir/mod.rs b/src/librustc/hir/mod.rs
index f76c29ea1c9..9149da459c2 100644
--- a/src/librustc/hir/mod.rs
+++ b/src/librustc/hir/mod.rs
@@ -856,7 +856,7 @@ pub enum UnsafeSource {
     UserProvided,
 }
 
-#[derive(Copy, Clone, PartialEq, Eq, RustcEncodable, RustcDecodable, Hash, Debug)]
+#[derive(Copy, Clone, PartialEq, Eq, PartialOrd, Ord, RustcEncodable, RustcDecodable, Hash, Debug)]
 pub struct BodyId {
     pub node_id: NodeId,
 }
diff --git a/src/librustc/middle/cstore.rs b/src/librustc/middle/cstore.rs
index ffd1ad8a0fa..7151e5226ca 100644
--- a/src/librustc/middle/cstore.rs
+++ b/src/librustc/middle/cstore.rs
@@ -33,6 +33,8 @@ use mir::Mir;
 use session::Session;
 use session::search_paths::PathKind;
 use util::nodemap::{NodeSet, DefIdMap};
+
+use std::collections::BTreeMap;
 use std::path::PathBuf;
 use std::rc::Rc;
 use syntax::ast;
@@ -250,6 +252,7 @@ pub trait CrateStore<'tcx> {
     // misc. metadata
     fn maybe_get_item_body<'a>(&'tcx self, tcx: TyCtxt<'a, 'tcx, 'tcx>, def: DefId)
                                -> Option<&'tcx hir::Body>;
+    fn item_body_nested_bodies(&self, def: DefId) -> BTreeMap<hir::BodyId, hir::Body>;
     fn const_is_rvalue_promotable_to_static(&self, def: DefId) -> bool;
 
     fn get_item_mir<'a>(&self, tcx: TyCtxt<'a, 'tcx, 'tcx>, def: DefId) -> Mir<'tcx>;
@@ -421,6 +424,9 @@ impl<'tcx> CrateStore<'tcx> for DummyCrateStore {
                                -> Option<&'tcx hir::Body> {
         bug!("maybe_get_item_body")
     }
+    fn item_body_nested_bodies(&self, def: DefId) -> BTreeMap<hir::BodyId, hir::Body> {
+        bug!("item_body_nested_bodies")
+    }
     fn const_is_rvalue_promotable_to_static(&self, def: DefId) -> bool {
         bug!("const_is_rvalue_promotable_to_static")
     }
diff --git a/src/librustc_metadata/astencode.rs b/src/librustc_metadata/astencode.rs
index 54824fc51d6..b27b164bd47 100644
--- a/src/librustc_metadata/astencode.rs
+++ b/src/librustc_metadata/astencode.rs
@@ -30,6 +30,7 @@ pub struct Ast<'tcx> {
     id_range: IdRange,
     body: Lazy<hir::Body>,
     side_tables: LazySeq<(ast::NodeId, TableEntry<'tcx>)>,
+    pub nested_bodies: LazySeq<hir::Body>,
     pub rvalue_promotable_to_static: bool,
 }
 
@@ -61,6 +62,16 @@ impl<'a, 'tcx> EncodeContext<'a, 'tcx> {
             visitor.count
         };
 
+        let nested_pos = self.position();
+        let nested_count = {
+            let mut visitor = NestedBodyEncodingVisitor {
+                ecx: self,
+                count: 0,
+            };
+            visitor.visit_body(body);
+            visitor.count
+        };
+
         let rvalue_promotable_to_static =
             self.tcx.rvalue_promotable_to_static.borrow()[&body.value.id];
 
@@ -68,6 +79,7 @@ impl<'a, 'tcx> EncodeContext<'a, 'tcx> {
             id_range: id_visitor.result(),
             body: Lazy::with_position(body_pos),
             side_tables: LazySeq::with_position_and_length(tables_pos, tables_count),
+            nested_bodies: LazySeq::with_position_and_length(nested_pos, nested_count),
             rvalue_promotable_to_static: rvalue_promotable_to_static
         })
     }
@@ -102,6 +114,25 @@ impl<'a, 'b, 'tcx> Visitor<'tcx> for SideTableEncodingIdVisitor<'a, 'b, 'tcx> {
     }
 }
 
+struct NestedBodyEncodingVisitor<'a, 'b: 'a, 'tcx: 'b> {
+    ecx: &'a mut EncodeContext<'b, 'tcx>,
+    count: usize,
+}
+
+impl<'a, 'b, 'tcx> Visitor<'tcx> for NestedBodyEncodingVisitor<'a, 'b, 'tcx> {
+    fn nested_visit_map<'this>(&'this mut self) -> NestedVisitorMap<'this, 'tcx> {
+        NestedVisitorMap::None
+    }
+
+    fn visit_nested_body(&mut self, body: hir::BodyId) {
+        let body = self.ecx.tcx.map.body(body);
+        body.encode(self.ecx).unwrap();
+        self.count += 1;
+
+        self.visit_body(body);
+    }
+}
+
 /// Decodes an item's body from its AST in the cdata's metadata and adds it to the
 /// ast-map.
 pub fn decode_body<'a, 'tcx>(cdata: &CrateMetadata,
diff --git a/src/librustc_metadata/cstore_impl.rs b/src/librustc_metadata/cstore_impl.rs
index a5173e00c72..64513fa41b2 100644
--- a/src/librustc_metadata/cstore_impl.rs
+++ b/src/librustc_metadata/cstore_impl.rs
@@ -36,6 +36,8 @@ use rustc::hir::svh::Svh;
 use rustc_back::target::Target;
 use rustc::hir;
 
+use std::collections::BTreeMap;
+
 impl<'tcx> CrateStore<'tcx> for cstore::CStore {
     fn describe_def(&self, def: DefId) -> Option<Def> {
         self.dep_graph.read(DepNode::MetaData(def));
@@ -455,6 +457,11 @@ impl<'tcx> CrateStore<'tcx> for cstore::CStore {
         inlined
     }
 
+    fn item_body_nested_bodies(&self, def: DefId) -> BTreeMap<hir::BodyId, hir::Body> {
+        self.dep_graph.read(DepNode::MetaData(def));
+        self.get_crate_data(def.krate).item_body_nested_bodies(def.index)
+    }
+
     fn const_is_rvalue_promotable_to_static(&self, def: DefId) -> bool {
         self.dep_graph.read(DepNode::MetaData(def));
         self.get_crate_data(def.krate).const_is_rvalue_promotable_to_static(def.index)
diff --git a/src/librustc_metadata/decoder.rs b/src/librustc_metadata/decoder.rs
index 81f63825919..c27e06c5022 100644
--- a/src/librustc_metadata/decoder.rs
+++ b/src/librustc_metadata/decoder.rs
@@ -32,6 +32,7 @@ use rustc::mir::Mir;
 
 use std::borrow::Cow;
 use std::cell::Ref;
+use std::collections::BTreeMap;
 use std::io;
 use std::mem;
 use std::str;
@@ -829,6 +830,12 @@ impl<'a, 'tcx> CrateMetadata {
         })
     }
 
+    pub fn item_body_nested_bodies(&self, id: DefIndex) -> BTreeMap<hir::BodyId, hir::Body> {
+        self.entry(id).ast.into_iter().flat_map(|ast| {
+            ast.decode(self).nested_bodies.decode(self).map(|body| (body.id(), body))
+        }).collect()
+    }
+
     pub fn const_is_rvalue_promotable_to_static(&self, id: DefIndex) -> bool {
         self.entry(id).ast.expect("const item missing `ast`")
             .decode(self).rvalue_promotable_to_static
diff --git a/src/librustdoc/clean/inline.rs b/src/librustdoc/clean/inline.rs
index a1c121ca424..cba5e1ba6f3 100644
--- a/src/librustdoc/clean/inline.rs
+++ b/src/librustdoc/clean/inline.rs
@@ -10,6 +10,8 @@
 
 //! Support for inlining external documentation into the current AST.
 
+use std::collections::BTreeMap;
+use std::io;
 use std::iter::once;
 
 use syntax::ast;
@@ -342,8 +344,7 @@ pub fn build_impl(cx: &DocContext, did: DefId, ret: &mut Vec<clean::Item>) {
         match item.kind {
             ty::AssociatedKind::Const => {
                 let default = if item.defaultness.has_value() {
-                    Some(hir::print::to_string(&cx.tcx.map, |s| s.print_expr(
-                        &tcx.sess.cstore.maybe_get_item_body(tcx, item.def_id).unwrap().value)))
+                    Some(print_inlined_const(cx, item.def_id))
                 } else {
                     None
                 };
@@ -473,11 +474,33 @@ fn build_module(cx: &DocContext, did: DefId) -> clean::Module {
     }
 }
 
+struct InlinedConst {
+    nested_bodies: BTreeMap<hir::BodyId, hir::Body>
+}
+
+impl hir::print::PpAnn for InlinedConst {
+    fn nested(&self, state: &mut hir::print::State, nested: hir::print::Nested)
+              -> io::Result<()> {
+        if let hir::print::Nested::Body(body) = nested {
+            state.print_expr(&self.nested_bodies[&body].value)
+        } else {
+            Ok(())
+        }
+    }
+}
+
+fn print_inlined_const(cx: &DocContext, did: DefId) -> String {
+    let body = cx.tcx.sess.cstore.maybe_get_item_body(cx.tcx, did).unwrap();
+    let inlined = InlinedConst {
+        nested_bodies: cx.tcx.sess.cstore.item_body_nested_bodies(did)
+    };
+    hir::print::to_string(&inlined, |s| s.print_expr(&body.value))
+}
+
 fn build_const(cx: &DocContext, did: DefId) -> clean::Constant {
     clean::Constant {
         type_: cx.tcx.item_type(did).clean(cx),
-        expr: hir::print::to_string(&cx.tcx.map, |s| s.print_expr(
-            &cx.tcx.sess.cstore.maybe_get_item_body(cx.tcx, did).unwrap().value))
+        expr: print_inlined_const(cx, did)
     }
 }