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

Author: willcrichton

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);
             }
         }

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

@@ -0,0 +1,17 @@
+-include ../../run-make-fulldeps/tools.mk
+
+OUTPUT_DIR := "$(TMPDIR)/rustdoc"
+
+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)/libfoobar_macro.so \
+		-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/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/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/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
+}