feat: auto-import traits when suggesting trait methods (#7037)

tooling/lsp/src/requests/completion.rs

@@ -706,6 +706,7 @@ impl<'a> NodeFinder<'a> {
                     function_completion_kind,
                     function_kind,
                     None, // attribute first type
+                    trait_id,
                     self_prefix,
                 );
                 if !completion_items.is_empty() {
@@ -758,6 +759,7 @@ impl<'a> NodeFinder<'a> {
                     function_completion_kind,
                     function_kind,
                     None, // attribute first type
+                    None, // trait_id (we are suggesting methods for `Trait::>|<` so no need to auto-import it)
                     self_prefix,
                 );
                 if !completion_items.is_empty() {

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

@@ -10,6 +10,13 @@ use noirc_frontend::{
     QuotedType, Type,
 };
 
+use crate::{
+    modules::relative_module_full_path,
+    use_segment_positions::{
+        use_completion_item_additional_text_edits, UseCompletionItemAdditionTextEditsRequest,
+    },
+};
+
 use super::{
     sort_text::{
         crate_or_module_sort_text, default_sort_text, new_sort_text, operator_sort_text,
@@ -75,6 +82,7 @@ impl<'a> NodeFinder<'a> {
                 function_completion_kind,
                 function_kind,
                 attribute_first_type.as_ref(),
+                None,  // trait_id
                 false, // self_prefix
             ),
             ModuleDefId::TypeId(struct_id) => vec![self.struct_completion_item(name, struct_id)],
@@ -144,13 +152,15 @@ impl<'a> NodeFinder<'a> {
         self.completion_item_with_doc_comments(ReferenceId::Global(global_id), completion_item)
     }
 
+    #[allow(clippy::too_many_arguments)]
     pub(super) fn function_completion_items(
         &self,
         name: &String,
         func_id: FuncId,
         function_completion_kind: FunctionCompletionKind,
         function_kind: FunctionKind,
         attribute_first_type: Option<&Type>,
+        trait_id: Option<TraitId>,
         self_prefix: bool,
     ) -> Vec<CompletionItem> {
         let func_meta = self.interner.function_meta(&func_id);
@@ -223,6 +233,7 @@ impl<'a> NodeFinder<'a> {
                 function_completion_kind,
                 function_kind,
                 attribute_first_type,
+                trait_id,
                 self_prefix,
                 is_macro_call,
             )
@@ -265,6 +276,7 @@ impl<'a> NodeFinder<'a> {
         function_completion_kind: FunctionCompletionKind,
         function_kind: FunctionKind,
         attribute_first_type: Option<&Type>,
+        trait_id: Option<TraitId>,
         self_prefix: bool,
         is_macro_call: bool,
     ) -> CompletionItem {
@@ -325,7 +337,7 @@ impl<'a> NodeFinder<'a> {
             completion_item
         };
 
-        let completion_item = match function_completion_kind {
+        let mut completion_item = match function_completion_kind {
             FunctionCompletionKind::Name => completion_item,
             FunctionCompletionKind::NameAndParameters => {
                 if has_arguments {
@@ -336,9 +348,51 @@ impl<'a> NodeFinder<'a> {
             }
         };
 
+        self.auto_import_trait_if_trait_method(func_id, trait_id, &mut completion_item);
+
         self.completion_item_with_doc_comments(ReferenceId::Function(func_id), completion_item)
     }
 
+    fn auto_import_trait_if_trait_method(
+        &self,
+        func_id: FuncId,
+        trait_id: Option<TraitId>,
+        completion_item: &mut CompletionItem,
+    ) -> Option<()> {
+        // If this is a trait method, check if the trait is in scope
+        let trait_ = self.interner.get_trait(trait_id?);
+        let module_data =
+            &self.def_maps[&self.module_id.krate].modules()[self.module_id.local_id.0];
+        if !module_data.scope().find_name(&trait_.name).is_none() {
+            return None;
+        }
+        // If not, automatically import it
+        let current_module_parent_id = self.module_id.parent(self.def_maps);
+        let module_full_path = relative_module_full_path(
+            ModuleDefId::FunctionId(func_id),
+            self.module_id,
+            current_module_parent_id,
+            self.interner,
+        )?;
+        let full_path = format!("{}::{}", module_full_path, trait_.name);
+        let mut label_details = completion_item.label_details.clone().unwrap();
+        label_details.detail = Some(format!("(use {})", full_path));
+        completion_item.label_details = Some(label_details);
+        completion_item.additional_text_edits = Some(use_completion_item_additional_text_edits(
+            UseCompletionItemAdditionTextEditsRequest {
+                full_path: &full_path,
+                files: self.files,
+                file: self.file,
+                lines: &self.lines,
+                nesting: self.nesting,
+                auto_import_line: self.auto_import_line,
+            },
+            &self.use_segment_positions,
+        ));
+
+        None
+    }
+
     fn compute_function_insert_text(
         &self,
         func_meta: &FuncMeta,

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

@@ -2930,4 +2930,96 @@ fn main() {
         let items = get_completions(src).await;
         assert_eq!(items.len(), 2);
     }
+
+    #[test]
+    async fn test_suggests_and_imports_trait_method_without_self() {
+        let src = r#"
+mod moo {
+    pub trait Foo {
+        fn foobar();
+    }
+
+    impl Foo for Field {
+        fn foobar() {}
+    }
+}
+
+fn main() {
+    Field::fooba>|<
+}
+        "#;
+        let mut items = get_completions(src).await;
+        assert_eq!(items.len(), 1);
+
+        let item = items.remove(0);
+        assert_eq!(item.label_details.unwrap().detail, Some("(use moo::Foo)".to_string()));
+
+        let new_code =
+            apply_text_edits(&src.replace(">|<", ""), &item.additional_text_edits.unwrap());
+
+        let expected = r#"use moo::Foo;
+
+mod moo {
+    pub trait Foo {
+        fn foobar();
+    }
+
+    impl Foo for Field {
+        fn foobar() {}
+    }
+}
+
+fn main() {
+    Field::fooba
+}
+        "#;
+        assert_eq!(new_code, expected);
+    }
+
+    #[test]
+    async fn test_suggests_and_imports_trait_method_with_self() {
+        let src = r#"
+mod moo {
+    pub trait Foo {
+        fn foobar(self);
+    }
+
+    impl Foo for Field {
+        fn foobar(self) {}
+    }
+}
+
+fn main() {
+    let x: Field = 1;
+    x.fooba>|<
+}
+        "#;
+        let mut items = get_completions(src).await;
+        assert_eq!(items.len(), 1);
+
+        let item = items.remove(0);
+        assert_eq!(item.label_details.unwrap().detail, Some("(use moo::Foo)".to_string()));
+
+        let new_code =
+            apply_text_edits(&src.replace(">|<", ""), &item.additional_text_edits.unwrap());
+
+        let expected = r#"use moo::Foo;
+
+mod moo {
+    pub trait Foo {
+        fn foobar(self);
+    }
+
+    impl Foo for Field {
+        fn foobar(self) {}
+    }
+}
+
+fn main() {
+    let x: Field = 1;
+    x.fooba
+}
+        "#;
+        assert_eq!(new_code, expected);
+    }
 }