diff --git a/compiler/noirc_frontend/src/hir/def_collector/dc_crate.rs b/compiler/noirc_frontend/src/hir/def_collector/dc_crate.rs index 9081cb1ccc..7f6509b9f1 100644 --- a/compiler/noirc_frontend/src/hir/def_collector/dc_crate.rs +++ b/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, + ); + } } } diff --git a/compiler/noirc_frontend/src/node_interner.rs b/compiler/noirc_frontend/src/node_interner.rs index 0ec033a399..ae2cf224cb 100644 --- a/compiler/noirc_frontend/src/node_interner.rs +++ b/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>, + + /// 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>, } /// 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 { diff --git a/tooling/lsp/src/requests/code_action/import_trait.rs b/tooling/lsp/src/requests/code_action/import_trait.rs index 1e731aa563..dd500c334a 100644 --- a/tooling/lsp/src/requests/code_action/import_trait.rs +++ b/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>|| 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() { diff --git a/tooling/lsp/src/requests/completion/completion_items.rs b/tooling/lsp/src/requests/completion/completion_items.rs index 86bb308ee3..8cb63d8c17 100644 --- a/tooling/lsp/src/requests/completion/completion_items.rs +++ b/tooling/lsp/src/requests/completion/completion_items.rs @@ -11,7 +11,7 @@ use noirc_frontend::{ }; 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, }, @@ -22,7 +22,7 @@ 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, }; impl<'a> NodeFinder<'a> { @@ -160,7 +160,7 @@ impl<'a> NodeFinder<'a> { function_completion_kind: FunctionCompletionKind, function_kind: FunctionKind, attribute_first_type: Option<&Type>, - trait_id: Option, + trait_info: Option<(TraitId, Option<&TraitReexport>)>, self_prefix: bool, ) -> Vec { let func_meta = self.interner.function_meta(&func_id); @@ -233,7 +233,7 @@ impl<'a> NodeFinder<'a> { function_completion_kind, function_kind, attribute_first_type, - trait_id, + trait_info, self_prefix, is_macro_call, ) @@ -276,7 +276,7 @@ impl<'a> NodeFinder<'a> { function_completion_kind: FunctionCompletionKind, function_kind: FunctionKind, attribute_first_type: Option<&Type>, - trait_id: Option, + trait_info: Option<(TraitId, Option<&TraitReexport>)>, self_prefix: bool, is_macro_call: bool, ) -> CompletionItem { @@ -348,7 +348,7 @@ 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) } @@ -356,25 +356,43 @@ impl<'a> NodeFinder<'a> { fn auto_import_trait_if_trait_method( &self, func_id: FuncId, - trait_id: Option, + 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, trait_reexport) = trait_info?; + + let trait_name = if let Some(trait_reexport) = trait_reexport { + trait_reexport.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) = trait_reexport { + 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); diff --git a/tooling/lsp/src/requests/completion/tests.rs b/tooling/lsp/src/requests/completion/tests.rs index e6cfebddb0..8ff568e3c2 100644 --- a/tooling/lsp/src/requests/completion/tests.rs +++ b/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 diff --git a/tooling/lsp/src/requests/mod.rs b/tooling/lsp/src/requests/mod.rs index d81108c2ec..80f4a167a0 100644 --- a/tooling/lsp/src/requests/mod.rs +++ b/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;