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

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/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();