Fix ICE when rustdoc is scraping examples inside of a proc macro

src/librustdoc/html/highlight.rs

@@ -270,12 +270,18 @@ struct Decorations {
 
 impl Decorations {
     fn new(info: DecorationInfo) -> Self {
-        let (starts, ends) = info
+        // Extract tuples (start, end, kind) into separate sequences of (start, kind) and (end).
+        let (mut starts, mut ends): (Vec<_>, Vec<_>) = info
             .0
             .into_iter()
             .map(|(kind, ranges)| ranges.into_iter().map(move |(lo, hi)| ((lo, kind), hi)))
             .flatten()
             .unzip();
+
+        // Sort the sequences in document order.
+        starts.sort_by_key(|(lo, _)| *lo);
+        ends.sort();
+
         Decorations { starts, ends }
     }
 }

src/librustdoc/scrape_examples.rs

@@ -10,7 +10,6 @@ use rustc_data_structures::fx::FxHashMap;
 use rustc_hir::{
     self as hir,
     intravisit::{self, Visitor},
-    HirId,
 };
 use rustc_interface::interface;
 use rustc_macros::{Decodable, Encodable};
@@ -83,15 +82,10 @@ crate struct CallLocation {
 
 impl CallLocation {
     fn new(
-        tcx: TyCtxt<'_>,
         expr_span: rustc_span::Span,
-        expr_id: HirId,
+        enclosing_item_span: rustc_span::Span,
         source_file: &SourceFile,
     ) -> Self {
-        let enclosing_item_span =
-            tcx.hir().span_with_body(tcx.hir().get_parent_item(expr_id)).source_callsite();
-        assert!(enclosing_item_span.contains(expr_span));
-
         CallLocation {
             call_expr: SyntaxRange::new(expr_span, source_file),
             enclosing_item: SyntaxRange::new(enclosing_item_span, source_file),
@@ -168,13 +162,29 @@ where
         // If this span comes from a macro expansion, then the source code may not actually show
         // a use of the given item, so it would be a poor example. Hence, we skip all uses in macros.
         if span.from_expansion() {
+            trace!("Rejecting expr from macro: {:?}", span);
+            return;
+        }
+
+        // If the enclosing item has a span coming from a proc macro, then we also don't want to include
+        // the example.
+        let enclosing_item_span = tcx.hir().span_with_body(tcx.hir().get_parent_item(ex.hir_id));
+        if enclosing_item_span.from_expansion() {
+            trace!("Rejecting expr ({:?}) from macro item: {:?}", span, enclosing_item_span);
             return;
         }
 
+        assert!(
+            enclosing_item_span.contains(span),
+            "Attempted to scrape call at [{:?}] whose enclosing item [{:?}] doesn't contain the span of the call.",
+            span,
+            enclosing_item_span
+        );
+
         // Save call site if the function resolves to a concrete definition
         if let ty::FnDef(def_id, _) = ty.kind() {
-            // Ignore functions not from the crate being documented
             if self.target_crates.iter().all(|krate| *krate != def_id.krate) {
+                trace!("Rejecting expr from crate not being documented: {:?}", span);
                 return;
             }
 
@@ -198,7 +208,8 @@ where
                 let fn_key = tcx.def_path_hash(*def_id);
                 let fn_entries = self.calls.entry(fn_key).or_default();
 
-                let location = CallLocation::new(tcx, span, ex.hir_id, &file);
+                trace!("Including expr: {:?}", span);
+                let location = CallLocation::new(span, enclosing_item_span, &file);
                 fn_entries.entry(abs_path).or_insert_with(mk_call_data).locations.push(location);
             }
         }
@@ -240,6 +251,13 @@ crate fn run(
         let mut finder = FindCalls { calls: &mut calls, tcx, map: tcx.hir(), cx, target_crates };
         tcx.hir().visit_all_item_likes(&mut finder.as_deep_visitor());
 
+        // Sort call locations within a given file in document order
+        for fn_calls in calls.values_mut() {
+            for file_calls in fn_calls.values_mut() {
+                file_calls.locations.sort_by_key(|loc| loc.call_expr.byte_span.0);
+            }
+        }
+
         // Save output to provided path
         let mut encoder = FileEncoder::new(options.output_path).map_err(|e| e.to_string())?;
         calls.encode(&mut encoder).map_err(|e| e.to_string())?;

src/test/run-make-fulldeps/rustdoc-scrape-examples-macros/Makefile

@@ -0,0 +1,18 @@
+-include ../../run-make-fulldeps/tools.mk
+
+OUTPUT_DIR := "$(TMPDIR)/rustdoc"
+DYLIB_NAME := $(shell echo | $(RUSTC) --crate-name foobar_macro --crate-type dylib --print file-names -)
+
+all:
+	$(RUSTC) src/proc.rs --crate-name foobar_macro --edition=2021 --crate-type proc-macro --emit=dep-info,link
+
+	$(RUSTC) src/lib.rs --crate-name foobar --edition=2021 --crate-type lib --emit=dep-info,link
+
+	$(RUSTDOC) examples/ex.rs --crate-name ex --crate-type bin --output $(OUTPUT_DIR) \
+		--extern foobar=$(TMPDIR)/libfoobar.rlib --extern foobar_macro=$(TMPDIR)/$(DYLIB_NAME) \
+		-Z unstable-options --scrape-examples-output-path $(TMPDIR)/ex.calls --scrape-examples-target-crate foobar
+
+	$(RUSTDOC) src/lib.rs --crate-name foobar --crate-type lib --output $(OUTPUT_DIR) \
+		-Z unstable-options --with-examples $(TMPDIR)/ex.calls
+
+	$(HTMLDOCCK) $(OUTPUT_DIR) src/lib.rs

src/test/run-make-fulldeps/rustdoc-scrape-examples-macros/examples/ex.rs

@@ -0,0 +1,27 @@
+extern crate foobar;
+extern crate foobar_macro;
+
+use foobar::*;
+use foobar_macro::*;
+
+a_proc_macro!(); // no
+
+#[an_attr_macro]
+fn a() {
+  f(); // no
+}
+
+#[an_attr_macro(with_span)]
+fn b() {
+  f(); // yes
+}
+
+fn c() {
+  a_rules_macro!(f()); // yes
+}
+
+fn d() {
+  a_rules_macro!(()); // no
+}
+
+fn main(){}

src/test/run-make-fulldeps/rustdoc-scrape-examples-macros/src/lib.rs

@@ -0,0 +1,12 @@
+// Scraped example should only include line numbers for items b and c in ex.rs
+// @!has foobar/fn.f.html '//*[@class="line-numbers"]' '14'
+// @has foobar/fn.f.html '//*[@class="line-numbers"]' '15'
+// @has foobar/fn.f.html '//*[@class="line-numbers"]' '21'
+// @!has foobar/fn.f.html '//*[@class="line-numbers"]' '22'
+
+pub fn f() {}
+
+#[macro_export]
+macro_rules! a_rules_macro {
+  ($e:expr) => { ($e, foobar::f()); }
+}

src/test/run-make-fulldeps/rustdoc-scrape-examples-macros/src/proc.rs

@@ -0,0 +1,39 @@
+extern crate proc_macro;
+use proc_macro::*;
+
+#[proc_macro]
+pub fn a_proc_macro(_item: TokenStream) -> TokenStream {
+    "fn ex() { foobar::f(); }".parse().unwrap()
+}
+
+// inserts foobar::f() to the end of the function
+#[proc_macro_attribute]
+pub fn an_attr_macro(attr: TokenStream, item: TokenStream) -> TokenStream {
+    let new_call: TokenStream = "foobar::f();".parse().unwrap();
+
+    let mut tokens = item.into_iter();
+
+    let fn_tok = tokens.next().unwrap();
+    let ident_tok = tokens.next().unwrap();
+    let args_tok = tokens.next().unwrap();
+    let body = match tokens.next().unwrap() {
+        TokenTree::Group(g) => {
+            let new_g = Group::new(g.delimiter(), new_call);
+            let mut outer_g = Group::new(
+                g.delimiter(),
+                [TokenTree::Group(g.clone()), TokenTree::Group(new_g)].into_iter().collect(),
+            );
+
+            if attr.to_string() == "with_span" {
+                outer_g.set_span(g.span());
+            }
+
+            TokenTree::Group(outer_g)
+        }
+        _ => unreachable!(),
+    };
+
+    let tokens = vec![fn_tok, ident_tok, args_tok, body].into_iter().collect::<TokenStream>();
+
+    tokens
+}

src/test/run-make/rustdoc-scrape-examples-ordering/examples/ex1.rs

@@ -1,8 +1,10 @@
 fn main() {
-    foobar::ok();
+    foobar::ok(0);
 
     // this is a
 
+    //  ..
+
     // BIG
 
     // item

src/test/run-make/rustdoc-scrape-examples-ordering/examples/ex2.rs

@@ -1,4 +1,8 @@
 fn main() {
-    foobar::ok();
+    foobar::ok(1);
     // small item
 }
+
+fn f() {
+    foobar::ok(2);
+}

src/test/run-make/rustdoc-scrape-examples-ordering/src/lib.rs

@@ -1,4 +1,7 @@
 // @has foobar/fn.ok.html '//*[@class="docblock scraped-example-list"]' 'ex2'
 // @has foobar/fn.ok.html '//*[@class="more-scraped-examples"]' 'ex1'
+// @has foobar/fn.ok.html '//*[@class="highlight focus"]' '1'
+// @has foobar/fn.ok.html '//*[@class="highlight"]' '2'
+// @has foobar/fn.ok.html '//*[@class="highlight focus"]' '0'
 
-pub fn ok() {}
+pub fn ok(_x: i32) {}