fix(LSP): use generic self type to narrow down methods to complete

compiler/noirc_frontend/src/node_interner.rs

@@ -218,12 +218,12 @@
     interned_statement_kinds: noirc_arena::Arena<StatementKind>,
 
     // Interned `UnresolvedTypeData`s during comptime code.
     interned_unresolved_type_datas: noirc_arena::Arena<UnresolvedTypeData>,
 
     // Interned `Pattern`s during comptime code.
     interned_patterns: noirc_arena::Arena<Pattern>,
 
     /// Determins whether to run in LSP mode. In LSP mode references are tracked.
     pub(crate) lsp_mode: bool,
 
     /// Store the location of the references in the graph.
@@ -666,7 +666,7 @@
             quoted_types: Default::default(),
             interned_expression_kinds: Default::default(),
             interned_statement_kinds: Default::default(),
             interned_unresolved_type_datas: Default::default(),
             interned_patterns: Default::default(),
             lsp_mode: false,
             location_indices: LocationIndices::default(),
@@ -2201,11 +2201,11 @@
         &mut self,
         typ: UnresolvedTypeData,
     ) -> InternedUnresolvedTypeData {
         InternedUnresolvedTypeData(self.interned_unresolved_type_datas.insert(typ))
     }
 
     pub fn get_unresolved_type_data(&self, id: InternedUnresolvedTypeData) -> &UnresolvedTypeData {
         &self.interned_unresolved_type_datas[id.0]
     }
 
     /// Returns the type of an operator (which is always a function), along with its return type.
@@ -2351,7 +2351,7 @@
     }
 
     /// Select the 1 matching method with an object type matching `typ`
-    fn find_matching_method(
+    pub fn find_matching_method(
         &self,
         typ: &Type,
         has_self_param: bool,

tooling/lsp/src/requests/completion.rs

@@ -46,7 +46,7 @@
 use super::process_request;
 
 mod auto_import;
 mod builtins;
 mod completion_items;
 mod kinds;
 mod sort_text;
@@ -242,7 +242,7 @@
         let mut idents: Vec<Ident> = Vec::new();
 
         // Find in which ident we are in, and in which part of it
         // (it could be that we are completting in the middle of an ident)
         for segment in &path.segments {
             let ident = &segment.ident;
 
@@ -645,55 +645,53 @@
 
         let struct_id = get_type_struct_id(typ);
         let is_primitive = typ.is_primitive();
+        let has_self_param = matches!(function_kind, FunctionKind::SelfType(..));
 
         for (name, methods) in methods_by_name {
-            for (func_id, method_type) in methods.iter() {
-                if function_kind == FunctionKind::Any {
-                    if let Some(method_type) = method_type {
-                        if method_type.unify(typ).is_err() {
-                            continue;
-                        }
-                    }
-                }
-
-                if let Some(struct_id) = struct_id {
-                    let modifiers = self.interner.function_modifiers(&func_id);
-                    let visibility = modifiers.visibility;
-                    if !struct_member_is_visible(
-                        struct_id,
-                        visibility,
-                        self.module_id,
-                        self.def_maps,
-                    ) {
-                        continue;
-                    }
-                }
+            let Some(func_id) =
+                methods.find_matching_method(typ, has_self_param, self.interner).or_else(|| {
+                    // Also try to find a method assuming typ is `&mut typ`:
+                    // we want to suggest methods that take `&mut self` even though a variable might not
+                    // be mutable, so a user can know they need to mark it as mutable.
+                    let typ = Type::MutableReference(Box::new(typ.clone()));
+                    methods.find_matching_method(&typ, has_self_param, self.interner)
+                })
+            else {
+                continue;
+            };
 
-                if is_primitive
-                    && !method_call_is_visible(
-                        typ,
-                        func_id,
-                        self.module_id,
-                        self.interner,
-                        self.def_maps,
-                    )
-                {
+            if let Some(struct_id) = struct_id {
+                let modifiers = self.interner.function_modifiers(&func_id);
+                let visibility = modifiers.visibility;
+                if !struct_member_is_visible(struct_id, visibility, self.module_id, self.def_maps) {
                     continue;
                 }
+            }
 
-                if name_matches(name, prefix) {
-                    let completion_items = self.function_completion_items(
-                        name,
-                        func_id,
-                        function_completion_kind,
-                        function_kind,
-                        None, // attribute first type
-                        self_prefix,
-                    );
-                    if !completion_items.is_empty() {
-                        self.completion_items.extend(completion_items);
-                        self.suggested_module_def_ids.insert(ModuleDefId::FunctionId(func_id));
-                    }
+            if is_primitive
+                && !method_call_is_visible(
+                    typ,
+                    func_id,
+                    self.module_id,
+                    self.interner,
+                    self.def_maps,
+                )
+            {
+                continue;
+            }
+
+            if name_matches(name, prefix) {
+                let completion_items = self.function_completion_items(
+                    name,
+                    func_id,
+                    function_completion_kind,
+                    function_kind,
+                    None, // attribute first type
+                    self_prefix,
+                );
+                if !completion_items.is_empty() {
+                    self.completion_items.extend(completion_items);
+                    self.suggested_module_def_ids.insert(ModuleDefId::FunctionId(func_id));
                 }
             }
         }
@@ -1769,8 +1767,8 @@
 ///
 /// For example:
 ///
 /// // "merk" and "ro" match "merkle" and "root" and are in order
 /// name_matches("compute_merkle_root", "merk_ro") == true
 ///
 /// // "ro" matches "root", but "merkle" comes before it, so no match
 /// name_matches("compute_merkle_root", "ro_mer") == false

tooling/lsp/src/requests/completion/tests.rs

@@ -134,7 +134,7 @@
     #[test]
     async fn test_use_first_segment() {
         let src = r#"
             mod foobaz {}
             mod foobar {}
             use foob>|<
         "#;
@@ -2780,4 +2780,37 @@
         )
         .await;
     }
+
+    #[test]
+    async fn test_suggests_methods_based_on_type_generics() {
+        let src = r#"
+        struct Foo<T> {
+            t: T,
+        }
+
+        impl Foo<Field> {
+            fn bar_baz(_self: Self) -> Field {
+                5
+            }
+        }
+
+        impl Foo<u32> {
+            fn bar(_self: Self) -> Field {
+                5
+            }
+
+            fn baz(_self: Self) -> Field {
+                6
+            }
+        }
+
+        fn main() -> pub Field {
+            let foo: Foo<Field> = Foo { t: 5 };
+            foo.b>|<
+        }
+        "#;
+        let items = get_completions(src).await;
+        assert_eq!(items.len(), 1);
+        assert!(items[0].label == "bar_baz()");
+    }
 }