From 8f959f20ee0fecd644054ffed334c378f9ae20f5 Mon Sep 17 00:00:00 2001 From: Kirill Bulatov Date: Tue, 11 Feb 2020 15:21:12 +0200 Subject: [PATCH] Trait location draft --- crates/ra_assists/src/handlers/auto_import.rs | 129 ++++++++++++++++-- crates/ra_hir/src/code_model.rs | 18 +-- 2 files changed, 127 insertions(+), 20 deletions(-) diff --git a/crates/ra_assists/src/handlers/auto_import.rs b/crates/ra_assists/src/handlers/auto_import.rs index a25f0650d6..a9778fab77 100644 --- a/crates/ra_assists/src/handlers/auto_import.rs +++ b/crates/ra_assists/src/handlers/auto_import.rs @@ -5,7 +5,13 @@ use crate::{ assist_ctx::{Assist, AssistCtx}, insert_use_statement, AssistId, }; -use hir::{db::HirDatabase, Adt, ModPath, Module, ModuleDef, PathResolution, SourceAnalyzer}; +use ast::{FnDefOwner, ModuleItem, ModuleItemOwner}; +use hir::{ + db::{DefDatabase, HirDatabase}, + Adt, AssocContainerId, Crate, Function, HasSource, InFile, ModPath, Module, ModuleDef, + PathResolution, SourceAnalyzer, SourceBinder, Trait, +}; +use rustc_hash::FxHashSet; use std::collections::BTreeSet; // Assist: auto_import @@ -135,21 +141,88 @@ impl ImportCandidate { ImportsLocator::new(db) .find_imports(&self.get_search_query()) .into_iter() - .filter_map(|module_def| match self { + .map(|module_def| match self { ImportCandidate::TraitFunction(function_callee, _) => { - if let ModuleDef::Function(function) = module_def { - dbg!(function); - todo!() - } else { - None + let mut applicable_traits = Vec::new(); + if let ModuleDef::Function(located_function) = module_def { + let trait_candidates = Self::get_trait_candidates( + db, + located_function, + module_with_name_to_import.krate(), + ) + .into_iter() + .map(|trait_candidate| trait_candidate.into()) + .collect(); + + function_callee.ty(db).iterate_path_candidates( + db, + module_with_name_to_import.krate(), + &trait_candidates, + None, + |_, assoc| { + if let AssocContainerId::TraitId(trait_id) = assoc.container(db) { + applicable_traits.push( + module_with_name_to_import + .find_use_path(db, ModuleDef::Trait(trait_id.into())), + ); + }; + None::<()> + }, + ); } + applicable_traits } - _ => module_with_name_to_import.find_use_path(db, module_def), + _ => vec![module_with_name_to_import.find_use_path(db, module_def)], }) + .flatten() + .filter_map(std::convert::identity) .filter(|use_path| !use_path.segments.is_empty()) .take(20) .collect::>() } + + fn get_trait_candidates( + db: &RootDatabase, + called_function: Function, + root_crate: Crate, + ) -> FxHashSet { + let mut source_binder = SourceBinder::new(db); + root_crate + .dependencies(db) + .into_iter() + .map(|dependency| db.crate_def_map(dependency.krate.into())) + .chain(std::iter::once(db.crate_def_map(root_crate.into()))) + .map(|crate_def_map| { + crate_def_map + .modules + .iter() + .filter_map(|(_, module_data)| module_data.declaration_source(db)) + .filter_map(|in_file_module| { + Some((in_file_module.file_id, in_file_module.value.item_list()?.items())) + }) + .map(|(file_id, item_list)| { + let mut if_file_trait_defs = Vec::new(); + for module_item in item_list { + if let ModuleItem::TraitDef(trait_def) = module_item { + if let Some(item_list) = trait_def.item_list() { + if item_list + .functions() + .any(|fn_def| fn_def == called_function.source(db).value) + { + if_file_trait_defs.push(InFile::new(file_id, trait_def)) + } + } + } + } + if_file_trait_defs + }) + .flatten() + .filter_map(|in_file_trait_def| source_binder.to_def(in_file_trait_def)) + .collect::>() + }) + .flatten() + .collect() + } } #[cfg(test)] @@ -452,7 +525,45 @@ mod tests { } #[test] - #[ignore] // TODO kb + fn not_applicable_for_imported_trait() { + check_assist_not_applicable( + auto_import, + r" + mod test_mod { + pub trait TestTrait { + fn test_method(&self); + fn test_function(); + } + + pub trait TestTrait2 { + fn test_method(&self); + fn test_function(); + } + pub enum TestEnum { + One, + Two, + } + + impl TestTrait2 for TestEnum { + fn test_method(&self) {} + fn test_function() {} + } + + impl TestTrait for TestEnum { + fn test_method(&self) {} + fn test_function() {} + } + } + + use test_mod::TestTrait2; + fn main() { + test_mod::TestEnum::test_function<|>; + } + ", + ) + } + + #[test] fn trait_method() { check_assist( auto_import, diff --git a/crates/ra_hir/src/code_model.rs b/crates/ra_hir/src/code_model.rs index 4fb679f6dd..73158b8bd4 100644 --- a/crates/ra_hir/src/code_model.rs +++ b/crates/ra_hir/src/code_model.rs @@ -123,7 +123,7 @@ impl_froms!( ); pub use hir_def::{ - attr::Attrs, item_scope::ItemInNs, visibility::Visibility, AssocItemId, AssocItemLoc, + attr::Attrs, item_scope::ItemInNs, visibility::Visibility, AssocContainerId, AssocItemId, }; use rustc_hash::FxHashSet; @@ -696,16 +696,12 @@ impl AssocItem { AssocItem::TypeAlias(t) => t.module(db), } } - pub fn container(self, db: &impl DefDatabase) -> AssocItemContainer { - let container = match self { - AssocItem::Function(it) => it.id.lookup(db).container, - AssocItem::Const(it) => it.id.lookup(db).container, - AssocItem::TypeAlias(it) => it.id.lookup(db).container, - }; - match container { - AssocContainerId::TraitId(id) => AssocItemContainer::Trait(id.into()), - AssocContainerId::ImplId(id) => AssocItemContainer::ImplBlock(id.into()), - AssocContainerId::ContainerId(_) => panic!("invalid AssocItem"), + + pub fn container(self, db: &impl DefDatabase) -> AssocContainerId { + match self { + AssocItem::Function(f) => f.id.lookup(db).container, + AssocItem::Const(c) => c.id.lookup(db).container, + AssocItem::TypeAlias(t) => t.id.lookup(db).container, } } }