Change how suggested lifetime args are computed.

1 files changed, 48 insertions, 20 deletions

diff --git a/compiler/rustc_hir_analysis/src/structured_errors/wrong_number_of_generic_args.rs b/compiler/rustc_hir_analysis/src/structured_errors/wrong_number_of_generic_args.rs
index 9c77387c238..4451db19f5c 100644
--- a/compiler/rustc_hir_analysis/src/structured_errors/wrong_number_of_generic_args.rs
+++ b/compiler/rustc_hir_analysis/src/structured_errors/wrong_number_of_generic_args.rs
@@ -296,25 +296,35 @@ impl<'a, 'tcx> WrongNumberOfGenericArgs<'a, 'tcx> {
     ) -> String {
         debug!(?path_hir_id);
 
+        // If there was already a lifetime among the arguments, just replicate that one.
+        if let Some(lt) = self.gen_args.args.iter().find_map(|arg| match arg {
+            hir::GenericArg::Lifetime(lt) => Some(lt),
+            _ => None,
+        }) {
+            return std::iter::repeat(lt.to_string())
+                .take(num_params_to_take)
+                .collect::<Vec<_>>()
+                .join(", ");
+        }
+
         let mut ret = Vec::new();
+        let mut ty_id = None;
         for (id, node) in self.tcx.hir().parent_iter(path_hir_id) {
             debug!(?id);
-            let params = if let Some(generics) = node.generics() {
-                generics.params
-            } else if let hir::Node::Ty(ty) = node
-                && let hir::TyKind::BareFn(bare_fn) = ty.kind
-            {
-                bare_fn.generic_params
-            } else {
-                &[]
-            };
-            ret.extend(params.iter().filter_map(|p| {
-                let hir::GenericParamKind::Lifetime { kind: hir::LifetimeParamKind::Explicit }
-                    = p.kind
-                else { return None };
-                let hir::ParamName::Plain(name) = p.name else { return None };
-                Some(name.to_string())
-            }));
+            if let hir::Node::Ty(_) = node {
+                ty_id = Some(id);
+            }
+
+            // Suggest `'_` when in function parameter or elided function return.
+            if let Some(fn_decl) = node.fn_decl() && let Some(ty_id) = ty_id {
+                let in_arg = fn_decl.inputs.iter().any(|t| t.hir_id == ty_id);
+                let in_ret = matches!(fn_decl.output, hir::FnRetTy::Return(ty) if ty.hir_id == ty_id);
+
+                if in_arg || (in_ret && fn_decl.lifetime_elision_allowed) {
+                    return std::iter::repeat("'_".to_owned()).take(num_params_to_take).collect::<Vec<_>>().join(", ");
+                }
+            }
+
             // Suggest `'static` when in const/static item-like.
             if let hir::Node::Item(hir::Item {
                 kind: hir::ItemKind::Static { .. } | hir::ItemKind::Const { .. },
@@ -334,11 +344,29 @@ impl<'a, 'tcx> WrongNumberOfGenericArgs<'a, 'tcx> {
             })
             | hir::Node::AnonConst(..) = node
             {
-                ret.extend(
-                    std::iter::repeat("'static".to_owned())
-                        .take(num_params_to_take.saturating_sub(ret.len())),
-                );
+                return std::iter::repeat("'static".to_owned())
+                    .take(num_params_to_take.saturating_sub(ret.len()))
+                    .collect::<Vec<_>>()
+                    .join(", ");
             }
+
+            let params = if let Some(generics) = node.generics() {
+                generics.params
+            } else if let hir::Node::Ty(ty) = node
+                && let hir::TyKind::BareFn(bare_fn) = ty.kind
+            {
+                bare_fn.generic_params
+            } else {
+                &[]
+            };
+            ret.extend(params.iter().filter_map(|p| {
+                let hir::GenericParamKind::Lifetime { kind: hir::LifetimeParamKind::Explicit }
+                    = p.kind
+                else { return None };
+                let hir::ParamName::Plain(name) = p.name else { return None };
+                Some(name.to_string())
+            }));
+
             if ret.len() >= num_params_to_take {
                 return ret[..num_params_to_take].join(", ");
             }