From acf5f43639c489022b5ff3273fc6324dbf2ef017 Mon Sep 17 00:00:00 2001 From: Kirill Bulatov Date: Wed, 12 Feb 2020 14:41:34 +0200 Subject: Refactor the code --- crates/ra_assists/src/handlers/auto_import.rs | 189 +++++++++++++------------- 1 file changed, 93 insertions(+), 96 deletions(-) diff --git a/crates/ra_assists/src/handlers/auto_import.rs b/crates/ra_assists/src/handlers/auto_import.rs index c564f5027..86615d659 100644 --- a/crates/ra_assists/src/handlers/auto_import.rs +++ b/crates/ra_assists/src/handlers/auto_import.rs @@ -36,15 +36,8 @@ use std::collections::BTreeSet; // # pub mod std { pub mod collections { pub struct HashMap { } } } // ``` pub(crate) fn auto_import(ctx: AssistCtx) -> Option { - let auto_import_assets = if let Some(path_under_caret) = ctx.find_node_at_offset::() - { - AutoImportAssets::for_regular_path(path_under_caret, &ctx)? - } else { - AutoImportAssets::for_method_call(ctx.find_node_at_offset()?, &ctx)? - }; - - let proposed_imports = auto_import_assets - .search_for_imports(ctx.db, auto_import_assets.module_with_name_to_import); + let auto_import_assets = AutoImportAssets::new(&ctx)?; + let proposed_imports = auto_import_assets.search_for_imports(ctx.db); if proposed_imports.is_empty() { return None; } @@ -54,7 +47,6 @@ pub(crate) fn auto_import(ctx: AssistCtx) -> Option { } else { auto_import_assets.get_import_group_message() }; - let mut group = ctx.add_assist_group(assist_group_name); for import in proposed_imports { group.add_assist(AssistId("auto_import"), format!("Import `{}`", &import), |edit| { @@ -77,6 +69,14 @@ struct AutoImportAssets { } impl AutoImportAssets { + fn new(ctx: &AssistCtx) -> Option { + if let Some(path_under_caret) = ctx.find_node_at_offset::() { + Self::for_regular_path(path_under_caret, &ctx) + } else { + Self::for_method_call(ctx.find_node_at_offset()?, &ctx) + } + } + fn for_method_call(method_call: ast::MethodCallExpr, ctx: &AssistCtx) -> Option { let syntax_under_caret = method_call.syntax().to_owned(); let source_analyzer = ctx.source_analyzer(&syntax_under_caret, None); @@ -111,36 +111,33 @@ impl AutoImportAssets { }) } - fn get_search_query(&self) -> String { + fn get_search_query(&self) -> &str { match &self.import_candidate { - ImportCandidate::UnqualifiedName(name_ref) - | ImportCandidate::QualifierStart(name_ref) => name_ref.syntax().to_string(), - ImportCandidate::TraitFunction(_, trait_function) => { - trait_function.syntax().to_string() - } - ImportCandidate::TraitMethod(_, trait_method) => trait_method.syntax().to_string(), + ImportCandidate::UnqualifiedName(name) => name, + ImportCandidate::QualifierStart(qualifier_start) => qualifier_start, + ImportCandidate::TraitFunction(_, trait_function_name) => trait_function_name, + ImportCandidate::TraitMethod(_, trait_method_name) => trait_method_name, } } fn get_import_group_message(&self) -> String { match &self.import_candidate { - ImportCandidate::UnqualifiedName(name_ref) - | ImportCandidate::QualifierStart(name_ref) => format!("Import {}", name_ref.syntax()), - ImportCandidate::TraitFunction(_, trait_function) => { - format!("Import a trait for function {}", trait_function.syntax()) + ImportCandidate::UnqualifiedName(name) => format!("Import {}", name), + ImportCandidate::QualifierStart(qualifier_start) => { + format!("Import {}", qualifier_start) } - ImportCandidate::TraitMethod(_, trait_method) => { - format!("Import a trait for method {}", trait_method.syntax()) + ImportCandidate::TraitFunction(_, trait_function_name) => { + format!("Import a trait for function {}", trait_function_name) + } + ImportCandidate::TraitMethod(_, trait_method_name) => { + format!("Import a trait for method {}", trait_method_name) } } } - fn search_for_imports( - &self, - db: &RootDatabase, - module_with_name_to_import: Module, - ) -> BTreeSet { + fn search_for_imports(&self, db: &RootDatabase) -> BTreeSet { let _p = profile("auto_import::search_for_imports"); + let current_crate = self.module_with_name_to_import.krate(); ImportsLocator::new(db) .find_imports(&self.get_search_query()) .into_iter() @@ -148,49 +145,46 @@ impl AutoImportAssets { ImportCandidate::TraitFunction(function_callee, _) => { 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.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::<()> - }, - ); + let trait_candidates: FxHashSet<_> = + Self::get_trait_candidates(db, located_function, current_crate) + .into_iter() + .map(|trait_candidate| trait_candidate.into()) + .collect(); + if !trait_candidates.is_empty() { + function_callee.iterate_path_candidates( + db, + current_crate, + &trait_candidates, + None, + |_, assoc| { + if let AssocContainerId::TraitId(trait_id) = assoc.container(db) + { + applicable_traits.push( + self.module_with_name_to_import.find_use_path( + db, + ModuleDef::Trait(trait_id.into()), + ), + ); + }; + None::<()> + }, + ); + }; } applicable_traits } ImportCandidate::TraitMethod(function_callee, _) => { let mut applicable_traits = Vec::new(); if let ModuleDef::Function(located_function) = module_def { - let trait_candidates: FxHashSet<_> = Self::get_trait_candidates( - db, - located_function, - module_with_name_to_import.krate(), - ) - .into_iter() - .map(|trait_candidate| trait_candidate.into()) - .collect(); - + let trait_candidates: FxHashSet<_> = + Self::get_trait_candidates(db, located_function, current_crate) + .into_iter() + .map(|trait_candidate| trait_candidate.into()) + .collect(); if !trait_candidates.is_empty() { function_callee.iterate_method_candidates( db, - module_with_name_to_import.krate(), + current_crate, &trait_candidates, None, |_, funciton| { @@ -198,7 +192,7 @@ impl AutoImportAssets { funciton.container(db) { applicable_traits.push( - module_with_name_to_import.find_use_path( + self.module_with_name_to_import.find_use_path( db, ModuleDef::Trait(trait_id.into()), ), @@ -211,7 +205,7 @@ impl AutoImportAssets { } applicable_traits } - _ => vec![module_with_name_to_import.find_use_path(db, module_def)], + _ => vec![self.module_with_name_to_import.find_use_path(db, module_def)], }) .flatten() .filter_map(std::convert::identity) @@ -235,22 +229,19 @@ impl AutoImportAssets { crate_def_map .modules .iter() - .map(|(_, module_data)| { - let mut traits = Vec::new(); - for module_def_id in module_data.scope.declarations() { - if let ModuleDef::Trait(trait_candidate) = module_def_id.into() { - if trait_candidate - .items(db) - .into_iter() - .any(|item| item == AssocItem::Function(called_function)) - { - traits.push(trait_candidate) - } - } + .map(|(_, module_data)| module_data.scope.declarations()) + .flatten() + .filter_map(|module_def_id| match module_def_id.into() { + ModuleDef::Trait(trait_candidate) + if trait_candidate + .items(db) + .into_iter() + .any(|item| item == AssocItem::Function(called_function)) => + { + Some(trait_candidate) } - traits + _ => None, }) - .flatten() .collect::>() }) .flatten() @@ -259,12 +250,20 @@ impl AutoImportAssets { } #[derive(Debug)] -// TODO kb rustdocs enum ImportCandidate { - UnqualifiedName(ast::NameRef), - QualifierStart(ast::NameRef), - TraitFunction(Type, ast::PathSegment), - TraitMethod(Type, ast::NameRef), + /// Simple name like 'HashMap' + UnqualifiedName(String), + /// First part of the qualified name. + /// For 'std::collections::HashMap', that will be 'std'. + QualifierStart(String), + /// A trait function that has no self parameter. + /// For 'test_mod::TestEnum::test_function', `Type` is the `test_mod::TestEnum` expression type + /// and `String`is the `test_function` + TraitFunction(Type, String), + /// A trait method with self parameter. + /// For 'test_enum.test_method()', `Type` is the `test_enum` expression type + /// and `String` is the `test_method` + TraitMethod(Type, String), } impl ImportCandidate { @@ -278,7 +277,7 @@ impl ImportCandidate { } Some(Self::TraitMethod( source_analyzer.type_of(db, &method_call.expr()?)?, - method_call.name_ref()?, + method_call.name_ref()?.syntax().to_string(), )) } @@ -299,36 +298,34 @@ impl ImportCandidate { if let Some(qualifier_start_resolution) = source_analyzer.resolve_path(db, &qualifier_start_path) { - let qualifier_resolution = if &qualifier_start_path == path_under_caret { + let qualifier_resolution = if qualifier_start_path == qualifier { qualifier_start_resolution } else { source_analyzer.resolve_path(db, &qualifier)? }; if let PathResolution::Def(ModuleDef::Adt(function_callee)) = qualifier_resolution { - Some(ImportCandidate::TraitFunction(function_callee.ty(db), segment)) + Some(ImportCandidate::TraitFunction( + function_callee.ty(db), + segment.syntax().to_string(), + )) } else { None } } else { - Some(ImportCandidate::QualifierStart(qualifier_start)) + Some(ImportCandidate::QualifierStart(qualifier_start.syntax().to_string())) } } else { - if source_analyzer.resolve_path(db, path_under_caret).is_none() { - Some(ImportCandidate::UnqualifiedName( - segment.syntax().descendants().find_map(ast::NameRef::cast)?, - )) - } else { - None - } + Some(ImportCandidate::UnqualifiedName( + segment.syntax().descendants().find_map(ast::NameRef::cast)?.syntax().to_string(), + )) } } } #[cfg(test)] mod tests { - use crate::helpers::{check_assist, check_assist_not_applicable, check_assist_target}; - use super::*; + use crate::helpers::{check_assist, check_assist_not_applicable, check_assist_target}; #[test] fn applicable_when_found_an_import() { -- cgit v1.2.3