Merge pull request #20812 from A4-Tacks/comp-trait-self-param

Add self param completions for trait assoc fn

Author: ChayimFriedman2

crates/ide-completion/src/completions.rs

@@ -630,7 +630,7 @@ fn enum_variants_with_paths(
     acc: &mut Completions,
     ctx: &CompletionContext<'_>,
     enum_: hir::Enum,
-    impl_: &Option<ast::Impl>,
+    impl_: Option<&ast::Impl>,
     cb: impl Fn(&mut Completions, &CompletionContext<'_>, hir::Variant, hir::ModPath),
 ) {
     let mut process_variant = |variant: Variant| {
@@ -644,7 +644,7 @@ fn enum_variants_with_paths(
 
     let variants = enum_.variants(ctx.db);
 
-    if let Some(impl_) = impl_.as_ref().and_then(|impl_| ctx.sema.to_def(impl_))
+    if let Some(impl_) = impl_.and_then(|impl_| ctx.sema.to_def(impl_))
         && impl_.self_ty(ctx.db).as_adt() == Some(hir::Adt::Enum(enum_))
     {
         variants.iter().for_each(|variant| process_variant(*variant));

crates/ide-completion/src/completions/expr.rs

@@ -297,7 +297,7 @@ pub(crate) fn complete_expr_path(
                             acc,
                             ctx,
                             e,
-                            impl_,
+                            impl_.as_ref(),
                             |acc, ctx, variant, path| {
                                 acc.add_qualified_enum_variant(ctx, path_ctx, variant, path)
                             },

crates/ide-completion/src/completions/fn_param.rs

@@ -2,6 +2,7 @@
 
 use hir::HirDisplay;
 use ide_db::FxHashMap;
+use itertools::Either;
 use syntax::{
     AstNode, Direction, SyntaxKind, TextRange, TextSize, algo,
     ast::{self, HasModuleItem},
@@ -24,8 +25,8 @@ pub(crate) fn complete_fn_param(
     ctx: &CompletionContext<'_>,
     pattern_ctx: &PatternContext,
 ) -> Option<()> {
-    let (ParamContext { param_list, kind, .. }, impl_) = match pattern_ctx {
-        PatternContext { param_ctx: Some(kind), impl_, .. } => (kind, impl_),
+    let (ParamContext { param_list, kind, .. }, impl_or_trait) = match pattern_ctx {
+        PatternContext { param_ctx: Some(kind), impl_or_trait, .. } => (kind, impl_or_trait),
         _ => return None,
     };
 
@@ -45,7 +46,7 @@ pub(crate) fn complete_fn_param(
 
     match kind {
         ParamKind::Function(function) => {
-            fill_fn_params(ctx, function, param_list, impl_, add_new_item_to_acc);
+            fill_fn_params(ctx, function, param_list, impl_or_trait, add_new_item_to_acc);
         }
         ParamKind::Closure(closure) => {
             let stmt_list = closure.syntax().ancestors().find_map(ast::StmtList::cast)?;
@@ -62,7 +63,7 @@ fn fill_fn_params(
     ctx: &CompletionContext<'_>,
     function: &ast::Fn,
     param_list: &ast::ParamList,
-    impl_: &Option<ast::Impl>,
+    impl_or_trait: &Option<Either<ast::Impl, ast::Trait>>,
     mut add_new_item_to_acc: impl FnMut(&str),
 ) {
     let mut file_params = FxHashMap::default();
@@ -107,7 +108,7 @@ fn fill_fn_params(
     }
     remove_duplicated(&mut file_params, param_list.params());
     let self_completion_items = ["self", "&self", "mut self", "&mut self"];
-    if should_add_self_completions(ctx.token.text_range().start(), param_list, impl_) {
+    if should_add_self_completions(ctx.token.text_range().start(), param_list, impl_or_trait) {
         self_completion_items.into_iter().for_each(&mut add_new_item_to_acc);
     }
 
@@ -161,9 +162,9 @@ fn remove_duplicated(
 fn should_add_self_completions(
     cursor: TextSize,
     param_list: &ast::ParamList,
-    impl_: &Option<ast::Impl>,
+    impl_or_trait: &Option<Either<ast::Impl, ast::Trait>>,
 ) -> bool {
-    if impl_.is_none() || param_list.self_param().is_some() {
+    if impl_or_trait.is_none() || param_list.self_param().is_some() {
         return false;
     }
     match param_list.params().next() {

crates/ide-completion/src/completions/pattern.rs

@@ -70,7 +70,7 @@ pub(crate) fn complete_pattern(
             acc,
             ctx,
             e,
-            &pattern_ctx.impl_,
+            pattern_ctx.impl_or_trait.as_ref().and_then(|it| it.as_ref().left()),
             |acc, ctx, variant, path| {
                 acc.add_qualified_variant_pat(ctx, pattern_ctx, variant, path);
             },

crates/ide-completion/src/context.rs

@@ -15,6 +15,7 @@ use ide_db::{
     FilePosition, FxHashMap, FxHashSet, RootDatabase, famous_defs::FamousDefs,
     helpers::is_editable_crate,
 };
+use itertools::Either;
 use syntax::{
     AstNode, Edition, SmolStr,
     SyntaxKind::{self, *},
@@ -282,7 +283,7 @@ pub(crate) struct PatternContext {
     pub(crate) mut_token: Option<SyntaxToken>,
     /// The record pattern this name or ref is a field of
     pub(crate) record_pat: Option<ast::RecordPat>,
-    pub(crate) impl_: Option<ast::Impl>,
+    pub(crate) impl_or_trait: Option<Either<ast::Impl, ast::Trait>>,
     /// List of missing variants in a match expr
     pub(crate) missing_variants: Vec<hir::Variant>,
 }

crates/ide-completion/src/context/analysis.rs

@@ -1302,7 +1302,8 @@ fn classify_name_ref<'db>(
             .is_some_and(|it| it.semicolon_token().is_none())
             || after_incomplete_let && incomplete_expr_stmt.unwrap_or(true) && !before_else_kw;
         let in_value = it.parent().and_then(Either::<ast::LetStmt, ast::ArgList>::cast).is_some();
-        let impl_ = fetch_immediate_impl(sema, original_file, expr.syntax());
+        let impl_ = fetch_immediate_impl_or_trait(sema, original_file, expr.syntax())
+            .and_then(Either::left);
 
         let in_match_guard = match it.parent().and_then(ast::MatchArm::cast) {
             Some(arm) => arm
@@ -1755,27 +1756,29 @@ fn pattern_context_for(
         mut_token,
         ref_token,
         record_pat: None,
-        impl_: fetch_immediate_impl(sema, original_file, pat.syntax()),
+        impl_or_trait: fetch_immediate_impl_or_trait(sema, original_file, pat.syntax()),
         missing_variants,
     }
 }
 
-fn fetch_immediate_impl(
+fn fetch_immediate_impl_or_trait(
     sema: &Semantics<'_, RootDatabase>,
     original_file: &SyntaxNode,
     node: &SyntaxNode,
-) -> Option<ast::Impl> {
+) -> Option<Either<ast::Impl, ast::Trait>> {
     let mut ancestors = ancestors_in_file_compensated(sema, original_file, node)?
         .filter_map(ast::Item::cast)
         .filter(|it| !matches!(it, ast::Item::MacroCall(_)));
 
     match ancestors.next()? {
         ast::Item::Const(_) | ast::Item::Fn(_) | ast::Item::TypeAlias(_) => (),
-        ast::Item::Impl(it) => return Some(it),
+        ast::Item::Impl(it) => return Some(Either::Left(it)),
+        ast::Item::Trait(it) => return Some(Either::Right(it)),
         _ => return None,
     }
     match ancestors.next()? {
-        ast::Item::Impl(it) => Some(it),
+        ast::Item::Impl(it) => Some(Either::Left(it)),
+        ast::Item::Trait(it) => Some(Either::Right(it)),
         _ => None,
     }
 }

crates/ide-completion/src/tests/fn_param.rs

@@ -85,7 +85,11 @@ pub(crate) trait SourceRoot {
 }
 "#,
         expect![[r#"
+            bn &mut self
+            bn &self
             bn file_id: usize
+            bn mut self
+            bn self
             kw mut
             kw ref
         "#]],
@@ -183,6 +187,44 @@ impl A {
     )
 }
 
+#[test]
+fn in_trait_only_param() {
+    check(
+        r#"
+trait A {
+    fn foo(file_id: usize) {}
+    fn new($0) {}
+}
+"#,
+        expect![[r#"
+            bn &mut self
+            bn &self
+            bn file_id: usize
+            bn mut self
+            bn self
+            kw mut
+            kw ref
+        "#]],
+    )
+}
+
+#[test]
+fn in_trait_after_self() {
+    check(
+        r#"
+trait A {
+    fn foo(file_id: usize) {}
+    fn new(self, $0) {}
+}
+"#,
+        expect![[r#"
+            bn file_id: usize
+            kw mut
+            kw ref
+        "#]],
+    )
+}
+
 // doesn't complete qux due to there being no expression after
 // see source_analyzer::adjust comment
 #[test]