Insert ref for completions

Author: adamrk

crates/completion/src/completion_context.rs

@@ -246,6 +246,19 @@ impl<'a> CompletionContext<'a> {
         }
     }
 
+    pub(crate) fn active_name_and_type(&self) -> Option<(String, Type)> {
+        if let Some(record_field) = &self.record_field_syntax {
+            mark::hit!(record_field_type_match);
+            let (struct_field, _local) = self.sema.resolve_record_field(record_field)?;
+            Some((struct_field.name(self.db).to_string(), struct_field.signature_ty(self.db)))
+        } else if let Some(active_parameter) = &self.active_parameter {
+            mark::hit!(active_param_type_match);
+            Some((active_parameter.name.clone(), active_parameter.ty.clone()))
+        } else {
+            None
+        }
+    }
+
     fn fill_keyword_patterns(&mut self, file_with_fake_ident: &SyntaxNode, offset: TextSize) {
         let fake_ident_token = file_with_fake_ident.token_at_offset(offset).right_biased().unwrap();
         let syntax_element = NodeOrToken::Token(fake_ident_token);

crates/completion/src/completion_item.rs

@@ -2,7 +2,7 @@
 
 use std::fmt;
 
-use hir::Documentation;
+use hir::{Documentation, Mutability};
 use syntax::TextRange;
 use text_edit::TextEdit;
 
@@ -56,6 +56,10 @@ pub struct CompletionItem {
 
     /// Score is useful to pre select or display in better order completion items
     score: Option<CompletionScore>,
+
+    /// Indicates that a reference or mutable reference to this variable is a
+    /// possible match.
+    ref_match: Option<(Mutability, CompletionScore)>,
 }
 
 // We use custom debug for CompletionItem to make snapshot tests more readable.
@@ -194,6 +198,7 @@ impl CompletionItem {
             deprecated: None,
             trigger_call_info: None,
             score: None,
+            ref_match: None,
         }
     }
     /// What user sees in pop-up in the UI.
@@ -240,10 +245,15 @@ impl CompletionItem {
     pub fn trigger_call_info(&self) -> bool {
         self.trigger_call_info
     }
+
+    pub fn ref_match(&self) -> Option<(Mutability, CompletionScore)> {
+        self.ref_match
+    }
 }
 
 /// A helper to make `CompletionItem`s.
 #[must_use]
+#[derive(Clone)]
 pub(crate) struct Builder {
     source_range: TextRange,
     completion_kind: CompletionKind,
@@ -258,6 +268,7 @@ pub(crate) struct Builder {
     deprecated: Option<bool>,
     trigger_call_info: Option<bool>,
     score: Option<CompletionScore>,
+    ref_match: Option<(Mutability, CompletionScore)>,
 }
 
 impl Builder {
@@ -288,6 +299,7 @@ impl Builder {
             deprecated: self.deprecated.unwrap_or(false),
             trigger_call_info: self.trigger_call_info.unwrap_or(false),
             score: self.score,
+            ref_match: self.ref_match,
         }
     }
     pub(crate) fn lookup_by(mut self, lookup: impl Into<String>) -> Builder {
@@ -350,6 +362,13 @@ impl Builder {
         self.trigger_call_info = Some(true);
         self
     }
+    pub(crate) fn set_ref_match(
+        mut self,
+        ref_match: Option<(Mutability, CompletionScore)>,
+    ) -> Builder {
+        self.ref_match = ref_match;
+        self
+    }
 }
 
 impl<'a> Into<CompletionItem> for Builder {

crates/completion/src/presentation.rs

@@ -1,7 +1,7 @@
 //! This modules takes care of rendering various definitions as completion items.
 //! It also handles scoring (sorting) completions.
 
-use hir::{HasAttrs, HasSource, HirDisplay, ModPath, ScopeDef, StructKind, Type};
+use hir::{HasAttrs, HasSource, HirDisplay, ModPath, Mutability, ScopeDef, StructKind, Type};
 use itertools::Itertools;
 use syntax::{ast::NameOwner, display::*};
 use test_utils::mark;
@@ -107,9 +107,16 @@ impl Completions {
             }
         };
 
+        let mut ref_match = None;
         if let ScopeDef::Local(local) = resolution {
-            if let Some(score) = compute_score(ctx, &local.ty(ctx.db), &local_name) {
-                completion_item = completion_item.set_score(score);
+            if let Some((active_name, active_type)) = ctx.active_name_and_type() {
+                let ty = local.ty(ctx.db);
+                if let Some(score) =
+                    compute_score_from_active(&active_type, &active_name, &ty, &local_name)
+                {
+                    completion_item = completion_item.set_score(score);
+                }
+                ref_match = refed_type_matches(&active_type, &active_name, &ty, &local_name);
             }
         }
 
@@ -131,7 +138,7 @@ impl Completions {
             }
         }
 
-        completion_item.kind(kind).set_documentation(docs).add_to(self)
+        completion_item.kind(kind).set_documentation(docs).set_ref_match(ref_match).add_to(self)
     }
 
     pub(crate) fn add_macro(
@@ -342,25 +349,15 @@ impl Completions {
     }
 }
 
-pub(crate) fn compute_score(
-    ctx: &CompletionContext,
+fn compute_score_from_active(
+    active_type: &Type,
+    active_name: &str,
     ty: &Type,
     name: &str,
 ) -> Option<CompletionScore> {
-    let (active_name, active_type) = if let Some(record_field) = &ctx.record_field_syntax {
-        mark::hit!(record_field_type_match);
-        let (struct_field, _local) = ctx.sema.resolve_record_field(record_field)?;
-        (struct_field.name(ctx.db).to_string(), struct_field.signature_ty(ctx.db))
-    } else if let Some(active_parameter) = &ctx.active_parameter {
-        mark::hit!(active_param_type_match);
-        (active_parameter.name.clone(), active_parameter.ty.clone())
-    } else {
-        return None;
-    };
-
     // Compute score
     // For the same type
-    if &active_type != ty {
+    if active_type != ty {
         return None;
     }
 
@@ -373,6 +370,24 @@ pub(crate) fn compute_score(
 
     Some(res)
 }
+fn refed_type_matches(
+    active_type: &Type,
+    active_name: &str,
+    ty: &Type,
+    name: &str,
+) -> Option<(Mutability, CompletionScore)> {
+    let derefed_active = active_type.remove_ref()?;
+    let score = compute_score_from_active(&derefed_active, &active_name, &ty, &name)?;
+    Some((
+        if active_type.is_mutable_reference() { Mutability::Mut } else { Mutability::Shared },
+        score,
+    ))
+}
+
+fn compute_score(ctx: &CompletionContext, ty: &Type, name: &str) -> Option<CompletionScore> {
+    let (active_name, active_type) = ctx.active_name_and_type()?;
+    compute_score_from_active(&active_type, &active_name, ty, name)
+}
 
 enum Params {
     Named(Vec<String>),

crates/rust-analyzer/src/handlers.rs

@@ -570,7 +570,7 @@ pub(crate) fn handle_completion(
     let line_endings = snap.file_line_endings(position.file_id);
     let items: Vec<CompletionItem> = items
         .into_iter()
-        .map(|item| to_proto::completion_item(&line_index, line_endings, item))
+        .flat_map(|item| to_proto::completion_item(&line_index, line_endings, item))
         .collect();
 
     Ok(Some(items.into()))

crates/rust-analyzer/src/to_proto.rs

@@ -160,7 +160,13 @@ pub(crate) fn completion_item(
     line_index: &LineIndex,
     line_endings: LineEndings,
     completion_item: CompletionItem,
-) -> lsp_types::CompletionItem {
+) -> Vec<lsp_types::CompletionItem> {
+    fn set_score(res: &mut lsp_types::CompletionItem, label: &str) {
+        res.preselect = Some(true);
+        // HACK: sort preselect items first
+        res.sort_text = Some(format!(" {}", label));
+    }
+
     let mut additional_text_edits = Vec::new();
     let mut text_edit = None;
     // LSP does not allow arbitrary edits in completion, so we have to do a
@@ -200,9 +206,7 @@ pub(crate) fn completion_item(
     };
 
     if completion_item.score().is_some() {
-        res.preselect = Some(true);
-        // HACK: sort preselect items first
-        res.sort_text = Some(format!(" {}", completion_item.label()));
+        set_score(&mut res, completion_item.label());
     }
 
     if completion_item.deprecated() {
@@ -217,9 +221,22 @@ pub(crate) fn completion_item(
         });
     }
 
-    res.insert_text_format = Some(insert_text_format(completion_item.insert_text_format()));
+    let mut all_results = match completion_item.ref_match() {
+        Some(ref_match) => {
+            let mut refed = res.clone();
+            let (mutability, _score) = ref_match;
+            let label = format!("&{}{}", mutability.as_keyword_for_ref(), refed.label);
+            set_score(&mut refed, &label);
+            refed.label = label;
+            vec![res, refed]
+        }
+        None => vec![res],
+    };
 
-    res
+    for mut r in all_results.iter_mut() {
+        r.insert_text_format = Some(insert_text_format(completion_item.insert_text_format()));
+    }
+    all_results
 }
 
 pub(crate) fn signature_help(
@@ -775,6 +792,48 @@ mod tests {
 
     use super::*;
 
+    #[test]
+    fn test_completion_with_ref() {
+        let fixture = r#"
+        struct Foo;
+        fn foo(arg: &Foo) {}
+        fn main() {
+            let arg = Foo;
+            foo(<|>)
+        }"#;
+
+        let (offset, text) = test_utils::extract_offset(fixture);
+        let line_index = LineIndex::new(&text);
+        let (analysis, file_id) = Analysis::from_single_file(text);
+        let completions: Vec<(String, Option<String>)> = analysis
+            .completions(
+                &ide::CompletionConfig::default(),
+                base_db::FilePosition { file_id, offset },
+            )
+            .unwrap()
+            .unwrap()
+            .into_iter()
+            .filter(|c| c.label().ends_with("arg"))
+            .map(|c| completion_item(&line_index, LineEndings::Unix, c))
+            .flat_map(|comps| comps.into_iter().map(|c| (c.label, c.sort_text)))
+            .collect();
+        expect_test::expect![[r#"
+            [
+                (
+                    "arg",
+                    None,
+                ),
+                (
+                    "&arg",
+                    Some(
+                        " &arg",
+                    ),
+                ),
+            ]
+        "#]]
+        .assert_debug_eq(&completions);
+    }
+
     #[test]
     fn conv_fold_line_folding_only_fixup() {
         let text = r#"mod a;