handle more cases

3 files changed, 128 insertions, 98 deletions

diff --git a/crates/ide-completion/src/completions/type.rs b/crates/ide-completion/src/completions/type.rs
index de855242ee3..69b1e1fd11f 100644
--- a/crates/ide-completion/src/completions/type.rs
+++ b/crates/ide-completion/src/completions/type.rs
@@ -173,7 +173,7 @@ pub(crate) fn complete_type_path(
                     });
                     return;
                 }
-                TypeLocation::GenericArgList(Some((arg_list, generic_param))) => {
+                TypeLocation::GenericArgList(Some((arg_list, _))) => {
                     let in_assoc_type_arg = ctx
                         .original_token
                         .parent_ancestors()
diff --git a/crates/ide-completion/src/context/analysis.rs b/crates/ide-completion/src/context/analysis.rs
index 7fd3147a739..4bd429c5723 100644
--- a/crates/ide-completion/src/context/analysis.rs
+++ b/crates/ide-completion/src/context/analysis.rs
@@ -726,20 +726,30 @@ fn classify_name_ref(
             arg.syntax().parent().and_then(ast::GenericArgList::cast),
         )
         .map(|args| {
-            // Determine the index of the parameter in the `GenericArgList`
-            // (subtract 1 because `siblings` includes the node itself)
-            let param_idx = arg.syntax().siblings(Direction::Prev).count() - 1;
-            let parent = args.syntax().parent();
-            let param = parent.and_then(|parent| {
-                match_ast! {
+            let param = (|| {
+                let parent = args.syntax().parent()?;
+                let params = match_ast! {
                     match parent {
                         ast::PathSegment(segment) => {
                             match sema.resolve_path(&segment.parent_path().top_path())? {
                                 hir::PathResolution::Def(def) => match def {
                                     hir::ModuleDef::Function(func) => {
-                                        let src = func.source(sema.db)?;
-                                        let params = src.value.generic_param_list()?;
-                                        params.generic_params().nth(param_idx)
+                                        func.source(sema.db)?.value.generic_param_list()
+                                    }
+                                    hir::ModuleDef::Adt(adt) => {
+                                        adt.source(sema.db)?.value.generic_param_list()
+                                    }
+                                    hir::ModuleDef::Variant(variant) => {
+                                        variant.parent_enum(sema.db).source(sema.db)?.value.generic_param_list()
+                                    }
+                                    hir::ModuleDef::Trait(trait_) => {
+                                        trait_.source(sema.db)?.value.generic_param_list()
+                                    }
+                                    hir::ModuleDef::TraitAlias(trait_) => {
+                                        trait_.source(sema.db)?.value.generic_param_list()
+                                    }
+                                    hir::ModuleDef::TypeAlias(ty_) => {
+                                        ty_.source(sema.db)?.value.generic_param_list()
                                     }
                                     _ => None,
                                 },
@@ -748,9 +758,7 @@ fn classify_name_ref(
                         },
                         ast::MethodCallExpr(call) => {
                             let func = sema.resolve_method_call(&call)?;
-                            let src = func.source(sema.db)?;
-                            let params = src.value.generic_param_list()?;
-                            params.generic_params().nth(param_idx)
+                            func.source(sema.db)?.value.generic_param_list()
                         },
                         ast::AssocTypeArg(arg) => {
                             let trait_ = ast::PathSegment::cast(arg.syntax().parent()?.parent()?)?;
@@ -765,9 +773,7 @@ fn classify_name_ref(
                                             },
                                             _ => None,
                                         })?;
-                                        let src = assoc_ty.source(sema.db)?;
-                                        let params = src.value.generic_param_list()?;
-                                        params.generic_params().nth(param_idx)
+                                        assoc_ty.source(sema.db)?.value.generic_param_list()
                                     }
                                     _ => None,
                                 },
@@ -776,8 +782,12 @@ fn classify_name_ref(
                         },
                         _ => None,
                     }
-                }
-            });
+                }?;
+                // Determine the index of the parameter in the `GenericArgList`
+                // (subtract 1 because `siblings` includes the node itself)
+                let param_idx = arg.syntax().siblings(Direction::Prev).count() - 1;
+                params.generic_params().nth(param_idx)
+            })();
             (args, param)
         });
         TypeLocation::GenericArgList(location)
diff --git a/crates/ide-completion/src/tests/type_pos.rs b/crates/ide-completion/src/tests/type_pos.rs
index 67110b3d7fb..27bcf4c630a 100644
--- a/crates/ide-completion/src/tests/type_pos.rs
+++ b/crates/ide-completion/src/tests/type_pos.rs
@@ -722,16 +722,31 @@ pub struct S;
 
 #[test]
 fn completes_const_and_type_generics_separately() {
-    check(
+    let type_completion_cases = [
+        // Function
         r#"
-    struct Foo;
-    const X: usize = 0;
-    fn foo<T, const N: usize>() {}
-    fn main() {
-        foo::<F$0, _>();
-    }
-    "#,
-        expect![[r#"
+struct Foo;
+const X: usize = 0;
+fn foo<T, const N: usize>() {}
+fn main() {
+    foo::<F$0, _>();
+}
+        "#,
+        // Method
+        r#"
+const X: usize = 0;
+struct Foo;
+impl Foo { fn bar<const N: usize, T>(self) {} }
+fn main() {
+    Foo.bar::<_, $0>();
+}
+        "#,
+    ];
+
+    for case in type_completion_cases {
+        check(
+            case,
+            expect![[r#"
                 en Enum
                 ma makro!(…) macro_rules! makro
                 md module
@@ -745,51 +760,18 @@ fn completes_const_and_type_generics_separately() {
                 kw crate::
                 kw self::
             "#]],
-    );
-    check(
-        r#"
-    struct Foo;
-    const X: usize = 0;
-    fn foo<T, const N: usize>() {}
-    fn main() {
-        foo::<_, $0>();
+        )
     }
-    "#,
-        expect![[r#"
-                ct CONST
-                ct X
-                ma makro!(…) macro_rules! makro
-                kw crate::
-                kw self::
-            "#]],
-    );
 
     check(
         r#"
 const X: usize = 0;
 struct Foo;
-impl Foo { fn bar<const N: usize, T>(self) {} }
-fn main() {
-    Foo.bar::<X$0, _>();
-}
-"#,
-        expect![[r#"
-            ct CONST
-            ct X
-            ma makro!(…) macro_rules! makro
-            kw crate::
-            kw self::
-        "#]],
-    );
-    check(
-        r#"
-const X: usize = 0;
-struct Foo;
-impl Foo { fn bar<const N: usize, T>(self) {} }
-fn main() {
-    Foo.bar::<_, $0>();
+trait Bar {
+    type Baz<T, const X: usize>;
 }
-"#,
+fn foo(_: impl Bar<Baz<F$0, 0> = ()>) {}
+        "#,
         expect![[r#"
             en Enum
             ma makro!(…) macro_rules! makro
@@ -798,6 +780,7 @@ fn main() {
             st Record
             st Tuple
             st Unit
+            tt Bar
             tt Trait
             un Union
             bt u32
@@ -806,47 +789,84 @@ fn main() {
         "#]],
     );
 
-    check(
+    let const_completion_cases = [
+        // Function params
+        r#"
+struct Foo;
+const X: usize = 0;
+fn foo<T, const N: usize>() {}
+fn main() {
+    foo::<_, $0>();
+}
+        "#,
+        // Method params
         r#"
 const X: usize = 0;
 struct Foo;
-trait Bar {
-    type Baz<T, const X: usize>;
+impl Foo { fn bar<const N: usize, T>(self) {} }
+fn main() {
+    Foo.bar::<X$0, _>();
 }
-fn foo<T: Bar<Baz<(), $0> = ()>>() {}
-"#,
-        expect![[r#"
-            ct CONST
-            ct X
-            ma makro!(…) macro_rules! makro
-            kw crate::
-            kw self::
-        "#]],
-    );
-    check(
+        "#,
+        // Associated type params
         r#"
 const X: usize = 0;
 struct Foo;
 trait Bar {
     type Baz<T, const X: usize>;
 }
-fn foo<T: Bar<Baz<F$0, 0> = ()>>() {}
-"#,
-        expect![[r#"
-            en Enum
-            ma makro!(…) macro_rules! makro
-            md module
-            st Foo
-            st Record
-            st Tuple
-            st Unit
-            tt Bar
-            tt Trait
-            tp T
-            un Union
-            bt u32
-            kw crate::
-            kw self::
-        "#]],
-    );
+fn foo<T: Bar<Baz<(), $0> = ()>>() {}
+        "#,
+        // Type params
+        r#"
+const X: usize = 0;
+struct Foo<T, const N: usize>(T);
+fn main() {
+    let _: Foo::<_, $0> = todo!();
+}
+        "#,
+        // Enum variant params
+        r#"
+const X: usize = 0;
+struct Foo<T, const N: usize>(T);
+type Bar<const X: usize, U> = Foo<U, X>;
+fn main() {
+    let _: Bar::<X$0, _> = todo!();
+}
+        "#,
+        r#"
+const X: usize = 0;
+enum Foo<T, const N: usize> { A(T), B }
+fn main() {
+    Foo::B::<(), $0>;
+}
+        "#,
+        // Trait params
+        r#"
+const X: usize = 0;
+trait Foo<T, const N: usize> {}
+impl Foo<(), $0> for () {}
+        "#,
+        // Trait alias params
+        r#"
+#![feature(trait_alias)]
+const X: usize = 0;
+trait Foo<T, const N: usize> {}
+trait Bar<const M: usize, U> = Foo<U, M>;
+fn foo<T: Bar<X$0, ()>>() {}
+        "#,
+    ];
+
+    for case in const_completion_cases {
+        check(
+            case,
+            expect![[r#"
+                ct CONST
+                ct X
+                ma makro!(…) macro_rules! makro
+                kw crate::
+                kw self::
+            "#]],
+        );
+    }
 }