def_collector: Fully visit async functions

Author: petrochenkov

src/librustc_resolve/def_collector.rs

@@ -2,7 +2,7 @@ use log::debug;
 use rustc::hir::map::definitions::*;
 use rustc_ast::ast::*;
 use rustc_ast::token::{self, Token};
-use rustc_ast::visit;
+use rustc_ast::visit::{self, FnKind};
 use rustc_expand::expand::AstFragment;
 use rustc_hir::def_id::DefIndex;
 use rustc_span::hygiene::ExpnId;
@@ -38,42 +38,6 @@ impl<'a> DefCollector<'a> {
         self.parent_def = orig_parent_def;
     }
 
-    fn visit_async_fn(
-        &mut self,
-        id: NodeId,
-        name: Name,
-        span: Span,
-        header: &FnHeader,
-        generics: &'a Generics,
-        decl: &'a FnDecl,
-        body: Option<&'a Block>,
-    ) {
-        let (closure_id, return_impl_trait_id) = match header.asyncness {
-            Async::Yes { span: _, closure_id, return_impl_trait_id } => {
-                (closure_id, return_impl_trait_id)
-            }
-            _ => unreachable!(),
-        };
-
-        // For async functions, we need to create their inner defs inside of a
-        // closure to match their desugared representation.
-        let fn_def_data = DefPathData::ValueNs(name);
-        let fn_def = self.create_def(id, fn_def_data, span);
-        return self.with_parent(fn_def, |this| {
-            this.create_def(return_impl_trait_id, DefPathData::ImplTrait, span);
-
-            visit::walk_generics(this, generics);
-            visit::walk_fn_decl(this, decl);
-
-            let closure_def = this.create_def(closure_id, DefPathData::ClosureExpr, span);
-            this.with_parent(closure_def, |this| {
-                if let Some(body) = body {
-                    visit::walk_block(this, body);
-                }
-            })
-        });
-    }
-
     fn collect_field(&mut self, field: &'a StructField, index: Option<usize>) {
         let index = |this: &Self| {
             index.unwrap_or_else(|| {
@@ -117,17 +81,6 @@ impl<'a> visit::Visitor<'a> for DefCollector<'a> {
             | ItemKind::ExternCrate(..)
             | ItemKind::ForeignMod(..)
             | ItemKind::TyAlias(..) => DefPathData::TypeNs(i.ident.name),
-            ItemKind::Fn(_, sig, generics, body) if sig.header.asyncness.is_async() => {
-                return self.visit_async_fn(
-                    i.id,
-                    i.ident.name,
-                    i.span,
-                    &sig.header,
-                    generics,
-                    &sig.decl,
-                    body.as_deref(),
-                );
-            }
             ItemKind::Static(..) | ItemKind::Const(..) | ItemKind::Fn(..) => {
                 DefPathData::ValueNs(i.ident.name)
             }
@@ -154,6 +107,27 @@ impl<'a> visit::Visitor<'a> for DefCollector<'a> {
         });
     }
 
+    fn visit_fn(&mut self, fn_kind: FnKind<'a>, span: Span, _: NodeId) {
+        if let FnKind::Fn(_, _, sig, _, body) = fn_kind {
+            if let Async::Yes { closure_id, return_impl_trait_id, .. } = sig.header.asyncness {
+                self.create_def(return_impl_trait_id, DefPathData::ImplTrait, span);
+
+                // For async functions, we need to create their inner defs inside of a
+                // closure to match their desugared representation. Besides that,
+                // we must mirror everything that `visit::walk_fn` below does.
+                self.visit_fn_header(&sig.header);
+                visit::walk_fn_decl(self, &sig.decl);
+                if let Some(body) = body {
+                    let closure_def = self.create_def(closure_id, DefPathData::ClosureExpr, span);
+                    self.with_parent(closure_def, |this| this.visit_block(body));
+                }
+                return;
+            }
+        }
+
+        visit::walk_fn(self, fn_kind, span);
+    }
+
     fn visit_use_tree(&mut self, use_tree: &'a UseTree, id: NodeId, _nested: bool) {
         self.create_def(id, DefPathData::Misc, use_tree.span);
         visit::walk_use_tree(self, use_tree, id);
@@ -215,19 +189,6 @@ impl<'a> visit::Visitor<'a> for DefCollector<'a> {
 
     fn visit_assoc_item(&mut self, i: &'a AssocItem, ctxt: visit::AssocCtxt) {
         let def_data = match &i.kind {
-            AssocItemKind::Fn(_, FnSig { header, decl }, generics, body)
-                if header.asyncness.is_async() =>
-            {
-                return self.visit_async_fn(
-                    i.id,
-                    i.ident.name,
-                    i.span,
-                    header,
-                    generics,
-                    decl,
-                    body.as_deref(),
-                );
-            }
             AssocItemKind::Fn(..) | AssocItemKind::Const(..) => DefPathData::ValueNs(i.ident.name),
             AssocItemKind::TyAlias(..) => DefPathData::TypeNs(i.ident.name),
             AssocItemKind::Macro(..) => return self.visit_macro_invoc(i.id),

src/test/ui/async-await/expansion-in-attrs.rs

@@ -0,0 +1,13 @@
+// check-pass
+// edition:2018
+
+macro_rules! with_doc {
+    ($doc: expr) => {
+        #[doc = $doc]
+        async fn f() {}
+    };
+}
+
+with_doc!(concat!(""));
+
+fn main() {}