Implement thin-arrow completion in fn return position

crates/ide-completion/src/completions.rs

@@ -33,7 +33,7 @@ use crate::{
     CompletionContext, CompletionItem, CompletionItemKind,
     context::{
         DotAccess, ItemListKind, NameContext, NameKind, NameRefContext, NameRefKind,
-        PathCompletionCtx, PathKind, PatternContext, TypeLocation, Visible,
+        PathCompletionCtx, PathKind, PatternContext, TypeAscriptionTarget, TypeLocation, Visible,
     },
     item::Builder,
     render::{
@@ -111,11 +111,22 @@ impl Completions {
         }
     }
 
-    pub(crate) fn add_type_keywords(&mut self, ctx: &CompletionContext<'_>) {
-        self.add_keyword_snippet(ctx, "fn", "fn($1)");
-        self.add_keyword_snippet(ctx, "dyn", "dyn $0");
-        self.add_keyword_snippet(ctx, "impl", "impl $0");
-        self.add_keyword_snippet(ctx, "for", "for<$1>");
+    pub(crate) fn add_type_keywords(
+        &mut self,
+        ctx: &CompletionContext<'_>,
+        required_thin_arrow: bool,
+    ) {
+        let mut add_keyword = |kw, snippet| {
+            if required_thin_arrow {
+                self.add_keyword_snippet(ctx, kw, snippet);
+            } else {
+                self.add_keyword_snippet(ctx, kw, &snippet[3..]);
+            }
+        };
+        add_keyword("fn", "-> fn($1)");
+        add_keyword("dyn", "-> dyn $0");
+        add_keyword("impl", "-> impl $0");
+        add_keyword("for", "-> for<$1>");
     }
 
     pub(crate) fn add_super_keyword(
@@ -746,6 +757,12 @@ pub(super) fn complete_name_ref(
                             field::complete_field_list_tuple_variant(acc, ctx, path_ctx);
                         }
                         TypeLocation::TypeAscription(ascription) => {
+                            if let TypeAscriptionTarget::RetType { item: Some(item), .. } =
+                                ascription
+                                && path_ctx.required_thin_arrow()
+                            {
+                                keyword::complete_for_and_where(acc, ctx, &item.clone().into());
+                            }
                             r#type::complete_ascribed_type(acc, ctx, path_ctx, ascription);
                         }
                         TypeLocation::GenericArg { .. }

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

@@ -204,8 +204,9 @@ pub(crate) fn complete_type_path(
                 _ => {}
             };
 
+            // FIXME: "crate::" -> "-> crate::" when required_thin_arrow
             acc.add_nameref_keywords_with_colon(ctx);
-            acc.add_type_keywords(ctx);
+            acc.add_type_keywords(ctx, path_ctx.required_thin_arrow());
             ctx.process_all_names(&mut |name, def, doc_aliases| {
                 if scope_def_applicable(def) {
                     acc.add_path_resolution(ctx, path_ctx, name, def, doc_aliases);
@@ -228,14 +229,14 @@ pub(crate) fn complete_ascribed_type(
         TypeAscriptionTarget::Let(pat) | TypeAscriptionTarget::FnParam(pat) => {
             ctx.sema.type_of_pat(pat.as_ref()?)
         }
-        TypeAscriptionTarget::Const(exp) | TypeAscriptionTarget::RetType(exp) => {
+        TypeAscriptionTarget::Const(exp) | TypeAscriptionTarget::RetType { body: exp, .. } => {
             ctx.sema.type_of_expr(exp.as_ref()?)
         }
     }?
     .adjusted();
     if !ty.is_unknown() {
         let ty_string = ty.display_source_code(ctx.db, ctx.module.into(), true).ok()?;
-        acc.add(render_type_inference(ty_string, ctx));
+        acc.add(render_type_inference(ty_string, ctx, path_ctx));
     }
     None
 }

crates/ide-completion/src/context.rs

@@ -101,6 +101,30 @@ impl PathCompletionCtx<'_> {
             }
         )
     }
+
+    pub(crate) fn required_thin_arrow(&self) -> bool {
+        let PathKind::Type {
+            location:
+                TypeLocation::TypeAscription(TypeAscriptionTarget::RetType {
+                    item: Some(ref fn_item),
+                    ..
+                }),
+        } = self.kind
+        else {
+            return false;
+        };
+        let Some(ret_type) = fn_item.ret_type() else { return true };
+        if ret_type.thin_arrow_token().is_some() {
+            return false;
+        }
+        match ret_type.ty() {
+            Some(ast::Type::PathType(path_ty)) => {
+                path_ty.path().is_some_and(|path| path.as_single_name_ref().is_some())
+            }
+            Some(_) => false,
+            None => true,
+        }
+    }
 }
 
 /// The kind of path we are completing right now.
@@ -230,7 +254,7 @@ impl TypeLocation {
 pub(crate) enum TypeAscriptionTarget {
     Let(Option<ast::Pat>),
     FnParam(Option<ast::Pat>),
-    RetType(Option<ast::Expr>),
+    RetType { body: Option<ast::Expr>, item: Option<ast::Fn> },
     Const(Option<ast::Expr>),
 }
 

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

@@ -1200,15 +1200,14 @@ fn classify_name_ref<'db>(
                     let original = ast::Const::cast(name.syntax().parent()?)?;
                     TypeLocation::TypeAscription(TypeAscriptionTarget::Const(original.body()))
                 },
-                ast::RetType(it) => {
-                    it.thin_arrow_token()?;
+                ast::RetType(_) => {
                     let parent = match ast::Fn::cast(parent.parent()?) {
                         Some(it) => it.param_list(),
                         None => ast::ClosureExpr::cast(parent.parent()?)?.param_list(),
                     };
 
                     let parent = find_opt_node_in_file(original_file, parent)?.syntax().parent()?;
-                    TypeLocation::TypeAscription(TypeAscriptionTarget::RetType(match_ast! {
+                    let body = match_ast! {
                         match parent {
                             ast::ClosureExpr(it) => {
                                 it.body()
@@ -1218,7 +1217,9 @@ fn classify_name_ref<'db>(
                             },
                             _ => return None,
                         }
-                    }))
+                    };
+                    let item = ast::Fn::cast(parent);
+                    TypeLocation::TypeAscription(TypeAscriptionTarget::RetType { body, item })
                 },
                 ast::Param(it) => {
                     it.colon_token()?;

crates/ide-completion/src/item.rs

@@ -478,7 +478,7 @@ impl CompletionItem {
 
 /// A helper to make `CompletionItem`s.
 #[must_use]
-#[derive(Clone)]
+#[derive(Debug, Clone)]
 pub(crate) struct Builder {
     source_range: TextRange,
     imports_to_add: SmallVec<[LocatedImport; 1]>,

crates/ide-completion/src/render.rs

@@ -221,13 +221,17 @@ pub(crate) fn render_tuple_field(
 pub(crate) fn render_type_inference(
     ty_string: String,
     ctx: &CompletionContext<'_>,
+    path_ctx: &PathCompletionCtx<'_>,
 ) -> CompletionItem {
     let mut builder = CompletionItem::new(
         CompletionItemKind::InferredType,
         ctx.source_range(),
-        ty_string,
+        &ty_string,
         ctx.edition,
     );
+    if path_ctx.required_thin_arrow() {
+        builder.insert_text(format!("-> {ty_string}"));
+    }
     builder.set_relevance(CompletionRelevance {
         type_match: Some(CompletionRelevanceTypeMatch::Exact),
         exact_name_match: true,
@@ -423,11 +427,13 @@ fn render_resolution_path(
     let db = completion.db;
     let config = completion.config;
     let requires_import = import_to_add.is_some();
+    let arrow = if path_ctx.required_thin_arrow() { "-> " } else { "" };
 
-    let name = local_name.display_no_db(ctx.completion.edition).to_smolstr();
+    let name = local_name.display(db, completion.edition).to_smolstr();
+    let prefix = format_args!("{arrow}{name}");
     let mut item = render_resolution_simple_(ctx, &local_name, import_to_add, resolution);
-    if local_name.needs_escape(completion.edition) {
-        item.insert_text(local_name.display_no_db(completion.edition).to_smolstr());
+    if local_name.needs_escape(completion.edition) || !arrow.is_empty() {
+        item.insert_text(prefix.to_smolstr());
     }
     // Add `<>` for generic types
     let type_path_no_ty_args = matches!(
@@ -448,7 +454,7 @@ fn render_resolution_path(
             item.lookup_by(name.clone())
                 .label(SmolStr::from_iter([&name, "<…>"]))
                 .trigger_call_info()
-                .insert_snippet(cap, format!("{}<$0>", local_name.display(db, completion.edition)));
+                .insert_snippet(cap, format!("{prefix}<$0>"));
         }
     }
 

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

@@ -116,8 +116,23 @@ fn completes_where() {
     check_with_base_items(
         r"fn func() $0",
         expect![[r#"
-        kw where
-    "#]],
+            en Enum                    Enum
+            ma makro!(…) macro_rules! makro
+            md module
+            st Record                Record
+            st Tuple                  Tuple
+            st Unit                    Unit
+            tt Trait
+            un Union                  Union
+            bt u32                      u32
+            kw crate::
+            kw dyn
+            kw fn
+            kw for
+            kw impl
+            kw self::
+            kw where
+        "#]],
     );
     check_with_base_items(
         r"enum Enum $0",
@@ -243,6 +258,19 @@ impl Copy for S where $0
     );
 }
 
+#[test]
+fn fn_item_where_kw() {
+    check_edit(
+        "where",
+        r#"
+fn foo() $0
+"#,
+        r#"
+fn foo() where $0
+"#,
+    );
+}
+
 #[test]
 fn test_is_not_considered_macro() {
     check_with_base_items(