From 21eef0dd9f4afe0602b418b468d84058bcf5128c Mon Sep 17 00:00:00 2001
From: Ary Borenszweig <asterite@gmail.com>
Date: Thu, 2 Jan 2025 14:08:24 -0300
Subject: [PATCH] feat(LSP): suggest trait methods from where clauses (#6915)

---
 .../src/hir/def_collector/dc_mod.rs           |  3 +
 tooling/lsp/src/requests/completion.rs        | 57 ++++++++++++++++++-
 tooling/lsp/src/requests/completion/tests.rs  | 18 ++++++
 3 files changed, 76 insertions(+), 2 deletions(-)

diff --git a/compiler/noirc_frontend/src/hir/def_collector/dc_mod.rs b/compiler/noirc_frontend/src/hir/def_collector/dc_mod.rs
index 8f46d5429bc..eff48ce22a6 100644
--- a/compiler/noirc_frontend/src/hir/def_collector/dc_mod.rs
+++ b/compiler/noirc_frontend/src/hir/def_collector/dc_mod.rs
@@ -500,6 +500,9 @@ impl<'a> ModCollector<'a> {
                             .def_interner
                             .push_function_definition(func_id, modifiers, trait_id.0, location);
 
+                        let referenced = ReferenceId::Function(func_id);
+                        context.def_interner.add_definition_location(referenced, Some(trait_id.0));
+
                         if !trait_item.doc_comments.is_empty() {
                             context.def_interner.set_doc_comments(
                                 ReferenceId::Function(func_id),
diff --git a/tooling/lsp/src/requests/completion.rs b/tooling/lsp/src/requests/completion.rs
index 3762ba9cf8d..fd6ef60af82 100644
--- a/tooling/lsp/src/requests/completion.rs
+++ b/tooling/lsp/src/requests/completion.rs
@@ -31,7 +31,7 @@ use noirc_frontend::{
         },
     },
     hir_def::traits::Trait,
-    node_interner::{NodeInterner, ReferenceId, StructId},
+    node_interner::{FuncId, NodeInterner, ReferenceId, StructId},
     parser::{Item, ItemKind, ParsedSubModule},
     token::{MetaAttribute, Token, Tokens},
     Kind, ParsedModule, StructType, Type, TypeBinding,
@@ -120,6 +120,8 @@ struct NodeFinder<'a> {
     use_segment_positions: UseSegmentPositions,
     self_type: Option<Type>,
     in_comptime: bool,
+    /// The function we are in, if any
+    func_id: Option<FuncId>,
 }
 
 impl<'a> NodeFinder<'a> {
@@ -165,6 +167,7 @@ impl<'a> NodeFinder<'a> {
             use_segment_positions: UseSegmentPositions::default(),
             self_type: None,
             in_comptime: false,
+            func_id: None,
         }
     }
 
@@ -639,6 +642,13 @@ impl<'a> NodeFinder<'a> {
         function_completion_kind: FunctionCompletionKind,
         self_prefix: bool,
     ) {
+        self.complete_trait_constraints_methods(
+            typ,
+            prefix,
+            function_kind,
+            function_completion_kind,
+        );
+
         let Some(methods_by_name) = self.interner.get_type_methods(typ) else {
             return;
         };
@@ -697,6 +707,31 @@ impl<'a> NodeFinder<'a> {
         }
     }
 
+    fn complete_trait_constraints_methods(
+        &mut self,
+        typ: &Type,
+        prefix: &str,
+        function_kind: FunctionKind,
+        function_completion_kind: FunctionCompletionKind,
+    ) {
+        let Some(func_id) = self.func_id else {
+            return;
+        };
+
+        let func_meta = self.interner.function_meta(&func_id);
+        for constraint in &func_meta.trait_constraints {
+            if *typ == constraint.typ {
+                let trait_ = self.interner.get_trait(constraint.trait_bound.trait_id);
+                self.complete_trait_methods(
+                    trait_,
+                    prefix,
+                    function_kind,
+                    function_completion_kind,
+                );
+            }
+        }
+    }
+
     fn complete_trait_methods(
         &mut self,
         trait_: &Trait,
@@ -1125,8 +1160,17 @@ impl<'a> Visitor for NodeFinder<'a> {
         let old_in_comptime = self.in_comptime;
         self.in_comptime = noir_function.def.is_comptime;
 
+        if let Some(ReferenceId::Function(func_id)) = self
+            .interner
+            .reference_at_location(Location::new(noir_function.name_ident().span(), self.file))
+        {
+            self.func_id = Some(func_id);
+        }
+
         noir_function.def.body.accept(Some(span), self);
 
+        self.func_id = None;
+
         self.in_comptime = old_in_comptime;
         self.type_parameters = old_type_parameters;
         self.self_type = None;
@@ -1207,7 +1251,7 @@ impl<'a> Visitor for NodeFinder<'a> {
 
     fn visit_trait_item_function(
         &mut self,
-        _name: &Ident,
+        name: &Ident,
         generics: &UnresolvedGenerics,
         parameters: &[(Ident, UnresolvedType)],
         return_type: &noirc_frontend::ast::FunctionReturnType,
@@ -1232,7 +1276,16 @@ impl<'a> Visitor for NodeFinder<'a> {
             for (name, _) in parameters {
                 self.local_variables.insert(name.to_string(), name.span());
             }
+
+            if let Some(ReferenceId::Function(func_id)) =
+                self.interner.reference_at_location(Location::new(name.span(), self.file))
+            {
+                self.func_id = Some(func_id);
+            }
+
             body.accept(None, self);
+
+            self.func_id = None;
         };
 
         self.type_parameters = old_type_parameters;
diff --git a/tooling/lsp/src/requests/completion/tests.rs b/tooling/lsp/src/requests/completion/tests.rs
index 97c7ad86d5a..65157090767 100644
--- a/tooling/lsp/src/requests/completion/tests.rs
+++ b/tooling/lsp/src/requests/completion/tests.rs
@@ -2861,4 +2861,22 @@ fn main() {
         assert_eq!(items.len(), 1);
         assert!(items[0].label == "bar_baz()");
     }
+
+    #[test]
+    async fn test_suggests_trait_method_from_where_clause_in_function() {
+        let src = r#"
+        trait Foo {
+            fn foo(self) -> i32;
+        }
+
+        fn something<T>(x: T) -> i32
+        where
+            T: Foo,
+        {
+            x.fo>|<
+        }
+        "#;
+        let items = get_completions(src).await;
+        assert_eq!(items.len(), 1);
+    }
 }