Add support for macro expansion in rustdoc source code pages

3 files changed, 116 insertions, 7 deletions

diff --git a/src/librustdoc/html/render/context.rs b/src/librustdoc/html/render/context.rs
index e4fca09d64f..33dec1c24e6 100644
--- a/src/librustdoc/html/render/context.rs
+++ b/src/librustdoc/html/render/context.rs
@@ -12,7 +12,7 @@ use rustc_hir::def_id::{DefIdMap, LOCAL_CRATE};
 use rustc_middle::ty::TyCtxt;
 use rustc_session::Session;
 use rustc_span::edition::Edition;
-use rustc_span::{FileName, Symbol, sym};
+use rustc_span::{BytePos, FileName, Symbol, sym};
 use tracing::info;
 
 use super::print_item::{full_path, print_item, print_item_path};
@@ -29,6 +29,7 @@ use crate::formats::cache::Cache;
 use crate::formats::item_type::ItemType;
 use crate::html::escape::Escape;
 use crate::html::markdown::{self, ErrorCodes, IdMap, plain_text_summary};
+use crate::html::render::ExpandedCode;
 use crate::html::render::write_shared::write_shared;
 use crate::html::url_parts_builder::UrlPartsBuilder;
 use crate::html::{layout, sources, static_files};
@@ -139,6 +140,7 @@ pub(crate) struct SharedContext<'tcx> {
     /// Correspondence map used to link types used in the source code pages to allow to click on
     /// links to jump to the type's definition.
     pub(crate) span_correspondence_map: FxHashMap<rustc_span::Span, LinkFromSrc>,
+    pub(crate) expanded_codes: FxHashMap<BytePos, Vec<ExpandedCode>>,
     /// The [`Cache`] used during rendering.
     pub(crate) cache: Cache,
     pub(crate) call_locations: AllCallLocations,
@@ -548,7 +550,7 @@ impl<'tcx> FormatRenderer<'tcx> for Context<'tcx> {
             }
         }
 
-        let (local_sources, matches) = collect_spans_and_sources(
+        let (local_sources, matches, expanded_codes) = collect_spans_and_sources(
             tcx,
             &krate,
             &src_root,
@@ -579,6 +581,7 @@ impl<'tcx> FormatRenderer<'tcx> for Context<'tcx> {
             cache,
             call_locations,
             should_merge: options.should_merge,
+            expanded_codes,
         };
 
         let dst = output;
diff --git a/src/librustdoc/html/render/mod.rs b/src/librustdoc/html/render/mod.rs
index 8d7f0577506..0986b9ca47b 100644
--- a/src/librustdoc/html/render/mod.rs
+++ b/src/librustdoc/html/render/mod.rs
@@ -63,7 +63,7 @@ use serde::{Serialize, Serializer};
 use tracing::{debug, info};
 
 pub(crate) use self::context::*;
-pub(crate) use self::span_map::{LinkFromSrc, collect_spans_and_sources};
+pub(crate) use self::span_map::{ExpandedCode, LinkFromSrc, collect_spans_and_sources};
 pub(crate) use self::write_shared::*;
 use crate::clean::{self, ItemId, RenderedLink};
 use crate::display::{Joined as _, MaybeDisplay as _};
diff --git a/src/librustdoc/html/render/span_map.rs b/src/librustdoc/html/render/span_map.rs
index 846d3ad310c..842da838049 100644
--- a/src/librustdoc/html/render/span_map.rs
+++ b/src/librustdoc/html/render/span_map.rs
@@ -30,12 +30,26 @@ pub(crate) enum LinkFromSrc {
     Doc(DefId),
 }
 
+/// Contains information about macro expansion in the source code pages.
+#[derive(Debug)]
+pub(crate) struct ExpandedCode {
+    /// The line where the macro expansion starts.
+    pub(crate) start_line: u32,
+    /// The line where the macro expansion ends.
+    pub(crate) end_line: u32,
+    /// The source code of the expanded macro.
+    pub(crate) code: String,
+    /// The span of macro callsite.
+    pub(crate) span: Span,
+}
+
 /// This function will do at most two things:
 ///
 /// 1. Generate a `span` correspondence map which links an item `span` to its definition `span`.
 /// 2. Collect the source code files.
 ///
-/// It returns the `krate`, the source code files and the `span` correspondence map.
+/// It returns the source code files, the `span` correspondence map and the expanded macros
+/// correspondence map.
 ///
 /// Note about the `span` correspondence map: the keys are actually `(lo, hi)` of `span`s. We don't
 /// need the `span` context later on, only their position, so instead of keeping a whole `Span`, we
@@ -46,17 +60,23 @@ pub(crate) fn collect_spans_and_sources(
     src_root: &Path,
     include_sources: bool,
     generate_link_to_definition: bool,
-) -> (FxIndexMap<PathBuf, String>, FxHashMap<Span, LinkFromSrc>) {
+) -> (
+    FxIndexMap<PathBuf, String>,
+    FxHashMap<Span, LinkFromSrc>,
+    FxHashMap<BytePos, Vec<ExpandedCode>>,
+) {
     if include_sources {
         let mut visitor = SpanMapVisitor { tcx, matches: FxHashMap::default() };
+        let mut expanded_visitor = ExpandedCodeVisitor { tcx, expanded_codes: Vec::new() };
 
+        tcx.hir_walk_toplevel_module(&mut expanded_visitor);
         if generate_link_to_definition {
             tcx.hir_walk_toplevel_module(&mut visitor);
         }
         let sources = sources::collect_local_sources(tcx, src_root, krate);
-        (sources, visitor.matches)
+        (sources, visitor.matches, expanded_visitor.compute_expanded())
     } else {
-        (Default::default(), Default::default())
+        (Default::default(), Default::default(), Default::default())
     }
 }
 
@@ -292,3 +312,89 @@ impl<'tcx> Visitor<'tcx> for SpanMapVisitor<'tcx> {
         intravisit::walk_item(self, item);
     }
 }
+
+/// Contains temporary information of macro expanded code.
+///
+/// As we go through the HIR visitor, if any span overlaps with another, they will
+/// both be merged.
+struct ExpandedCodeInfo {
+    /// Callsite of the macro.
+    span: Span,
+    /// Expanded macro source code.
+    code: String,
+}
+
+/// HIR visitor which retrieves expanded macro.
+///
+/// Once done, the `expanded_codes` will be transformed into a vec of [`ExpandedCode`]
+/// which contains more information needed when running the source code highlighter.
+pub struct ExpandedCodeVisitor<'tcx> {
+    tcx: TyCtxt<'tcx>,
+    expanded_codes: Vec<ExpandedCodeInfo>,
+}
+
+impl<'tcx> ExpandedCodeVisitor<'tcx> {
+    fn handle_new_span<F: Fn(TyCtxt<'_>) -> String>(&mut self, new_span: Span, f: F) {
+        if new_span.is_dummy() || !new_span.from_expansion() {
+            return;
+        }
+        let new_span = new_span.source_callsite();
+        if let Some(index) =
+            self.expanded_codes.iter().position(|info| info.span.overlaps(new_span))
+        {
+            if !self.expanded_codes[index].span.contains(new_span) {
+                // We replace the item.
+                let info = &mut self.expanded_codes[index];
+                info.span = new_span;
+                info.code = f(self.tcx);
+            }
+        } else {
+            // We add a new item.
+            self.expanded_codes.push(ExpandedCodeInfo {
+                span: new_span,
+                code: f(self.tcx),
+            });
+        }
+    }
+
+    fn compute_expanded(mut self) -> FxHashMap<BytePos, Vec<ExpandedCode>> {
+        self.expanded_codes.sort_unstable_by(|item1, item2| item1.span.cmp(&item2.span));
+        let source_map = self.tcx.sess.source_map();
+        let mut expanded: FxHashMap<BytePos, Vec<ExpandedCode>> = FxHashMap::default();
+        for ExpandedCodeInfo { span, code } in self.expanded_codes {
+            if let Ok(lines) = source_map.span_to_lines(span)
+                && !lines.lines.is_empty()
+            {
+                let mut out = String::new();
+                super::highlight::write_code(&mut out, &code, None, None, None);
+                let first = lines.lines.first().unwrap();
+                let end = lines.lines.last().unwrap();
+                expanded.entry(lines.file.start_pos).or_default().push(ExpandedCode {
+                    start_line: first.line_index as u32 + 1,
+                    end_line: end.line_index as u32 + 1,
+                    code: out,
+                    span,
+                });
+            }
+        }
+        expanded
+    }
+}
+
+impl<'tcx> Visitor<'tcx> for ExpandedCodeVisitor<'tcx> {
+    type NestedFilter = nested_filter::All;
+
+    fn maybe_tcx(&mut self) -> Self::MaybeTyCtxt {
+        self.tcx
+    }
+
+    fn visit_expr(&mut self, expr: &'tcx rustc_hir::Expr<'tcx>) {
+        self.handle_new_span(expr.span, |tcx| rustc_hir_pretty::expr_to_string(&tcx, expr));
+        intravisit::walk_expr(self, expr);
+    }
+
+    fn visit_item(&mut self, item: &'tcx rustc_hir::Item<'tcx>) {
+        self.handle_new_span(item.span, |tcx| rustc_hir_pretty::item_to_string(&tcx, item));
+        intravisit::walk_item(self, item);
+    }
+}