Visit generics inside visit_fn.

8 files changed, 65 insertions, 40 deletions

diff --git a/compiler/rustc_ast/src/visit.rs b/compiler/rustc_ast/src/visit.rs
index 3183612597d..d925c6dd354 100644
--- a/compiler/rustc_ast/src/visit.rs
+++ b/compiler/rustc_ast/src/visit.rs
@@ -35,7 +35,7 @@ pub enum FnCtxt {
 #[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, Option<&'a Block>),
+    Fn(FnCtxt, Ident, &'a FnSig, &'a Visibility, &'a Generics, Option<&'a Block>),
 
     /// E.g., `|x, y| body`.
     Closure(&'a FnDecl, &'a Expr),
@@ -44,7 +44,7 @@ pub enum FnKind<'a> {
 impl<'a> FnKind<'a> {
     pub fn header(&self) -> Option<&'a FnHeader> {
         match *self {
-            FnKind::Fn(_, _, sig, _, _) => Some(&sig.header),
+            FnKind::Fn(_, _, sig, _, _, _) => Some(&sig.header),
             FnKind::Closure(_, _) => None,
         }
     }
@@ -58,7 +58,7 @@ impl<'a> FnKind<'a> {
 
     pub fn decl(&self) -> &'a FnDecl {
         match self {
-            FnKind::Fn(_, _, sig, _, _) => &sig.decl,
+            FnKind::Fn(_, _, sig, _, _, _) => &sig.decl,
             FnKind::Closure(decl, _) => decl,
         }
     }
@@ -295,8 +295,8 @@ pub fn walk_item<'a, V: Visitor<'a>>(visitor: &mut V, item: &'a Item) {
             walk_list!(visitor, visit_expr, expr);
         }
         ItemKind::Fn(box Fn { defaultness: _, ref generics, ref sig, ref body }) => {
-            visitor.visit_generics(generics);
-            let kind = FnKind::Fn(FnCtxt::Free, item.ident, sig, &item.vis, body.as_deref());
+            let kind =
+                FnKind::Fn(FnCtxt::Free, item.ident, sig, &item.vis, generics, body.as_deref());
             visitor.visit_fn(kind, item.span, item.id)
         }
         ItemKind::Mod(_unsafety, ref mod_kind) => match mod_kind {
@@ -561,8 +561,7 @@ pub fn walk_foreign_item<'a, V: Visitor<'a>>(visitor: &mut V, item: &'a ForeignI
             walk_list!(visitor, visit_expr, expr);
         }
         ForeignItemKind::Fn(box Fn { defaultness: _, ref generics, ref sig, ref body }) => {
-            visitor.visit_generics(generics);
-            let kind = FnKind::Fn(FnCtxt::Foreign, ident, sig, vis, body.as_deref());
+            let kind = FnKind::Fn(FnCtxt::Foreign, ident, sig, vis, generics, body.as_deref());
             visitor.visit_fn(kind, span, id);
         }
         ForeignItemKind::TyAlias(box TyAlias { generics, bounds, ty, .. }) => {
@@ -644,7 +643,8 @@ pub fn walk_fn_decl<'a, V: Visitor<'a>>(visitor: &mut V, function_declaration: &
 
 pub fn walk_fn<'a, V: Visitor<'a>>(visitor: &mut V, kind: FnKind<'a>, _span: Span) {
     match kind {
-        FnKind::Fn(_, _, sig, _, body) => {
+        FnKind::Fn(_, _, sig, _, generics, body) => {
+            visitor.visit_generics(generics);
             visitor.visit_fn_header(&sig.header);
             walk_fn_decl(visitor, &sig.decl);
             walk_list!(visitor, visit_block, body);
@@ -667,8 +667,7 @@ pub fn walk_assoc_item<'a, V: Visitor<'a>>(visitor: &mut V, item: &'a AssocItem,
             walk_list!(visitor, visit_expr, expr);
         }
         AssocItemKind::Fn(box Fn { defaultness: _, ref generics, ref sig, ref body }) => {
-            visitor.visit_generics(generics);
-            let kind = FnKind::Fn(FnCtxt::Assoc(ctxt), ident, sig, vis, body.as_deref());
+            let kind = FnKind::Fn(FnCtxt::Assoc(ctxt), ident, sig, vis, generics, body.as_deref());
             visitor.visit_fn(kind, span, id);
         }
         AssocItemKind::TyAlias(box TyAlias { generics, bounds, ty, .. }) => {
diff --git a/compiler/rustc_ast_passes/src/ast_validation.rs b/compiler/rustc_ast_passes/src/ast_validation.rs
index 33727084ccf..8d815e95528 100644
--- a/compiler/rustc_ast_passes/src/ast_validation.rs
+++ b/compiler/rustc_ast_passes/src/ast_validation.rs
@@ -91,16 +91,18 @@ impl<'a> AstValidator<'a> {
         self.is_impl_trait_banned = old;
     }
 
-    fn with_tilde_const_allowed(&mut self, f: impl FnOnce(&mut Self)) {
-        let old = mem::replace(&mut self.is_tilde_const_allowed, true);
+    fn with_tilde_const(&mut self, allowed: bool, f: impl FnOnce(&mut Self)) {
+        let old = mem::replace(&mut self.is_tilde_const_allowed, allowed);
         f(self);
         self.is_tilde_const_allowed = old;
     }
 
+    fn with_tilde_const_allowed(&mut self, f: impl FnOnce(&mut Self)) {
+        self.with_tilde_const(true, f)
+    }
+
     fn with_banned_tilde_const(&mut self, f: impl FnOnce(&mut Self)) {
-        let old = mem::replace(&mut self.is_tilde_const_allowed, false);
-        f(self);
-        self.is_tilde_const_allowed = old;
+        self.with_tilde_const(false, f)
     }
 
     fn with_let_management(
@@ -1202,12 +1204,8 @@ impl<'a> Visitor<'a> for AstValidator<'a> {
                 }
                 self.visit_vis(&item.vis);
                 self.visit_ident(item.ident);
-                if let Const::Yes(_) = sig.header.constness {
-                    self.with_tilde_const_allowed(|this| this.visit_generics(generics));
-                } else {
-                    self.visit_generics(generics);
-                }
-                let kind = FnKind::Fn(FnCtxt::Free, item.ident, sig, &item.vis, body.as_deref());
+                let kind =
+                    FnKind::Fn(FnCtxt::Free, item.ident, sig, &item.vis, generics, body.as_deref());
                 self.visit_fn(kind, item.span, item.id);
                 walk_list!(self, visit_attribute, &item.attrs);
                 return; // Avoid visiting again.
@@ -1555,13 +1553,14 @@ impl<'a> Visitor<'a> for AstValidator<'a> {
             FnSig { span: sig_span, header: FnHeader { ext: Extern::Implicit, .. }, .. },
             _,
             _,
+            _,
         ) = fk
         {
             self.maybe_lint_missing_abi(*sig_span, id);
         }
 
         // Functions without bodies cannot have patterns.
-        if let FnKind::Fn(ctxt, _, sig, _, None) = fk {
+        if let FnKind::Fn(ctxt, _, sig, _, _, None) = fk {
             Self::check_decl_no_pat(&sig.decl, |span, ident, mut_ident| {
                 let (code, msg, label) = match ctxt {
                     FnCtxt::Foreign => (
@@ -1596,7 +1595,11 @@ impl<'a> Visitor<'a> for AstValidator<'a> {
             });
         }
 
-        visit::walk_fn(self, fk, span);
+        let tilde_const_allowed =
+            matches!(fk.header(), Some(FnHeader { constness: Const::Yes(_), .. }))
+                || matches!(fk.ctxt(), Some(FnCtxt::Assoc(_)));
+
+        self.with_tilde_const(tilde_const_allowed, |this| visit::walk_fn(this, fk, span));
     }
 
     fn visit_assoc_item(&mut self, item: &'a AssocItem, ctxt: AssocCtxt) {
@@ -1670,9 +1673,14 @@ impl<'a> Visitor<'a> for AstValidator<'a> {
             {
                 self.visit_vis(&item.vis);
                 self.visit_ident(item.ident);
-                self.with_tilde_const_allowed(|this| this.visit_generics(generics));
-                let kind =
-                    FnKind::Fn(FnCtxt::Assoc(ctxt), item.ident, sig, &item.vis, body.as_deref());
+                let kind = FnKind::Fn(
+                    FnCtxt::Assoc(ctxt),
+                    item.ident,
+                    sig,
+                    &item.vis,
+                    generics,
+                    body.as_deref(),
+                );
                 self.visit_fn(kind, item.span, item.id);
             }
             _ => self
diff --git a/compiler/rustc_lint/src/builtin.rs b/compiler/rustc_lint/src/builtin.rs
index e6f9246d732..77fe76af2de 100644
--- a/compiler/rustc_lint/src/builtin.rs
+++ b/compiler/rustc_lint/src/builtin.rs
@@ -441,6 +441,7 @@ impl EarlyLintPass for UnsafeCode {
             _,
             ast::FnSig { header: ast::FnHeader { unsafety: ast::Unsafe::Yes(_), .. }, .. },
             _,
+            _,
             body,
         ) = fk
         {
diff --git a/compiler/rustc_lint/src/early.rs b/compiler/rustc_lint/src/early.rs
index 7447f9f64b7..e7855f8057c 100644
--- a/compiler/rustc_lint/src/early.rs
+++ b/compiler/rustc_lint/src/early.rs
@@ -158,7 +158,7 @@ impl<'a, T: EarlyLintPass> ast_visit::Visitor<'a> for EarlyContextAndPass<'a, T>
 
         // Explicitly check for lints associated with 'closure_id', since
         // it does not have a corresponding AST node
-        if let ast_visit::FnKind::Fn(_, _, sig, _, _) = fk {
+        if let ast_visit::FnKind::Fn(_, _, sig, _, _, _) = fk {
             if let ast::Async::Yes { closure_id, .. } = sig.header.asyncness {
                 self.check_id(closure_id);
             }
diff --git a/compiler/rustc_resolve/src/def_collector.rs b/compiler/rustc_resolve/src/def_collector.rs
index 8ea5dca6f10..1e8cca6122c 100644
--- a/compiler/rustc_resolve/src/def_collector.rs
+++ b/compiler/rustc_resolve/src/def_collector.rs
@@ -133,8 +133,10 @@ impl<'a, 'b> visit::Visitor<'a> for DefCollector<'a, 'b> {
     }
 
     fn visit_fn(&mut self, fn_kind: FnKind<'a>, span: Span, _: NodeId) {
-        if let FnKind::Fn(_, _, sig, _, body) = fn_kind {
+        if let FnKind::Fn(_, _, sig, _, generics, body) = fn_kind {
             if let Async::Yes { closure_id, return_impl_trait_id, .. } = sig.header.asyncness {
+                self.visit_generics(generics);
+
                 let return_impl_trait_id =
                     self.create_def(return_impl_trait_id, DefPathData::ImplTrait, span);
 
diff --git a/compiler/rustc_resolve/src/late.rs b/compiler/rustc_resolve/src/late.rs
index 2f6ac4020bc..383ca958e6d 100644
--- a/compiler/rustc_resolve/src/late.rs
+++ b/compiler/rustc_resolve/src/late.rs
@@ -537,10 +537,12 @@ impl<'a: 'ast, 'ast> Visitor<'ast> for LateResolutionVisitor<'a, '_, 'ast> {
         let rib_kind = match fn_kind {
             // Bail if the function is foreign, and thus cannot validly have
             // a body, or if there's no body for some other reason.
-            FnKind::Fn(FnCtxt::Foreign, _, sig, ..) | FnKind::Fn(_, _, sig, .., None) => {
+            FnKind::Fn(FnCtxt::Foreign, _, sig, _, generics, _)
+            | FnKind::Fn(_, _, sig, _, generics, None) => {
                 // We don't need to deal with patterns in parameters, because
                 // they are not possible for foreign or bodiless functions.
                 self.visit_fn_header(&sig.header);
+                self.visit_generics(generics);
                 visit::walk_fn_decl(self, &sig.decl);
                 return;
             }
@@ -559,6 +561,10 @@ impl<'a: 'ast, 'ast> Visitor<'ast> for LateResolutionVisitor<'a, '_, 'ast> {
         self.with_rib(ValueNS, rib_kind, |this| {
             // Create a label rib for the function.
             this.with_label_rib(rib_kind, |this| {
+                if let FnKind::Fn(_, _, _, _, generics, _) = fn_kind {
+                    this.visit_generics(generics);
+                }
+
                 // Add each argument to the rib.
                 this.resolve_params(&declaration.inputs);
 
diff --git a/src/tools/rustfmt/src/items.rs b/src/tools/rustfmt/src/items.rs
index 92f423bbb62..ad2502b0418 100644
--- a/src/tools/rustfmt/src/items.rs
+++ b/src/tools/rustfmt/src/items.rs
@@ -204,12 +204,11 @@ impl<'a> FnSig<'a> {
 
     pub(crate) fn from_fn_kind(
         fn_kind: &'a visit::FnKind<'_>,
-        generics: &'a ast::Generics,
         decl: &'a ast::FnDecl,
         defaultness: ast::Defaultness,
     ) -> FnSig<'a> {
         match *fn_kind {
-            visit::FnKind::Fn(fn_ctxt, _, fn_sig, vis, _) => match fn_ctxt {
+            visit::FnKind::Fn(fn_ctxt, _, fn_sig, vis, generics, _) => match fn_ctxt {
                 visit::FnCtxt::Assoc(..) => {
                     let mut fn_sig = FnSig::from_method_sig(fn_sig, generics, vis);
                     fn_sig.defaultness = defaultness;
@@ -3180,8 +3179,14 @@ impl Rewrite for ast::ForeignItem {
                     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, Some(body)),
-                        generics,
+                        visit::FnKind::Fn(
+                            fn_ctxt,
+                            self.ident,
+                            sig,
+                            &self.vis,
+                            generics,
+                            Some(body),
+                        ),
                         &sig.decl,
                         self.span,
                         defaultness,
diff --git a/src/tools/rustfmt/src/visitor.rs b/src/tools/rustfmt/src/visitor.rs
index 3ebfa551d1c..1621eb406b1 100644
--- a/src/tools/rustfmt/src/visitor.rs
+++ b/src/tools/rustfmt/src/visitor.rs
@@ -382,7 +382,6 @@ impl<'b, 'a: 'b> FmtVisitor<'a> {
     pub(crate) fn visit_fn(
         &mut self,
         fk: visit::FnKind<'_>,
-        generics: &ast::Generics,
         fd: &ast::FnDecl,
         s: Span,
         defaultness: ast::Defaultness,
@@ -391,12 +390,12 @@ impl<'b, 'a: 'b> FmtVisitor<'a> {
         let indent = self.block_indent;
         let block;
         let rewrite = match fk {
-            visit::FnKind::Fn(_, ident, _, _, Some(ref b)) => {
+            visit::FnKind::Fn(_, ident, _, _, _, Some(ref b)) => {
                 block = b;
                 self.rewrite_fn_before_block(
                     indent,
                     ident,
-                    &FnSig::from_fn_kind(&fk, generics, fd, defaultness),
+                    &FnSig::from_fn_kind(&fk, fd, defaultness),
                     mk_sp(s.lo(), b.span.lo()),
                 )
             }
@@ -552,8 +551,14 @@ impl<'b, 'a: 'b> FmtVisitor<'a> {
                             _ => visit::FnCtxt::Foreign,
                         };
                         self.visit_fn(
-                            visit::FnKind::Fn(fn_ctxt, item.ident, sig, &item.vis, Some(body)),
-                            generics,
+                            visit::FnKind::Fn(
+                                fn_ctxt,
+                                item.ident,
+                                sig,
+                                &item.vis,
+                                generics,
+                                Some(body),
+                            ),
                             &sig.decl,
                             item.span,
                             defaultness,
@@ -642,8 +647,7 @@ impl<'b, 'a: 'b> FmtVisitor<'a> {
                     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, Some(body)),
-                        generics,
+                        visit::FnKind::Fn(fn_ctxt, ai.ident, sig, &ai.vis, generics, Some(body)),
                         &sig.decl,
                         ai.span,
                         defaultness,