Do macro expansion at AST level rather than HIR

Author: GuillaumeGomez

src/librustdoc/core.rs

@@ -31,6 +31,7 @@ use crate::clean::inline::build_trait;
 use crate::clean::{self, ItemId};
 use crate::config::{Options as RustdocOptions, OutputFormat, RenderOptions};
 use crate::formats::cache::Cache;
+use crate::html::macro_expansion::{ExpandedCode, source_macro_expansion};
 use crate::passes;
 use crate::passes::Condition::*;
 use crate::passes::collect_intra_doc_links::LinkCollector;
@@ -335,11 +336,19 @@ pub(crate) fn run_global_ctxt(
     show_coverage: bool,
     render_options: RenderOptions,
     output_format: OutputFormat,
-) -> (clean::Crate, RenderOptions, Cache) {
+) -> (clean::Crate, RenderOptions, Cache, FxHashMap<rustc_span::BytePos, Vec<ExpandedCode>>) {
     // Certain queries assume that some checks were run elsewhere
     // (see https://github.com/rust-lang/rust/pull/73566#issuecomment-656954425),
     // so type-check everything other than function bodies in this crate before running lints.
 
+    let expanded_macros = {
+        // We need for these variables to be removed to ensure that the `Crate` won't be "stolen"
+        // anymore.
+        let (_resolver, krate) = &*tcx.resolver_for_lowering().borrow();
+
+        source_macro_expansion(&krate, &render_options, output_format, tcx.sess.source_map())
+    };
+
     // NOTE: this does not call `tcx.analysis()` so that we won't
     // typeck function bodies or run the default rustc lints.
     // (see `override_queries` in the `config`)
@@ -453,7 +462,7 @@ pub(crate) fn run_global_ctxt(
 
     tcx.dcx().abort_if_errors();
 
-    (krate, ctxt.render_options, ctxt.cache)
+    (krate, ctxt.render_options, ctxt.cache, expanded_macros)
 }
 
 /// Due to <https://github.com/rust-lang/rust/pull/73566>,

src/librustdoc/formats/renderer.rs

@@ -31,15 +31,6 @@ pub(crate) trait FormatRenderer<'tcx>: Sized {
     /// reset the information between each call to `item` by using `restore_module_data`.
     type ModuleData;
 
-    /// Sets up any state required for the renderer. When this is called the cache has already been
-    /// populated.
-    fn init(
-        krate: clean::Crate,
-        options: RenderOptions,
-        cache: Cache,
-        tcx: TyCtxt<'tcx>,
-    ) -> Result<(Self, clean::Crate), Error>;
-
     /// This method is called right before call [`Self::item`]. This method returns a type
     /// containing information that needs to be reset after the [`Self::item`] method has been
     /// called with the [`Self::restore_module_data`] method.
@@ -107,18 +98,23 @@ fn run_format_inner<'tcx, T: FormatRenderer<'tcx>>(
 }
 
 /// Main method for rendering a crate.
-pub(crate) fn run_format<'tcx, T: FormatRenderer<'tcx>>(
+pub(crate) fn run_format<
+    'tcx,
+    T: FormatRenderer<'tcx>,
+    F: FnOnce(clean::Crate, RenderOptions, Cache, TyCtxt<'tcx>) -> Result<(T, clean::Crate), Error>,
+>(
     krate: clean::Crate,
     options: RenderOptions,
     cache: Cache,
     tcx: TyCtxt<'tcx>,
+    init: F,
 ) -> Result<(), Error> {
     let prof = &tcx.sess.prof;
 
     let emit_crate = options.should_emit_crate();
     let (mut format_renderer, krate) = prof
         .verbose_generic_activity_with_arg("create_renderer", T::descr())
-        .run(|| T::init(krate, options, cache, tcx))?;
+        .run(|| init(krate, options, cache, tcx))?;
 
     if !emit_crate {
         return Ok(());

src/librustdoc/html/highlight.rs

@@ -18,7 +18,8 @@ use rustc_span::{BytePos, DUMMY_SP, Span};
 use super::format::{self, write_str};
 use crate::clean::PrimitiveType;
 use crate::html::escape::EscapeBodyText;
-use crate::html::render::{Context, ExpandedCode, LinkFromSrc};
+use crate::html::macro_expansion::ExpandedCode;
+use crate::html::render::{Context, LinkFromSrc};
 
 /// This type is needed in case we want to render links on items to allow to go to their definition.
 pub(crate) struct HrefContext<'a, 'tcx> {

src/librustdoc/html/macro_expansion.rs

@@ -0,0 +1,132 @@
+use rustc_ast::{Crate, Expr, Item};
+use rustc_ast::visit::{Visitor, walk_crate, walk_expr, walk_item};
+use rustc_data_structures::fx::FxHashMap;
+use rustc_span::{BytePos, Span};
+use rustc_span::source_map::SourceMap;
+
+use crate::config::{OutputFormat, RenderOptions};
+
+/// It returns the expanded macros correspondence map.
+pub(crate) fn source_macro_expansion(krate: &Crate, render_options: &RenderOptions, output_format: OutputFormat, source_map: &SourceMap) -> FxHashMap<BytePos, Vec<ExpandedCode>> {
+    if output_format == OutputFormat::Html && !render_options.html_no_source && render_options.generate_macro_expansion {
+        let mut expanded_visitor = ExpandedCodeVisitor { expanded_codes: Vec::new(), source_map };
+        walk_crate(&mut expanded_visitor, krate);
+        expanded_visitor.compute_expanded()
+    } else {
+        Default::default()
+    }
+}
+
+/// 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,
+}
+
+/// 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,
+    /// Expanded span
+    expanded_span: Span,
+}
+
+/// 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(crate) struct ExpandedCodeVisitor<'ast> {
+    expanded_codes: Vec<ExpandedCodeInfo>,
+    source_map: &'ast SourceMap,
+}
+
+impl<'ast> ExpandedCodeVisitor<'ast> {
+    fn handle_new_span<F: Fn() -> String>(&mut self, new_span: Span, f: F) {
+        if new_span.is_dummy() || !new_span.from_expansion() {
+            return;
+        }
+        let callsite_span = new_span.source_callsite();
+        if let Some(index) =
+            self.expanded_codes.iter().position(|info| info.span.overlaps(callsite_span))
+        {
+            let info = &mut self.expanded_codes[index];
+            if new_span.contains(info.expanded_span) {
+                // We replace the item.
+                info.span = callsite_span;
+                info.expanded_span = new_span;
+                info.code = f();
+            } else {
+                // We push the new item after the existing one.
+                let expanded_code = &mut self.expanded_codes[index];
+                expanded_code.code.push('\n');
+                expanded_code.code.push_str(&f());
+                let lo = BytePos(expanded_code.expanded_span.lo().0.min(new_span.lo().0));
+                let hi = BytePos(expanded_code.expanded_span.hi().0.min(new_span.hi().0));
+                expanded_code.expanded_span = expanded_code.expanded_span.with_lo(lo).with_hi(hi);
+            }
+        } else {
+            // We add a new item.
+            self.expanded_codes.push(ExpandedCodeInfo {
+                span: callsite_span,
+                code: f(),
+                expanded_span: new_span,
+            });
+        }
+    }
+
+    fn compute_expanded(mut self) -> FxHashMap<BytePos, Vec<ExpandedCode>> {
+        self.expanded_codes.sort_unstable_by(|item1, item2| item1.span.cmp(&item2.span));
+        let mut expanded: FxHashMap<BytePos, Vec<ExpandedCode>> = FxHashMap::default();
+        for ExpandedCodeInfo { span, code, .. } in self.expanded_codes {
+            if let Ok(lines) = self.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
+    }
+}
+
+// We need to use the AST pretty printing because:
+//
+// 1. HIR pretty printing doesn't display accurately the code (like `impl Trait`).
+// 2. `SourceMap::snippet_opt` might fail if the source is not available.
+impl<'ast> Visitor<'ast> for ExpandedCodeVisitor<'ast> {
+    fn visit_expr(&mut self, expr: &'ast Expr) {
+        if expr.span.from_expansion() {
+            self.handle_new_span(expr.span, || rustc_ast_pretty::pprust::expr_to_string(expr));
+        } else {
+            walk_expr(self, expr);
+        }
+    }
+
+    fn visit_item(&mut self, item: &'ast Item) {
+        if item.span.from_expansion() {
+            self.handle_new_span(item.span, || rustc_ast_pretty::pprust::item_to_string(item));
+        } else {
+            walk_item(self, item);
+        }
+    }
+}

src/librustdoc/html/mod.rs

@@ -4,6 +4,7 @@ pub(crate) mod highlight;
 pub(crate) mod layout;
 mod length_limit;
 // used by the error-index generator, so it needs to be public
+pub(crate) mod macro_expansion;
 pub mod markdown;
 pub(crate) mod render;
 pub(crate) mod sources;

src/librustdoc/html/render/context.rs

@@ -29,8 +29,8 @@ use crate::formats::item_type::ItemType;
 use crate::html::escape::Escape;
 use crate::html::format::join_with_double_colon;
 use crate::html::layout::{self, BufDisplay};
+use crate::html::macro_expansion::ExpandedCode;
 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::{sources, static_files};
@@ -456,20 +456,13 @@ impl<'tcx> Context<'tcx> {
     }
 }
 
-/// Generates the documentation for `crate` into the directory `dst`
-impl<'tcx> FormatRenderer<'tcx> for Context<'tcx> {
-    fn descr() -> &'static str {
-        "html"
-    }
-
-    const RUN_ON_MODULE: bool = true;
-    type ModuleData = ContextInfo;
-
-    fn init(
+impl<'tcx> Context<'tcx> {
+    pub(crate) fn init(
         krate: clean::Crate,
         options: RenderOptions,
         cache: Cache,
         tcx: TyCtxt<'tcx>,
+        expanded_codes: FxHashMap<BytePos, Vec<ExpandedCode>>,
     ) -> Result<(Self, clean::Crate), Error> {
         // need to save a copy of the options for rendering the index page
         let md_opts = options.clone();
@@ -488,7 +481,6 @@ impl<'tcx> FormatRenderer<'tcx> for Context<'tcx> {
             generate_redirect_map,
             show_type_layout,
             generate_link_to_definition,
-            generate_macro_expansion,
             call_locations,
             no_emit_shared,
             html_no_source,
@@ -547,13 +539,12 @@ impl<'tcx> FormatRenderer<'tcx> for Context<'tcx> {
             }
         }
 
-        let (local_sources, matches, expanded_codes) = collect_spans_and_sources(
+        let (local_sources, matches) = collect_spans_and_sources(
             tcx,
             &krate,
             &src_root,
             include_sources,
             generate_link_to_definition,
-            generate_macro_expansion,
         );
 
         let (sender, receiver) = channel();
@@ -605,6 +596,16 @@ impl<'tcx> FormatRenderer<'tcx> for Context<'tcx> {
 
         Ok((cx, krate))
     }
+}
+
+/// Generates the documentation for `crate` into the directory `dst`
+impl<'tcx> FormatRenderer<'tcx> for Context<'tcx> {
+    fn descr() -> &'static str {
+        "html"
+    }
+
+    const RUN_ON_MODULE: bool = true;
+    type ModuleData = ContextInfo;
 
     fn save_module_data(&mut self) -> Self::ModuleData {
         self.deref_id_map.borrow_mut().clear();

src/librustdoc/html/render/mod.rs

@@ -64,7 +64,7 @@ use serde::{Serialize, Serializer};
 use tracing::{debug, info};
 
 pub(crate) use self::context::*;
-pub(crate) use self::span_map::{ExpandedCode, LinkFromSrc, collect_spans_and_sources};
+pub(crate) use self::span_map::{LinkFromSrc, collect_spans_and_sources};
 pub(crate) use self::write_shared::*;
 use crate::clean::{self, ItemId, RenderedLink};
 use crate::display::{Joined as _, MaybeDisplay as _};