feat(LSP): suggest auto-importing trait via reexport if trait is not …

…visible

tooling/lsp/src/requests/code_action/import_trait.rs

@@ -1,16 +1,19 @@
+use std::collections::HashSet;
+
 use noirc_errors::Location;
 use noirc_frontend::{
     ast::MethodCallExpression,
-    hir::{def_map::ModuleDefId, resolution::visibility::trait_member_is_visible},
-    hir_def::traits::Trait,
-    node_interner::ReferenceId,
+    hir::def_map::ModuleDefId,
+    node_interner::{ReferenceId, TraitId},
 };
 
 use crate::{
-    modules::relative_module_full_path,
+    modules::{relative_module_full_path, relative_module_id_path},
+    requests::TraitReexport,
     use_segment_positions::{
         use_completion_item_additional_text_edits, UseCompletionItemAdditionTextEditsRequest,
     },
+    visibility::module_def_id_is_visible,
 };
 
 use super::CodeActionFinder;
@@ -29,16 +32,7 @@ impl<'a> CodeActionFinder<'a> {
 
             let trait_impl = self.interner.get_trait_implementation(trait_impl_id);
             let trait_id = trait_impl.borrow().trait_id;
-            let trait_ = self.interner.get_trait(trait_id);
-
-            // Check if the trait is currently imported. If so, no need to suggest anything
-            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;
-            }
-
-            self.push_import_trait_code_action(trait_);
+            self.import_trait(trait_id);
             return;
         }
 
@@ -48,33 +42,86 @@ impl<'a> CodeActionFinder<'a> {
             return;
         };
 
-        for (func_id, trait_id) in
-            self.interner.lookup_trait_methods(&typ, &method_call.method_name.0.contents, true)
-        {
-            let visibility = self.interner.function_modifiers(&func_id).visibility;
-            if !trait_member_is_visible(trait_id, visibility, self.module_id, self.def_maps) {
-                continue;
-            }
+        let trait_methods =
+            self.interner.lookup_trait_methods(&typ, &method_call.method_name.0.contents, true);
+        let trait_ids: HashSet<_> = trait_methods.iter().map(|(_, trait_id)| *trait_id).collect();
 
-            let trait_ = self.interner.get_trait(trait_id);
-            self.push_import_trait_code_action(trait_);
+        for trait_id in trait_ids {
+            self.import_trait(trait_id);
         }
     }
 
-    fn push_import_trait_code_action(&mut self, trait_: &Trait) {
-        let trait_id = trait_.id;
-
+    fn import_trait(&mut self, trait_id: TraitId) {
+        // First check if the trait is visible
+        let trait_ = self.interner.get_trait(trait_id);
+        let visibility = trait_.visibility;
         let module_def_id = ModuleDefId::TraitId(trait_id);
-        let current_module_parent_id = self.module_id.parent(self.def_maps);
-        let Some(module_full_path) = relative_module_full_path(
+        let mut trait_reexport = None;
+
+        if !module_def_id_is_visible(
             module_def_id,
             self.module_id,
-            current_module_parent_id,
+            visibility,
+            None,
             self.interner,
-        ) else {
+            self.def_maps,
+        ) {
+            // If it's not, try to find a visible reexport of the trait
+            // that is visible from the current module
+            let Some((visible_module_id, name, _)) =
+                self.interner.get_trait_reexports(trait_id).iter().find(
+                    |(module_id, _, visibility)| {
+                        module_def_id_is_visible(
+                            module_def_id,
+                            self.module_id,
+                            *visibility,
+                            Some(*module_id),
+                            self.interner,
+                            self.def_maps,
+                        )
+                    },
+                )
+            else {
+                return;
+            };
+            trait_reexport = Some(TraitReexport { module_id: visible_module_id, name });
+        }
+
+        let trait_name = if let Some(trait_reexport) = &trait_reexport {
+            trait_reexport.name
+        } else {
+            &trait_.name
+        };
+
+        // Check if the trait is currently imported. If yes, no need to suggest anything
+        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;
+        }
+
+        let module_def_id = ModuleDefId::TraitId(trait_id);
+        let current_module_parent_id = self.module_id.parent(self.def_maps);
+        let module_full_path = if let Some(trait_reexport) = &trait_reexport {
+            relative_module_id_path(
+                *trait_reexport.module_id,
+                &self.module_id,
+                current_module_parent_id,
+                self.interner,
+            )
+        } else {
+            let Some(path) = relative_module_full_path(
+                module_def_id,
+                self.module_id,
+                current_module_parent_id,
+                self.interner,
+            ) else {
+                return;
+            };
+            path
         };
-        let full_path = format!("{}::{}", module_full_path, trait_.name);
+
+        let full_path = format!("{}::{}", module_full_path, trait_name);
 
         let title = format!("Import {}", full_path);
 
@@ -242,6 +289,118 @@ mod moo {
     }
 }
 
+fn main() {
+    let x: Field = 1;
+    x.foobar();
+}"#;
+
+        assert_code_action(title, src, expected).await;
+    }
+
+    #[test]
+    async fn test_import_trait_in_method_call_when_one_option_but_not_in_scope_via_reexport() {
+        let title = "Import moo::Bar";
+
+        let src = r#"mod moo {
+    mod nested {
+        pub trait Foo {
+            fn foobar(self);
+        }
+
+        impl Foo for Field {
+            fn foobar(self) {}
+        }
+    }
+
+    pub use nested::Foo as Bar;
+}
+
+fn main() {
+    let x: Field = 1;
+    x.foo>|<bar();
+}"#;
+
+        let expected = r#"use moo::Bar;
+
+mod moo {
+    mod nested {
+        pub trait Foo {
+            fn foobar(self);
+        }
+
+        impl Foo for Field {
+            fn foobar(self) {}
+        }
+    }
+
+    pub use nested::Foo as Bar;
+}
+
+fn main() {
+    let x: Field = 1;
+    x.foobar();
+}"#;
+
+        assert_code_action(title, src, expected).await;
+    }
+
+    #[test]
+    async fn test_import_trait_in_method_call_when_multiple_via_reexport() {
+        let title = "Import moo::Baz";
+
+        let src = r#"mod moo {
+    mod nested {
+        pub trait Foo {
+            fn foobar(self);
+        }
+
+        impl Foo for Field {
+            fn foobar(self) {}
+        }
+
+        pub trait Bar {
+            fn foobar(self);
+        }
+
+        impl Bar for Field {
+            fn foobar(self) {}
+        }
+    }
+
+    pub use nested::Foo as Baz;
+    pub use nested::Foo as Qux;
+}
+
+fn main() {
+    let x: Field = 1;
+    x.foo>|<bar();
+}"#;
+
+        let expected = r#"use moo::Baz;
+
+mod moo {
+    mod nested {
+        pub trait Foo {
+            fn foobar(self);
+        }
+
+        impl Foo for Field {
+            fn foobar(self) {}
+        }
+
+        pub trait Bar {
+            fn foobar(self);
+        }
+
+        impl Bar for Field {
+            fn foobar(self) {}
+        }
+    }
+
+    pub use nested::Foo as Baz;
+    pub use nested::Foo as Qux;
+}
+
 fn main() {
     let x: Field = 1;
     x.foobar();

tooling/lsp/src/requests/completion.rs

@@ -6,7 +6,7 @@ use std::{
 use async_lsp::ResponseError;
 use completion_items::{
     field_completion_item, simple_completion_item, snippet_completion_item,
-    trait_impl_method_completion_item, TraitReexport,
+    trait_impl_method_completion_item,
 };
 use convert_case::{Case, Casing};
 use fm::{FileId, FileMap, PathString};
@@ -44,7 +44,7 @@ use crate::{
     LspState,
 };
 
-use super::process_request;
+use super::{process_request, TraitReexport};
 
 mod auto_import;
 mod builtins;

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

@@ -3,7 +3,7 @@ use lsp_types::{
     InsertTextFormat, MarkupContent, MarkupKind,
 };
 use noirc_frontend::{
-    ast::{AttributeTarget, Ident},
+    ast::AttributeTarget,
     hir::def_map::{ModuleDefId, ModuleId},
     hir_def::{function::FuncMeta, stmt::HirPattern},
     node_interner::{FuncId, GlobalId, ReferenceId, StructId, TraitId, TypeAliasId},
@@ -22,15 +22,9 @@ use super::{
         crate_or_module_sort_text, default_sort_text, new_sort_text, operator_sort_text,
         self_mismatch_sort_text,
     },
-    FunctionCompletionKind, FunctionKind, NodeFinder, RequestedItems,
+    FunctionCompletionKind, FunctionKind, NodeFinder, RequestedItems, TraitReexport,
 };
 
-/// Represents a trait reexported from a given module with a name.
-pub(super) struct TraitReexport<'a> {
-    pub(super) module_id: &'a ModuleId,
-    pub(super) name: &'a Ident,
-}
-
 impl<'a> NodeFinder<'a> {
     pub(super) fn module_def_id_completion_items(
         &self,
@@ -366,10 +360,10 @@ impl<'a> NodeFinder<'a> {
         completion_item: &mut CompletionItem,
     ) -> Option<()> {
         // If this is a trait method, check if the trait is in scope
-        let (trait_id, reexport_data) = trait_info?;
+        let (trait_id, trait_reexport) = trait_info?;
 
-        let trait_name = if let Some(reexport_data) = reexport_data {
-            reexport_data.name
+        let trait_name = if let Some(trait_reexport) = trait_reexport {
+            trait_reexport.name
         } else {
             let trait_ = self.interner.get_trait(trait_id);
             &trait_.name
@@ -383,7 +377,7 @@ impl<'a> NodeFinder<'a> {
 
         // If not, automatically import it
         let current_module_parent_id = self.module_id.parent(self.def_maps);
-        let module_full_path = if let Some(reexport_data) = reexport_data {
+        let module_full_path = if let Some(reexport_data) = trait_reexport {
             relative_module_id_path(
                 *reexport_data.module_id,
                 &self.module_id,

tooling/lsp/src/requests/mod.rs

@@ -16,8 +16,9 @@ use lsp_types::{
 };
 use nargo_fmt::Config;
 
+use noirc_frontend::ast::Ident;
 use noirc_frontend::graph::CrateId;
-use noirc_frontend::hir::def_map::CrateDefMap;
+use noirc_frontend::hir::def_map::{CrateDefMap, ModuleId};
 use noirc_frontend::parser::ParserError;
 use noirc_frontend::usage_tracker::UsageTracker;
 use noirc_frontend::{graph::Dependency, node_interner::NodeInterner};
@@ -619,6 +620,12 @@ pub(crate) fn find_all_references(
         .unwrap_or_default()
 }
 
+/// Represents a trait reexported from a given module with a name.
+pub(crate) struct TraitReexport<'a> {
+    pub(super) module_id: &'a ModuleId,
+    pub(super) name: &'a Ident,
+}
+
 #[cfg(test)]
 mod initialization {
     use acvm::blackbox_solver::StubbedBlackBoxSolver;