Make Visitor::FnKind and MutVisitor::FnKind compatible

6 files changed, 114 insertions, 62 deletions

diff --git a/compiler/rustc_ast/src/mut_visit.rs b/compiler/rustc_ast/src/mut_visit.rs
index 44bb44cb728..fe5ec1e5ad3 100644
--- a/compiler/rustc_ast/src/mut_visit.rs
+++ b/compiler/rustc_ast/src/mut_visit.rs
@@ -23,7 +23,7 @@ use crate::ast::*;
 use crate::ptr::P;
 use crate::token::{self, Token};
 use crate::tokenstream::*;
-use crate::visit::{AssocCtxt, BoundKind};
+use crate::visit::{AssocCtxt, BoundKind, FnCtxt};
 
 pub trait ExpectOne<A: Array> {
     fn expect_one(self, err: &'static str) -> A::Item;
@@ -37,7 +37,15 @@ impl<A: Array> ExpectOne<A> for SmallVec<A> {
 }
 
 pub trait WalkItemKind {
-    fn walk(&mut self, span: Span, id: NodeId, visitor: &mut impl MutVisitor);
+    fn walk(
+        &mut self,
+        span: Span,
+        id: NodeId,
+        ident: &mut Ident,
+        visibility: &mut Visibility,
+        ctxt: AssocCtxt,
+        visitor: &mut impl MutVisitor,
+    );
 }
 
 pub trait MutVisitor: Sized {
@@ -114,9 +122,9 @@ pub trait MutVisitor: Sized {
     fn flat_map_assoc_item(
         &mut self,
         i: P<AssocItem>,
-        _ctxt: AssocCtxt,
+        ctxt: AssocCtxt,
     ) -> SmallVec<[P<AssocItem>; 1]> {
-        walk_flat_map_item(self, i)
+        walk_flat_map_assoc_item(self, i, ctxt)
     }
 
     fn visit_fn_decl(&mut self, d: &mut P<FnDecl>) {
@@ -880,7 +888,7 @@ fn walk_coroutine_kind<T: MutVisitor>(vis: &mut T, coroutine_kind: &mut Coroutin
 
 fn walk_fn<T: MutVisitor>(vis: &mut T, kind: FnKind<'_>) {
     match kind {
-        FnKind::Fn(FnSig { header, decl, span }, generics, body) => {
+        FnKind::Fn(_ctxt, _ident, FnSig { header, decl, span }, _visibility, generics, body) => {
             // Identifier and visibility are visited as a part of the item.
             vis.visit_fn_header(header);
             vis.visit_generics(generics);
@@ -890,8 +898,9 @@ fn walk_fn<T: MutVisitor>(vis: &mut T, kind: FnKind<'_>) {
             }
             vis.visit_span(span);
         }
-        FnKind::Closure(binder, decl, body) => {
+        FnKind::Closure(binder, coroutine_kind, decl, body) => {
             vis.visit_closure_binder(binder);
+            coroutine_kind.as_mut().map(|coroutine_kind| vis.visit_coroutine_kind(coroutine_kind));
             vis.visit_fn_decl(decl);
             vis.visit_expr(body);
         }
@@ -1083,13 +1092,24 @@ pub fn walk_item_kind(
     kind: &mut impl WalkItemKind,
     span: Span,
     id: NodeId,
+    ident: &mut Ident,
+    visibility: &mut Visibility,
+    ctxt: AssocCtxt,
     vis: &mut impl MutVisitor,
 ) {
-    kind.walk(span, id, vis)
+    kind.walk(span, id, ident, visibility, ctxt, vis)
 }
 
 impl WalkItemKind for ItemKind {
-    fn walk(&mut self, span: Span, id: NodeId, vis: &mut impl MutVisitor) {
+    fn walk(
+        &mut self,
+        span: Span,
+        id: NodeId,
+        ident: &mut Ident,
+        visibility: &mut Visibility,
+        _ctxt: AssocCtxt,
+        vis: &mut impl MutVisitor,
+    ) {
         match self {
             ItemKind::ExternCrate(_orig_name) => {}
             ItemKind::Use(use_tree) => vis.visit_use_tree(use_tree),
@@ -1102,7 +1122,11 @@ impl WalkItemKind for ItemKind {
             }
             ItemKind::Fn(box Fn { defaultness, generics, sig, body }) => {
                 visit_defaultness(vis, defaultness);
-                vis.visit_fn(FnKind::Fn(sig, generics, body), span, id);
+                vis.visit_fn(
+                    FnKind::Fn(FnCtxt::Free, ident, sig, visibility, generics, body),
+                    span,
+                    id,
+                );
             }
             ItemKind::Mod(safety, mod_kind) => {
                 visit_safety(vis, safety);
@@ -1201,14 +1225,26 @@ impl WalkItemKind for ItemKind {
 }
 
 impl WalkItemKind for AssocItemKind {
-    fn walk(&mut self, span: Span, id: NodeId, visitor: &mut impl MutVisitor) {
+    fn walk(
+        &mut self,
+        span: Span,
+        id: NodeId,
+        ident: &mut Ident,
+        visibility: &mut Visibility,
+        ctxt: AssocCtxt,
+        visitor: &mut impl MutVisitor,
+    ) {
         match self {
             AssocItemKind::Const(item) => {
                 visit_const_item(item, visitor);
             }
             AssocItemKind::Fn(box Fn { defaultness, generics, sig, body }) => {
                 visit_defaultness(visitor, defaultness);
-                visitor.visit_fn(FnKind::Fn(sig, generics, body), span, id);
+                visitor.visit_fn(
+                    FnKind::Fn(FnCtxt::Assoc(ctxt), ident, sig, visibility, generics, body),
+                    span,
+                    id,
+                );
             }
             AssocItemKind::Type(box TyAlias {
                 defaultness,
@@ -1288,24 +1324,39 @@ pub fn walk_crate<T: MutVisitor>(vis: &mut T, krate: &mut Crate) {
     vis.visit_span(inject_use_span);
 }
 
-/// Mutates one item, returning the item again.
 pub fn walk_flat_map_item<K: WalkItemKind>(
     visitor: &mut impl MutVisitor,
+    item: P<Item<K>>,
+) -> SmallVec<[P<Item<K>>; 1]> {
+    walk_flat_map_assoc_item(visitor, item, AssocCtxt::Trait /* ignored */)
+}
+
+pub fn walk_flat_map_assoc_item<K: WalkItemKind>(
+    visitor: &mut impl MutVisitor,
     mut item: P<Item<K>>,
+    ctxt: AssocCtxt,
 ) -> SmallVec<[P<Item<K>>; 1]> {
     let Item { ident, attrs, id, kind, vis, span, tokens } = item.deref_mut();
     visitor.visit_id(id);
     visit_attrs(visitor, attrs);
     visitor.visit_vis(vis);
     visitor.visit_ident(ident);
-    kind.walk(*span, *id, visitor);
+    kind.walk(*span, *id, ident, vis, ctxt, visitor);
     visit_lazy_tts(visitor, tokens);
     visitor.visit_span(span);
     smallvec![item]
 }
 
 impl WalkItemKind for ForeignItemKind {
-    fn walk(&mut self, span: Span, id: NodeId, visitor: &mut impl MutVisitor) {
+    fn walk(
+        &mut self,
+        span: Span,
+        id: NodeId,
+        ident: &mut Ident,
+        visibility: &mut Visibility,
+        _ctxt: AssocCtxt,
+        visitor: &mut impl MutVisitor,
+    ) {
         match self {
             ForeignItemKind::Static(box StaticItem { ty, mutability: _, expr, safety: _ }) => {
                 visitor.visit_ty(ty);
@@ -1313,7 +1364,11 @@ impl WalkItemKind for ForeignItemKind {
             }
             ForeignItemKind::Fn(box Fn { defaultness, generics, sig, body }) => {
                 visit_defaultness(visitor, defaultness);
-                visitor.visit_fn(FnKind::Fn(sig, generics, body), span, id);
+                visitor.visit_fn(
+                    FnKind::Fn(FnCtxt::Foreign, ident, sig, visibility, generics, body),
+                    span,
+                    id,
+                );
             }
             ForeignItemKind::TyAlias(box TyAlias {
                 defaultness,
@@ -1522,9 +1577,8 @@ pub fn walk_expr<T: MutVisitor>(vis: &mut T, Expr { kind, id, span, attrs, token
             fn_arg_span,
         }) => {
             visit_constness(vis, constness);
-            coroutine_kind.as_mut().map(|coroutine_kind| vis.visit_coroutine_kind(coroutine_kind));
             vis.visit_capture_by(capture_clause);
-            vis.visit_fn(FnKind::Closure(binder, fn_decl, body), *span, *id);
+            vis.visit_fn(FnKind::Closure(binder, coroutine_kind, fn_decl, body), *span, *id);
             vis.visit_span(fn_decl_span);
             vis.visit_span(fn_arg_span);
         }
@@ -1785,8 +1839,20 @@ impl<N: DummyAstNode, T: DummyAstNode> DummyAstNode for crate::ast_traits::AstNo
 #[derive(Debug)]
 pub enum FnKind<'a> {
     /// E.g., `fn foo()`, `fn foo(&self)`, or `extern "Abi" fn foo()`.
-    Fn(&'a mut FnSig, &'a mut Generics, &'a mut Option<P<Block>>),
+    Fn(
+        FnCtxt,
+        &'a mut Ident,
+        &'a mut FnSig,
+        &'a mut Visibility,
+        &'a mut Generics,
+        &'a mut Option<P<Block>>,
+    ),
 
     /// E.g., `|x, y| body`.
-    Closure(&'a mut ClosureBinder, &'a mut P<FnDecl>, &'a mut P<Expr>),
+    Closure(
+        &'a mut ClosureBinder,
+        &'a mut Option<CoroutineKind>,
+        &'a mut P<FnDecl>,
+        &'a mut P<Expr>,
+    ),
 }
diff --git a/compiler/rustc_ast/src/visit.rs b/compiler/rustc_ast/src/visit.rs
index 2f8115441de..01c7a815e00 100644
--- a/compiler/rustc_ast/src/visit.rs
+++ b/compiler/rustc_ast/src/visit.rs
@@ -66,7 +66,7 @@ impl BoundKind {
 #[derive(Copy, Clone, Debug)]
 pub enum FnKind<'a> {
     /// E.g., `fn foo()`, `fn foo(&self)`, or `extern "Abi" fn foo()`.
-    Fn(FnCtxt, Ident, &'a FnSig, &'a Visibility, &'a Generics, Option<&'a Block>),
+    Fn(FnCtxt, &'a Ident, &'a FnSig, &'a Visibility, &'a Generics, &'a Option<P<Block>>),
 
     /// E.g., `|x, y| body`.
     Closure(&'a ClosureBinder, &'a Option<CoroutineKind>, &'a FnDecl, &'a Expr),
@@ -357,7 +357,7 @@ impl WalkItemKind for ItemKind {
                 visit_opt!(visitor, visit_expr, expr);
             }
             ItemKind::Fn(box Fn { defaultness: _, generics, sig, body }) => {
-                let kind = FnKind::Fn(FnCtxt::Free, *ident, sig, vis, generics, body.as_deref());
+                let kind = FnKind::Fn(FnCtxt::Free, ident, sig, vis, generics, body);
                 try_visit!(visitor.visit_fn(kind, *span, *id));
             }
             ItemKind::Mod(_unsafety, mod_kind) => match mod_kind {
@@ -687,15 +687,15 @@ impl WalkItemKind for ForeignItemKind {
         _ctxt: AssocCtxt,
         visitor: &mut V,
     ) -> V::Result {
-        let &Item { id, span, ident, ref vis, .. } = item;
+        let Item { id, span, ident, vis, .. } = item;
         match self {
             ForeignItemKind::Static(box StaticItem { ty, mutability: _, expr, safety: _ }) => {
                 try_visit!(visitor.visit_ty(ty));
                 visit_opt!(visitor, visit_expr, expr);
             }
             ForeignItemKind::Fn(box Fn { defaultness: _, generics, sig, body }) => {
-                let kind = FnKind::Fn(FnCtxt::Foreign, ident, sig, vis, generics, body.as_deref());
-                try_visit!(visitor.visit_fn(kind, span, id));
+                let kind = FnKind::Fn(FnCtxt::Foreign, ident, sig, vis, generics, body);
+                try_visit!(visitor.visit_fn(kind, *span, *id));
             }
             ForeignItemKind::TyAlias(box TyAlias {
                 generics,
@@ -850,7 +850,7 @@ impl WalkItemKind for AssocItemKind {
         ctxt: AssocCtxt,
         visitor: &mut V,
     ) -> V::Result {
-        let &Item { id, span, ident, ref vis, .. } = item;
+        let Item { id, span, ident, vis, .. } = item;
         match self {
             AssocItemKind::Const(box ConstItem { defaultness: _, generics, ty, expr }) => {
                 try_visit!(visitor.visit_generics(generics));
@@ -858,9 +858,8 @@ impl WalkItemKind for AssocItemKind {
                 visit_opt!(visitor, visit_expr, expr);
             }
             AssocItemKind::Fn(box Fn { defaultness: _, generics, sig, body }) => {
-                let kind =
-                    FnKind::Fn(FnCtxt::Assoc(ctxt), ident, sig, vis, generics, body.as_deref());
-                try_visit!(visitor.visit_fn(kind, span, id));
+                let kind = FnKind::Fn(FnCtxt::Assoc(ctxt), ident, sig, vis, generics, body);
+                try_visit!(visitor.visit_fn(kind, *span, *id));
             }
             AssocItemKind::Type(box TyAlias {
                 generics,
@@ -891,7 +890,7 @@ impl WalkItemKind for AssocItemKind {
             }
             AssocItemKind::DelegationMac(box DelegationMac { qself, prefix, suffixes, body }) => {
                 try_visit!(walk_qself(visitor, qself));
-                try_visit!(visitor.visit_path(prefix, id));
+                try_visit!(visitor.visit_path(prefix, *id));
                 if let Some(suffixes) = suffixes {
                     for (ident, rename) in suffixes {
                         visitor.visit_ident(ident);
diff --git a/compiler/rustc_ast_passes/src/ast_validation.rs b/compiler/rustc_ast_passes/src/ast_validation.rs
index dee48586f34..07a6f4e5ee7 100644
--- a/compiler/rustc_ast_passes/src/ast_validation.rs
+++ b/compiler/rustc_ast_passes/src/ast_validation.rs
@@ -946,8 +946,7 @@ impl<'a> Visitor<'a> for AstValidator<'a> {
 
                 self.visit_vis(&item.vis);
                 self.visit_ident(&item.ident);
-                let kind =
-                    FnKind::Fn(FnCtxt::Free, item.ident, sig, &item.vis, generics, body.as_deref());
+                let kind = FnKind::Fn(FnCtxt::Free, &item.ident, sig, &item.vis, generics, body);
                 self.visit_fn(kind, item.span, item.id);
                 walk_list!(self, visit_attribute, &item.attrs);
                 return; // Avoid visiting again.
@@ -1476,14 +1475,8 @@ impl<'a> Visitor<'a> for AstValidator<'a> {
             {
                 self.visit_vis(&item.vis);
                 self.visit_ident(&item.ident);
-                let kind = FnKind::Fn(
-                    FnCtxt::Assoc(ctxt),
-                    item.ident,
-                    sig,
-                    &item.vis,
-                    generics,
-                    body.as_deref(),
-                );
+                let kind =
+                    FnKind::Fn(FnCtxt::Assoc(ctxt), &item.ident, sig, &item.vis, generics, body);
                 walk_list!(self, visit_attribute, &item.attrs);
                 self.visit_fn(kind, item.span, item.id);
             }
diff --git a/compiler/rustc_builtin_macros/src/test_harness.rs b/compiler/rustc_builtin_macros/src/test_harness.rs
index 95348453308..4a1ade77952 100644
--- a/compiler/rustc_builtin_macros/src/test_harness.rs
+++ b/compiler/rustc_builtin_macros/src/test_harness.rs
@@ -6,7 +6,7 @@ use rustc_ast as ast;
 use rustc_ast::entry::EntryPointType;
 use rustc_ast::mut_visit::*;
 use rustc_ast::ptr::P;
-use rustc_ast::visit::{Visitor, walk_item};
+use rustc_ast::visit::{AssocCtxt, Visitor, walk_item};
 use rustc_ast::{ModKind, attr};
 use rustc_errors::DiagCtxtHandle;
 use rustc_expand::base::{ExtCtxt, ResolverExpand};
@@ -144,7 +144,15 @@ impl<'a> MutVisitor for TestHarnessGenerator<'a> {
             item.kind
         {
             let prev_tests = mem::take(&mut self.tests);
-            walk_item_kind(&mut item.kind, item.span, item.id, self);
+            walk_item_kind(
+                &mut item.kind,
+                item.span,
+                item.id,
+                &mut item.ident,
+                &mut item.vis,
+                AssocCtxt::Trait, /* ignored */
+                self,
+            );
             self.add_test_cases(item.id, span, prev_tests);
         } else {
             // But in those cases, we emit a lint to warn the user of these missing tests.
diff --git a/src/tools/rustfmt/src/items.rs b/src/tools/rustfmt/src/items.rs
index 327a547b295..68c481bda6e 100644
--- a/src/tools/rustfmt/src/items.rs
+++ b/src/tools/rustfmt/src/items.rs
@@ -3449,21 +3449,14 @@ impl Rewrite for ast::ForeignItem {
                     ref generics,
                     ref body,
                 } = **fn_kind;
-                if let Some(ref body) = body {
+                if body.is_some() {
                     let mut visitor = FmtVisitor::from_context(context);
                     visitor.block_indent = shape.indent;
                     visitor.last_pos = self.span.lo();
                     let inner_attrs = inner_attributes(&self.attrs);
                     let fn_ctxt = visit::FnCtxt::Foreign;
                     visitor.visit_fn(
-                        visit::FnKind::Fn(
-                            fn_ctxt,
-                            self.ident,
-                            sig,
-                            &self.vis,
-                            generics,
-                            Some(body),
-                        ),
+                        visit::FnKind::Fn(fn_ctxt, &self.ident, sig, &self.vis, generics, body),
                         &sig.decl,
                         self.span,
                         defaultness,
diff --git a/src/tools/rustfmt/src/visitor.rs b/src/tools/rustfmt/src/visitor.rs
index 8102fe7ad8f..9b116b620b7 100644
--- a/src/tools/rustfmt/src/visitor.rs
+++ b/src/tools/rustfmt/src/visitor.rs
@@ -390,7 +390,7 @@ impl<'b, 'a: 'b> FmtVisitor<'a> {
                 block = b;
                 self.rewrite_fn_before_block(
                     indent,
-                    ident,
+                    *ident,
                     &FnSig::from_fn_kind(&fk, fd, defaultness),
                     mk_sp(s.lo(), b.span.lo()),
                 )
@@ -540,21 +540,14 @@ impl<'b, 'a: 'b> FmtVisitor<'a> {
                         ref generics,
                         ref body,
                     } = **fn_kind;
-                    if let Some(ref body) = body {
+                    if body.is_some() {
                         let inner_attrs = inner_attributes(&item.attrs);
                         let fn_ctxt = match sig.header.ext {
                             ast::Extern::None => visit::FnCtxt::Free,
                             _ => visit::FnCtxt::Foreign,
                         };
                         self.visit_fn(
-                            visit::FnKind::Fn(
-                                fn_ctxt,
-                                item.ident,
-                                sig,
-                                &item.vis,
-                                generics,
-                                Some(body),
-                            ),
+                            visit::FnKind::Fn(fn_ctxt, &item.ident, sig, &item.vis, generics, body),
                             &sig.decl,
                             item.span,
                             defaultness,
@@ -648,11 +641,11 @@ impl<'b, 'a: 'b> FmtVisitor<'a> {
                     ref generics,
                     ref body,
                 } = **fn_kind;
-                if let Some(ref body) = body {
+                if body.is_some() {
                     let inner_attrs = inner_attributes(&ai.attrs);
                     let fn_ctxt = visit::FnCtxt::Assoc(assoc_ctxt);
                     self.visit_fn(
-                        visit::FnKind::Fn(fn_ctxt, ai.ident, sig, &ai.vis, generics, Some(body)),
+                        visit::FnKind::Fn(fn_ctxt, &ai.ident, sig, &ai.vis, generics, body),
                         &sig.decl,
                         ai.span,
                         defaultness,