Auto merge of #12695 - xuhongxu96:fix-12140, r=jonas-schievink

Complete type param/associated type in trait generic arg per arg index

- Fix #12140
- Also fix tidy check does not work for marks in multiline

5 files changed, 347 insertions, 50 deletions

diff --git a/crates/hir-def/src/generics.rs b/crates/hir-def/src/generics.rs
index 0b2e78bdcfb..eec960aa7dd 100644
--- a/crates/hir-def/src/generics.rs
+++ b/crates/hir-def/src/generics.rs
@@ -47,6 +47,7 @@ pub struct LifetimeParamData {
 pub struct ConstParamData {
     pub name: Name,
     pub ty: Interned<TypeRef>,
+    pub has_default: bool,
 }
 
 #[derive(Copy, Clone, PartialEq, Eq, Debug, Hash)]
@@ -70,6 +71,13 @@ impl TypeOrConstParamData {
         }
     }
 
+    pub fn has_default(&self) -> bool {
+        match self {
+            TypeOrConstParamData::TypeParamData(x) => x.default.is_some(),
+            TypeOrConstParamData::ConstParamData(x) => x.has_default,
+        }
+    }
+
     pub fn type_param(&self) -> Option<&TypeParamData> {
         match self {
             TypeOrConstParamData::TypeParamData(x) => Some(x),
@@ -232,7 +240,11 @@ impl GenericParams {
                     let ty = const_param
                         .ty()
                         .map_or(TypeRef::Error, |it| TypeRef::from_ast(lower_ctx, it));
-                    let param = ConstParamData { name, ty: Interned::new(ty) };
+                    let param = ConstParamData {
+                        name,
+                        ty: Interned::new(ty),
+                        has_default: const_param.default_val().is_some(),
+                    };
                     self.type_or_consts.alloc(param.into());
                 }
             }
diff --git a/crates/hir/src/lib.rs b/crates/hir/src/lib.rs
index ef17f2a75e1..96424d087ef 100644
--- a/crates/hir/src/lib.rs
+++ b/crates/hir/src/lib.rs
@@ -41,6 +41,7 @@ use hir_def::{
     adt::{ReprKind, VariantData},
     body::{BodyDiagnostic, SyntheticSyntax},
     expr::{BindingAnnotation, LabelId, Pat, PatId},
+    generics::{TypeOrConstParamData, TypeParamProvenance},
     item_tree::ItemTreeNode,
     lang_item::LangItemTarget,
     nameres::{self, diagnostics::DefDiagnostic},
@@ -1707,6 +1708,26 @@ impl Trait {
     pub fn is_unsafe(&self, db: &dyn HirDatabase) -> bool {
         db.trait_data(self.id).is_unsafe
     }
+
+    pub fn type_or_const_param_count(
+        &self,
+        db: &dyn HirDatabase,
+        count_required_only: bool,
+    ) -> usize {
+        db.generic_params(GenericDefId::from(self.id))
+            .type_or_consts
+            .iter()
+            .filter(|(_, ty)| match ty {
+                TypeOrConstParamData::TypeParamData(ty)
+                    if ty.provenance != TypeParamProvenance::TypeParamList =>
+                {
+                    false
+                }
+                _ => true,
+            })
+            .filter(|(_, ty)| !count_required_only || !ty.has_default())
+            .count()
+    }
 }
 
 impl HasVisibility for Trait {
diff --git a/crates/ide-completion/src/completions/type.rs b/crates/ide-completion/src/completions/type.rs
index 616d8621543..0469d349519 100644
--- a/crates/ide-completion/src/completions/type.rs
+++ b/crates/ide-completion/src/completions/type.rs
@@ -1,7 +1,7 @@
 //! Completion of names from the current scope in type position.
 
 use hir::{HirDisplay, ScopeDef};
-use syntax::{ast, AstNode};
+use syntax::{ast, AstNode, SyntaxKind};
 
 use crate::{
     context::{PathCompletionCtx, Qualified, TypeAscriptionTarget, TypeLocation},
@@ -120,39 +120,83 @@ pub(crate) fn complete_type_path(
         }
         Qualified::Absolute => acc.add_crate_roots(ctx, path_ctx),
         Qualified::No => {
-            acc.add_nameref_keywords_with_colon(ctx);
-            if let TypeLocation::TypeBound = location {
-                ctx.process_all_names(&mut |name, res| {
-                    let add_resolution = match res {
-                        ScopeDef::ModuleDef(hir::ModuleDef::Macro(mac)) => mac.is_fn_like(ctx.db),
-                        ScopeDef::ModuleDef(
-                            hir::ModuleDef::Trait(_) | hir::ModuleDef::Module(_),
-                        ) => true,
-                        _ => false,
-                    };
-                    if add_resolution {
-                        acc.add_path_resolution(ctx, path_ctx, name, res);
-                    }
-                });
-                return;
-            }
-            if let TypeLocation::GenericArgList(Some(arg_list)) = location {
-                if let Some(path_seg) = arg_list.syntax().parent().and_then(ast::PathSegment::cast)
-                {
-                    if path_seg.syntax().ancestors().find_map(ast::TypeBound::cast).is_some() {
-                        if let Some(hir::PathResolution::Def(hir::ModuleDef::Trait(trait_))) =
-                            ctx.sema.resolve_path(&path_seg.parent_path())
+            match location {
+                TypeLocation::TypeBound => {
+                    acc.add_nameref_keywords_with_colon(ctx);
+                    ctx.process_all_names(&mut |name, res| {
+                        let add_resolution = match res {
+                            ScopeDef::ModuleDef(hir::ModuleDef::Macro(mac)) => {
+                                mac.is_fn_like(ctx.db)
+                            }
+                            ScopeDef::ModuleDef(
+                                hir::ModuleDef::Trait(_) | hir::ModuleDef::Module(_),
+                            ) => true,
+                            _ => false,
+                        };
+                        if add_resolution {
+                            acc.add_path_resolution(ctx, path_ctx, name, res);
+                        }
+                    });
+                    return;
+                }
+                TypeLocation::GenericArgList(Some(arg_list)) => {
+                    let in_assoc_type_arg = ctx
+                        .original_token
+                        .parent_ancestors()
+                        .any(|node| node.kind() == SyntaxKind::ASSOC_TYPE_ARG);
+
+                    if !in_assoc_type_arg {
+                        if let Some(path_seg) =
+                            arg_list.syntax().parent().and_then(ast::PathSegment::cast)
                         {
-                            trait_.items_with_supertraits(ctx.sema.db).into_iter().for_each(|it| {
-                                if let hir::AssocItem::TypeAlias(alias) = it {
-                                    cov_mark::hit!(complete_assoc_type_in_generics_list);
-                                    acc.add_type_alias_with_eq(ctx, alias)
+                            if path_seg
+                                .syntax()
+                                .ancestors()
+                                .find_map(ast::TypeBound::cast)
+                                .is_some()
+                            {
+                                if let Some(hir::PathResolution::Def(hir::ModuleDef::Trait(
+                                    trait_,
+                                ))) = ctx.sema.resolve_path(&path_seg.parent_path())
+                                {
+                                    let arg_idx = arg_list
+                                        .generic_args()
+                                        .filter(|arg| {
+                                            arg.syntax().text_range().end()
+                                                < ctx.original_token.text_range().start()
+                                        })
+                                        .count();
+
+                                    let n_required_params =
+                                        trait_.type_or_const_param_count(ctx.sema.db, true);
+                                    if arg_idx >= n_required_params {
+                                        trait_
+                                            .items_with_supertraits(ctx.sema.db)
+                                            .into_iter()
+                                            .for_each(|it| {
+                                                if let hir::AssocItem::TypeAlias(alias) = it {
+                                                    cov_mark::hit!(
+                                                        complete_assoc_type_in_generics_list
+                                                    );
+                                                    acc.add_type_alias_with_eq(ctx, alias);
+                                                }
+                                            });
+
+                                        let n_params =
+                                            trait_.type_or_const_param_count(ctx.sema.db, false);
+                                        if arg_idx >= n_params {
+                                            return; // only show assoc types
+                                        }
+                                    }
                                 }
-                            });
+                            }
                         }
                     }
                 }
-            }
+                _ => {}
+            };
+
+            acc.add_nameref_keywords_with_colon(ctx);
             ctx.process_all_names(&mut |name, def| {
                 if scope_def_applicable(def) {
                     acc.add_path_resolution(ctx, path_ctx, name, def);
diff --git a/crates/ide-completion/src/tests/type_pos.rs b/crates/ide-completion/src/tests/type_pos.rs
index 8943d303b60..fcd4743f74f 100644
--- a/crates/ide-completion/src/tests/type_pos.rs
+++ b/crates/ide-completion/src/tests/type_pos.rs
@@ -380,10 +380,26 @@ trait Trait2: Trait1 {
 fn foo<'lt, T: Trait2<$0>, const CONST_PARAM: usize>(_: T) {}
 "#,
         expect![[r#"
+            ta Foo =  (as Trait2)   type Foo
+            ta Super =  (as Trait1) type Super
+        "#]],
+    );
+    check(
+        r#"
+trait Trait1 {
+    type Super;
+}
+trait Trait2<T>: Trait1 {
+    type Foo;
+}
+
+fn foo<'lt, T: Trait2<$0>, const CONST_PARAM: usize>(_: T) {}
+"#,
+        expect![[r#"
             ct CONST
             cp CONST_PARAM
             en Enum
-            ma makro!(…)            macro_rules! makro
+            ma makro!(…)   macro_rules! makro
             md module
             st Record
             st Tuple
@@ -391,8 +407,6 @@ fn foo<'lt, T: Trait2<$0>, const CONST_PARAM: usize>(_: T) {}
             tt Trait
             tt Trait1
             tt Trait2
-            ta Foo =  (as Trait2)   type Foo
-            ta Super =  (as Trait1) type Super
             tp T
             un Union
             bt u32
@@ -472,3 +486,206 @@ fn func(_: Enum::$0) {}
         "#]],
     );
 }
+
+#[test]
+fn completes_type_parameter_or_associated_type() {
+    check(
+        r#"
+trait MyTrait<T, U> {
+    type Item1;
+    type Item2;
+};
+
+fn f(t: impl MyTrait<u$0
+"#,
+        expect![[r#"
+            ct CONST
+            en Enum
+            ma makro!(…) macro_rules! makro
+            md module
+            st Record
+            st Tuple
+            st Unit
+            tt MyTrait
+            tt Trait
+            un Union
+            bt u32
+            kw crate::
+            kw self::
+            kw super::
+        "#]],
+    );
+
+    check(
+        r#"
+trait MyTrait<T, U> {
+    type Item1;
+    type Item2;
+};
+
+fn f(t: impl MyTrait<u8, u$0
+"#,
+        expect![[r#"
+            ct CONST
+            en Enum
+            ma makro!(…) macro_rules! makro
+            md module
+            st Record
+            st Tuple
+            st Unit
+            tt MyTrait
+            tt Trait
+            un Union
+            bt u32
+            kw crate::
+            kw self::
+            kw super::
+        "#]],
+    );
+
+    check(
+        r#"
+trait MyTrait<T, U> {
+    type Item1;
+    type Item2;
+};
+
+fn f(t: impl MyTrait<u8, u8, I$0
+"#,
+        expect![[r#"
+            ta Item1 =  (as MyTrait) type Item1
+            ta Item2 =  (as MyTrait) type Item2
+        "#]],
+    );
+}
+
+#[test]
+fn completes_type_parameter_or_associated_type_with_default_value() {
+    check(
+        r#"
+trait MyTrait<T, U = u8> {
+    type Item1;
+    type Item2;
+};
+
+fn f(t: impl MyTrait<u$0
+"#,
+        expect![[r#"
+            ct CONST
+            en Enum
+            ma makro!(…) macro_rules! makro
+            md module
+            st Record
+            st Tuple
+            st Unit
+            tt MyTrait
+            tt Trait
+            un Union
+            bt u32
+            kw crate::
+            kw self::
+            kw super::
+        "#]],
+    );
+
+    check(
+        r#"
+trait MyTrait<T, U = u8> {
+    type Item1;
+    type Item2;
+};
+
+fn f(t: impl MyTrait<u8, u$0
+"#,
+        expect![[r#"
+            ct CONST
+            en Enum
+            ma makro!(…)             macro_rules! makro
+            md module
+            st Record
+            st Tuple
+            st Unit
+            tt MyTrait
+            tt Trait
+            ta Item1 =  (as MyTrait) type Item1
+            ta Item2 =  (as MyTrait) type Item2
+            un Union
+            bt u32
+            kw crate::
+            kw self::
+            kw super::
+        "#]],
+    );
+
+    check(
+        r#"
+trait MyTrait<T, U = u8> {
+    type Item1;
+    type Item2;
+};
+
+fn f(t: impl MyTrait<u8, u8, I$0
+"#,
+        expect![[r#"
+            ta Item1 =  (as MyTrait) type Item1
+            ta Item2 =  (as MyTrait) type Item2
+        "#]],
+    );
+}
+
+#[test]
+fn completes_types_after_associated_type() {
+    check(
+        r#"
+trait MyTrait {
+    type Item1;
+    type Item2;
+};
+
+fn f(t: impl MyTrait<Item1 = $0
+"#,
+        expect![[r#"
+            ct CONST
+            en Enum
+            ma makro!(…) macro_rules! makro
+            md module
+            st Record
+            st Tuple
+            st Unit
+            tt MyTrait
+            tt Trait
+            un Union
+            bt u32
+            kw crate::
+            kw self::
+            kw super::
+        "#]],
+    );
+
+    check(
+        r#"
+trait MyTrait {
+    type Item1;
+    type Item2;
+};
+
+fn f(t: impl MyTrait<Item1 = u8, Item2 = $0
+"#,
+        expect![[r#"
+            ct CONST
+            en Enum
+            ma makro!(…) macro_rules! makro
+            md module
+            st Record
+            st Tuple
+            st Unit
+            tt MyTrait
+            tt Trait
+            un Union
+            bt u32
+            kw crate::
+            kw self::
+            kw super::
+        "#]],
+    );
+}
diff --git a/crates/rust-analyzer/tests/slow-tests/tidy.rs b/crates/rust-analyzer/tests/slow-tests/tidy.rs
index 02ab6081296..12be7947fe0 100644
--- a/crates/rust-analyzer/tests/slow-tests/tidy.rs
+++ b/crates/rust-analyzer/tests/slow-tests/tidy.rs
@@ -471,17 +471,9 @@ struct TidyMarks {
 
 impl TidyMarks {
     fn visit(&mut self, _path: &Path, text: &str) {
-        for line in text.lines() {
-            if let Some(mark) = find_mark(line, "hit") {
-                self.hits.insert(mark.to_string());
-            }
-            if let Some(mark) = find_mark(line, "check") {
-                self.checks.insert(mark.to_string());
-            }
-            if let Some(mark) = find_mark(line, "check_count") {
-                self.checks.insert(mark.to_string());
-            }
-        }
+        find_marks(&mut self.hits, text, "hit");
+        find_marks(&mut self.checks, text, "check");
+        find_marks(&mut self.checks, text, "check_count");
     }
 
     fn finish(self) {
@@ -506,10 +498,21 @@ fn stable_hash(text: &str) -> u64 {
     hasher.finish()
 }
 
-fn find_mark<'a>(text: &'a str, mark: &'static str) -> Option<&'a str> {
-    let idx = text.find(mark)?;
-    let text = text[idx + mark.len()..].strip_prefix("!(")?;
-    let idx = text.find(|c: char| !(c.is_alphanumeric() || c == '_'))?;
-    let text = &text[..idx];
-    Some(text)
+fn find_marks(set: &mut HashSet<String>, text: &str, mark: &str) {
+    let mut text = text;
+    let mut prev_text = "";
+    while text != prev_text {
+        prev_text = text;
+        if let Some(idx) = text.find(mark) {
+            text = &text[idx + mark.len()..];
+            if let Some(stripped_text) = text.strip_prefix("!(") {
+                text = stripped_text.trim_start();
+                if let Some(idx2) = text.find(|c: char| !(c.is_alphanumeric() || c == '_')) {
+                    let mark_text = &text[..idx2];
+                    set.insert(mark_text.to_string());
+                    text = &text[idx2..];
+                }
+            }
+        }
+    }
 }