From 454b77b3ac4e99b1a272fdc5c36f8babb5781cec Mon Sep 17 00:00:00 2001
From: Ary Borenszweig <asterite@gmail.com>
Date: Tue, 26 Nov 2024 15:50:01 -0300
Subject: [PATCH] fix(LSP): use generic self type to narrow down methods to
 complete (#6617)

---
 compiler/noirc_frontend/src/node_interner.rs |  2 +-
 tooling/lsp/src/requests/completion.rs       | 84 ++++++++++----------
 tooling/lsp/src/requests/completion/tests.rs | 33 ++++++++
 3 files changed, 75 insertions(+), 44 deletions(-)

diff --git a/compiler/noirc_frontend/src/node_interner.rs b/compiler/noirc_frontend/src/node_interner.rs
index 736d37fe83f..6d70ea2fd6d 100644
--- a/compiler/noirc_frontend/src/node_interner.rs
+++ b/compiler/noirc_frontend/src/node_interner.rs
@@ -2351,7 +2351,7 @@ impl Methods {
     }
 
     /// 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,
diff --git a/tooling/lsp/src/requests/completion.rs b/tooling/lsp/src/requests/completion.rs
index d71f0414adc..3762ba9cf8d 100644
--- a/tooling/lsp/src/requests/completion.rs
+++ b/tooling/lsp/src/requests/completion.rs
@@ -645,55 +645,53 @@ impl<'a> NodeFinder<'a> {
 
         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));
                 }
             }
         }
diff --git a/tooling/lsp/src/requests/completion/tests.rs b/tooling/lsp/src/requests/completion/tests.rs
index 8cfb2a4b5ee..9306e38a48a 100644
--- a/tooling/lsp/src/requests/completion/tests.rs
+++ b/tooling/lsp/src/requests/completion/tests.rs
@@ -2780,4 +2780,37 @@ fn main() {
         )
         .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()");
+    }
 }