support AssocTypeArg and MethodCalls

2 files changed, 158 insertions, 34 deletions

diff --git a/crates/ide-completion/src/context/analysis.rs b/crates/ide-completion/src/context/analysis.rs
index cc9e4581bed..7fd3147a739 100644
--- a/crates/ide-completion/src/context/analysis.rs
+++ b/crates/ide-completion/src/context/analysis.rs
@@ -719,6 +719,70 @@ fn classify_name_ref(
         None
     };
 
+    let generic_arg_location = |arg: ast::GenericArg| {
+        let location = find_opt_node_in_file_compensated(
+            sema,
+            original_file,
+            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! {
+                    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)
+                                    }
+                                    _ => None,
+                                },
+                                _ => None,
+                            }
+                        },
+                        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)
+                        },
+                        ast::AssocTypeArg(arg) => {
+                            let trait_ = ast::PathSegment::cast(arg.syntax().parent()?.parent()?)?;
+                            match sema.resolve_path(&trait_.parent_path().top_path())? {
+                                hir::PathResolution::Def(def) => match def {
+                                    hir::ModuleDef::Trait(trait_) => {
+                                        let trait_items = trait_.items(sema.db);
+                                        let assoc_ty = trait_items.iter().find_map(|item| match item {
+                                            hir::AssocItem::TypeAlias(assoc_ty) => {
+                                                (assoc_ty.name(sema.db).as_str()? == arg.name_ref()?.text())
+                                                    .then_some(assoc_ty)
+                                            },
+                                            _ => None,
+                                        })?;
+                                        let src = assoc_ty.source(sema.db)?;
+                                        let params = src.value.generic_param_list()?;
+                                        params.generic_params().nth(param_idx)
+                                    }
+                                    _ => None,
+                                },
+                                _ => None,
+                            }
+                        },
+                        _ => None,
+                    }
+                }
+            });
+            (args, param)
+        });
+        TypeLocation::GenericArgList(location)
+    };
+
     let type_location = |node: &SyntaxNode| {
         let parent = node.parent()?;
         let res = match_ast! {
@@ -774,34 +838,8 @@ fn classify_name_ref(
                 ast::TypeBound(_) => TypeLocation::TypeBound,
                 // is this case needed?
                 ast::TypeBoundList(_) => TypeLocation::TypeBound,
-                ast::GenericArg(it) => {
-                    let location = find_opt_node_in_file_compensated(sema, original_file, it.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 = it.syntax().siblings(Direction::Prev).count() - 1;
-                            let param = args
-                                .syntax()
-                                .parent()
-                                .and_then(|p| ast::PathSegment::cast(p))
-                                .and_then(|segment| sema.resolve_path(&segment.parent_path().top_path()))
-                                .and_then(|resolved| {
-                                    match resolved {
-                                        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)
-                                            }
-                                            _ => None,
-                                        },
-                                        _ => None,
-                                    }
-                                });
-                            (args, param)
-                        });
-                    TypeLocation::GenericArgList(location)
-                },
+                ast::TypeArg(it) => generic_arg_location(ast::GenericArg::TypeArg(it)),
+                ast::GenericArg(it) => generic_arg_location(it),
                 // is this case needed?
                 ast::GenericArgList(it) => {
                     let location = find_opt_node_in_file_compensated(sema, original_file, Some(it))
diff --git a/crates/ide-completion/src/tests/type_pos.rs b/crates/ide-completion/src/tests/type_pos.rs
index 2a05aef17f9..67110b3d7fb 100644
--- a/crates/ide-completion/src/tests/type_pos.rs
+++ b/crates/ide-completion/src/tests/type_pos.rs
@@ -724,11 +724,70 @@ pub struct S;
 fn completes_const_and_type_generics_separately() {
     check(
         r#"
+    struct Foo;
+    const X: usize = 0;
+    fn foo<T, const N: usize>() {}
+    fn main() {
+        foo::<F$0, _>();
+    }
+    "#,
+        expect![[r#"
+                en Enum
+                ma makro!(…) macro_rules! makro
+                md module
+                st Foo
+                st Record
+                st Tuple
+                st Unit
+                tt Trait
+                un Union
+                bt u32
+                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;
-fn foo<T, const N: usize>() {}
+struct Foo;
+impl Foo { fn bar<const N: usize, T>(self) {} }
 fn main() {
-    foo::<F$0, _>();
+    Foo.bar::<_, $0>();
 }
 "#,
         expect![[r#"
@@ -746,14 +805,15 @@ fn main() {
             kw self::
         "#]],
     );
+
     check(
         r#"
-struct Foo;
 const X: usize = 0;
-fn foo<T, const N: usize>() {}
-fn main() {
-    foo::<_, $0>();
+struct Foo;
+trait Bar {
+    type Baz<T, const X: usize>;
 }
+fn foo<T: Bar<Baz<(), $0> = ()>>() {}
 "#,
         expect![[r#"
             ct CONST
@@ -763,4 +823,30 @@ fn main() {
             kw self::
         "#]],
     );
+    check(
+        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::
+        "#]],
+    );
 }