feat(LSP): auto-import trait reexport if trait is not visible

compiler/noirc_frontend/src/hir/def_collector/dc_crate.rs

@@ -3,7 +3,7 @@ use super::errors::{DefCollectorErrorKind, DuplicateType};
 use crate::elaborator::Elaborator;
 use crate::graph::CrateId;
 use crate::hir::comptime::InterpreterError;
-use crate::hir::def_map::{CrateDefMap, LocalModuleId, ModuleId};
+use crate::hir::def_map::{CrateDefMap, LocalModuleId, ModuleDefId, ModuleId};
 use crate::hir::resolution::errors::ResolverError;
 use crate::hir::type_check::TypeCheckError;
 use crate::locations::ReferencesTracker;
@@ -412,13 +412,24 @@ impl DefCollector {
                                 visibility,
                             );
 
-                            if visibility != ItemVisibility::Private {
+                            if context.def_interner.is_in_lsp_mode()
+                                && visibility != ItemVisibility::Private
+                            {
                                 context.def_interner.register_name_for_auto_import(
                                     name.to_string(),
                                     module_def_id,
                                     visibility,
                                     Some(defining_module),
                                 );
+
+                                if let ModuleDefId::TraitId(trait_id) = module_def_id {
+                                    context.def_interner.add_trait_reexport(
+                                        trait_id,
+                                        defining_module,
+                                        name.clone(),
+                                        visibility,
+                                    );
+                                }
                             }
                         }
 

compiler/noirc_frontend/src/node_interner.rs

@@ -267,6 +267,11 @@ pub struct NodeInterner {
 
     /// Captures the documentation comments for each module, struct, trait, function, etc.
     pub(crate) doc_comments: HashMap<ReferenceId, Vec<String>>,
+
+    /// Only for LSP: a map of trait ID to each module that pub or pub(crate) exports it.
+    /// In LSP this is used to offer importing the trait via one of these exports if
+    /// the trait is not visible where it's defined.
+    trait_reexports: HashMap<TraitId, Vec<(ModuleId, Ident, ItemVisibility)>>,
 }
 
 /// A dependency in the dependency graph may be a type or a definition.
@@ -680,6 +685,7 @@ impl Default for NodeInterner {
             comptime_scopes: vec![HashMap::default()],
             trait_impl_associated_types: HashMap::default(),
             doc_comments: HashMap::default(),
+            trait_reexports: HashMap::default(),
         }
     }
 }
@@ -2256,6 +2262,20 @@ impl NodeInterner {
             _ => None,
         }
     }
+
+    pub fn add_trait_reexport(
+        &mut self,
+        trait_id: TraitId,
+        module_id: ModuleId,
+        name: Ident,
+        visibility: ItemVisibility,
+    ) {
+        self.trait_reexports.entry(trait_id).or_default().push((module_id, name, visibility));
+    }
+
+    pub fn get_trait_reexports(&self, trait_id: TraitId) -> &[(ModuleId, Ident, ItemVisibility)] {
+        self.trait_reexports.get(&trait_id).map_or(&[], |exports| exports)
+    }
 }
 
 impl Methods {

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,
+    trait_impl_method_completion_item, TraitReexport,
 };
 use convert_case::{Case, Casing};
 use fm::{FileId, FileMap, PathString};
@@ -28,7 +28,6 @@ use noirc_frontend::{
         def_map::{CrateDefMap, LocalModuleId, ModuleDefId, ModuleId},
         resolution::visibility::{
             item_in_module_is_visible, method_call_is_visible, struct_member_is_visible,
-            trait_member_is_visible,
         },
     },
     hir_def::traits::Trait,
@@ -41,7 +40,8 @@ use sort_text::underscore_sort_text;
 
 use crate::{
     requests::to_lsp_location, trait_impl_method_stub_generator::TraitImplMethodStubGenerator,
-    use_segment_positions::UseSegmentPositions, utils, LspState,
+    use_segment_positions::UseSegmentPositions, utils, visibility::module_def_id_is_visible,
+    LspState,
 };
 
 use super::process_request;
@@ -679,12 +679,40 @@ impl<'a> NodeFinder<'a> {
                     }
                 }
 
+                let mut trait_reexport = None;
+
                 if let Some(trait_id) = trait_id {
                     let modifiers = self.interner.function_modifiers(&func_id);
                     let visibility = modifiers.visibility;
-                    if !trait_member_is_visible(trait_id, visibility, self.module_id, self.def_maps)
-                    {
-                        continue;
+                    let module_def_id = ModuleDefId::TraitId(trait_id);
+                    if !module_def_id_is_visible(
+                        module_def_id,
+                        self.module_id,
+                        visibility,
+                        None, // defining module
+                        self.interner,
+                        self.def_maps,
+                    ) {
+                        // 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 {
+                            continue;
+                        };
+
+                        trait_reexport = Some(TraitReexport { module_id: visible_module_id, name });
                     }
                 }
 
@@ -706,7 +734,7 @@ impl<'a> NodeFinder<'a> {
                     function_completion_kind,
                     function_kind,
                     None, // attribute first type
-                    trait_id,
+                    trait_id.map(|id| (id, trait_reexport.as_ref())),
                     self_prefix,
                 );
                 if !completion_items.is_empty() {

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

@@ -3,15 +3,15 @@ use lsp_types::{
     InsertTextFormat, MarkupContent, MarkupKind,
 };
 use noirc_frontend::{
-    ast::AttributeTarget,
+    ast::{AttributeTarget, Ident},
     hir::def_map::{ModuleDefId, ModuleId},
     hir_def::{function::FuncMeta, stmt::HirPattern},
     node_interner::{FuncId, GlobalId, ReferenceId, StructId, TraitId, TypeAliasId},
     QuotedType, Type,
 };
 
 use crate::{
-    modules::relative_module_full_path,
+    modules::{relative_module_full_path, relative_module_id_path},
     use_segment_positions::{
         use_completion_item_additional_text_edits, UseCompletionItemAdditionTextEditsRequest,
     },
@@ -25,6 +25,12 @@ use super::{
     FunctionCompletionKind, FunctionKind, NodeFinder, RequestedItems,
 };
 
+/// 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,
@@ -160,7 +166,7 @@ impl<'a> NodeFinder<'a> {
         function_completion_kind: FunctionCompletionKind,
         function_kind: FunctionKind,
         attribute_first_type: Option<&Type>,
-        trait_id: Option<TraitId>,
+        trait_info: Option<(TraitId, Option<&TraitReexport>)>,
         self_prefix: bool,
     ) -> Vec<CompletionItem> {
         let func_meta = self.interner.function_meta(&func_id);
@@ -233,7 +239,7 @@ impl<'a> NodeFinder<'a> {
                 function_completion_kind,
                 function_kind,
                 attribute_first_type,
-                trait_id,
+                trait_info,
                 self_prefix,
                 is_macro_call,
             )
@@ -276,7 +282,7 @@ impl<'a> NodeFinder<'a> {
         function_completion_kind: FunctionCompletionKind,
         function_kind: FunctionKind,
         attribute_first_type: Option<&Type>,
-        trait_id: Option<TraitId>,
+        trait_info: Option<(TraitId, Option<&TraitReexport>)>,
         self_prefix: bool,
         is_macro_call: bool,
     ) -> CompletionItem {
@@ -348,33 +354,51 @@ impl<'a> NodeFinder<'a> {
             }
         };
 
-        self.auto_import_trait_if_trait_method(func_id, trait_id, &mut completion_item);
+        self.auto_import_trait_if_trait_method(func_id, trait_info, &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>,
+        trait_info: Option<(TraitId, Option<&TraitReexport>)>,
         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 (trait_id, reexport_data) = trait_info?;
+
+        let trait_name = if let Some(reexport_data) = reexport_data {
+            reexport_data.name
+        } else {
+            let trait_ = self.interner.get_trait(trait_id);
+            &trait_.name
+        };
+
         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() {
+        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 module_full_path = if let Some(reexport_data) = reexport_data {
+            relative_module_id_path(
+                *reexport_data.module_id,
+                &self.module_id,
+                current_module_parent_id,
+                self.interner,
+            )
+        } else {
+            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);

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

@@ -3015,6 +3015,61 @@ mod moo {
     }
 }
 
+fn main() {
+    let x: Field = 1;
+    x.fooba
+}
+        "#;
+        assert_eq!(new_code, expected);
+    }
+
+    #[test]
+    async fn test_suggests_and_imports_trait_method_with_self_using_public_export() {
+        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.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::Bar)".to_string()));
+
+        let new_code =
+            apply_text_edits(&src.replace(">|<", ""), &item.additional_text_edits.unwrap());
+
+        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.fooba