Render closure fn trait kind in siganture help

Veykril · Veykril · commit 4b3d7f60390b · 2024-05-18T15:01:37.000+02:00
diff --git a/src/tools/rust-analyzer/crates/hir-ty/src/traits.rs b/src/tools/rust-analyzer/crates/hir-ty/src/traits.rs
@@ -1,5 +1,6 @@
 //! Trait solving using Chalk.
 
+use core::fmt;
 use std::env::var;
 
 use chalk_ir::{fold::TypeFoldable, DebruijnIndex, GoalData};
@@ -209,6 +210,16 @@ pub enum FnTrait {
     Fn,
 }
 
+impl fmt::Display for FnTrait {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        match self {
+            FnTrait::FnOnce => write!(f, "FnOnce"),
+            FnTrait::FnMut => write!(f, "FnMut"),
+            FnTrait::Fn => write!(f, "Fn"),
+        }
+    }
+}
+
 impl FnTrait {
     const fn lang_item(self) -> LangItem {
         match self {
diff --git a/src/tools/rust-analyzer/crates/hir/src/lib.rs b/src/tools/rust-analyzer/crates/hir/src/lib.rs
@@ -2199,7 +2199,7 @@ impl Param {
     pub fn as_local(&self, db: &dyn HirDatabase) -> Option<Local> {
         let parent = match self.func {
             Callee::Def(CallableDefId::FunctionId(it)) => DefWithBodyId::FunctionId(it),
-            Callee::Closure(closure) => db.lookup_intern_closure(closure.into()).0,
+            Callee::Closure(closure, _) => db.lookup_intern_closure(closure.into()).0,
             _ => return None,
         };
         let body = db.body(parent);
@@ -2237,7 +2237,7 @@ impl Param {
                 }
                 .map(|value| InFile { file_id, value })
             }
-            Callee::Closure(closure) => {
+            Callee::Closure(closure, _) => {
                 let InternedClosure(owner, expr_id) = db.lookup_intern_closure(closure.into());
                 let (_, source_map) = db.body_with_source_map(owner);
                 let ast @ InFile { file_id, value } = source_map.expr_syntax(expr_id).ok()?;
@@ -4316,16 +4316,23 @@ impl Type {
     }
 
     pub fn as_callable(&self, db: &dyn HirDatabase) -> Option<Callable> {
-        let mut the_ty = &self.ty;
         let callee = match self.ty.kind(Interner) {
-            TyKind::Ref(_, _, ty) if ty.as_closure().is_some() => {
-                the_ty = ty;
-                Callee::Closure(ty.as_closure().unwrap())
-            }
-            TyKind::Closure(id, _) => Callee::Closure(*id),
+            TyKind::Closure(id, subst) => Callee::Closure(*id, subst.clone()),
             TyKind::Function(_) => Callee::FnPtr,
             TyKind::FnDef(..) => Callee::Def(self.ty.callable_def(db)?),
-            _ => {
+            kind => {
+                // This branch shouldn't be necessary?
+                if let TyKind::Ref(_, _, ty) = kind {
+                    if let TyKind::Closure(closure, subst) = ty.kind(Interner) {
+                        let sig = ty.callable_sig(db)?;
+                        return Some(Callable {
+                            ty: self.clone(),
+                            sig,
+                            callee: Callee::Closure(*closure, subst.clone()),
+                            is_bound_method: false,
+                        });
+                    }
+                }
                 let sig = hir_ty::callable_sig_from_fnonce(&self.ty, self.env.clone(), db)?;
                 return Some(Callable {
                     ty: self.clone(),
@@ -4336,7 +4343,7 @@ impl Type {
             }
         };
 
-        let sig = the_ty.callable_sig(db)?;
+        let sig = self.ty.callable_sig(db)?;
         Some(Callable { ty: self.clone(), sig, callee, is_bound_method: false })
     }
 
@@ -4953,13 +4960,13 @@ pub struct Callable {
     sig: CallableSig,
     callee: Callee,
     /// Whether this is a method that was called with method call syntax.
-    pub(crate) is_bound_method: bool,
+    is_bound_method: bool,
 }
 
-#[derive(Copy, Clone, PartialEq, Eq, Hash, Debug)]
+#[derive(Clone, PartialEq, Eq, Hash, Debug)]
 enum Callee {
     Def(CallableDefId),
-    Closure(ClosureId),
+    Closure(ClosureId, Substitution),
     FnPtr,
     Other,
 }
@@ -4968,22 +4975,25 @@ pub enum CallableKind {
     Function(Function),
     TupleStruct(Struct),
     TupleEnumVariant(Variant),
-    Closure,
+    Closure(Closure),
     FnPtr,
     /// Some other type that implements `FnOnce`.
     Other,
 }
 
 impl Callable {
     pub fn kind(&self) -> CallableKind {
-        use Callee::*;
         match self.callee {
-            Def(CallableDefId::FunctionId(it)) => CallableKind::Function(it.into()),
-            Def(CallableDefId::StructId(it)) => CallableKind::TupleStruct(it.into()),
-            Def(CallableDefId::EnumVariantId(it)) => CallableKind::TupleEnumVariant(it.into()),
-            Closure(_) => CallableKind::Closure,
-            FnPtr => CallableKind::FnPtr,
-            Other => CallableKind::Other,
+            Callee::Def(CallableDefId::FunctionId(it)) => CallableKind::Function(it.into()),
+            Callee::Def(CallableDefId::StructId(it)) => CallableKind::TupleStruct(it.into()),
+            Callee::Def(CallableDefId::EnumVariantId(it)) => {
+                CallableKind::TupleEnumVariant(it.into())
+            }
+            Callee::Closure(id, ref subst) => {
+                CallableKind::Closure(Closure { id, subst: subst.clone() })
+            }
+            Callee::FnPtr => CallableKind::FnPtr,
+            Callee::Other => CallableKind::Other,
         }
     }
     pub fn receiver_param(&self, db: &dyn HirDatabase) -> Option<(SelfParam, Type)> {
@@ -5004,7 +5014,7 @@ impl Callable {
             .enumerate()
             .skip(if self.is_bound_method { 1 } else { 0 })
             .map(|(idx, ty)| (idx, self.ty.derived(ty.clone())))
-            .map(|(idx, ty)| Param { func: self.callee, idx, ty })
+            .map(|(idx, ty)| Param { func: self.callee.clone(), idx, ty })
             .collect()
     }
     pub fn return_type(&self) -> Type {
diff --git a/src/tools/rust-analyzer/crates/ide/src/signature_help.rs b/src/tools/rust-analyzer/crates/ide/src/signature_help.rs
@@ -201,7 +201,10 @@ fn signature_help_for_call(
                 variant.name(db).display(db)
             );
         }
-        hir::CallableKind::Closure | hir::CallableKind::FnPtr | hir::CallableKind::Other => (),
+        hir::CallableKind::Closure(closure) => {
+            format_to!(res.signature, "impl {}", closure.fn_trait(db));
+        }
+        hir::CallableKind::FnPtr | hir::CallableKind::Other => (),
     }
 
     res.signature.push('(');
@@ -245,7 +248,7 @@ fn signature_help_for_call(
             render(func.ret_type(db))
         }
         hir::CallableKind::Function(_)
-        | hir::CallableKind::Closure
+        | hir::CallableKind::Closure(_)
         | hir::CallableKind::FnPtr
         | hir::CallableKind::Other => render(callable.return_type()),
         hir::CallableKind::TupleStruct(_) | hir::CallableKind::TupleEnumVariant(_) => {}
@@ -1348,15 +1351,43 @@ fn test() { S.foo($0); }
             r#"
 struct S;
 fn foo(s: S) -> i32 { 92 }
+fn main() {
+    let _move = S;
+    (|s| {{_move}; foo(s)})($0)
+}
+        "#,
+            expect![[r#"
+                impl FnOnce(s: S) -> i32
+                            ^^^^
+            "#]],
+        );
+        check(
+            r#"
+struct S;
+fn foo(s: S) -> i32 { 92 }
 fn main() {
     (|s| foo(s))($0)
 }
         "#,
             expect![[r#"
-                (s: S) -> i32
-                 ^^^^
+                impl Fn(s: S) -> i32
+                        ^^^^
             "#]],
-        )
+        );
+        check(
+            r#"
+struct S;
+fn foo(s: S) -> i32 { 92 }
+fn main() {
+    let mut mutate = 0;
+    (|s| { mutate = 1; foo(s) })($0)
+}
+        "#,
+            expect![[r#"
+                impl FnMut(s: S) -> i32
+                           ^^^^
+            "#]],
+        );
     }
 
     #[test]
