Skip to content

Commit acf5f43

Browse files
SomeoneToIgnoreSomeoneToIgnore
committed
Refactor the code
1 parent 07058cb commit acf5f43

File tree

1 file changed

+93
-96
lines changed

1 file changed

+93
-96
lines changed

‎crates/ra_assists/src/handlers/auto_import.rs‎

Lines changed: 93 additions & 96 deletions
Original file line numberDiff line numberDiff line change
@@ -36,15 +36,8 @@ use std::collections::BTreeSet;
3636
// # pub mod std { pub mod collections { pub struct HashMap { } } }
3737
// ```
3838
pub(crate) fn auto_import(ctx: AssistCtx) -> Option<Assist> {
39-
let auto_import_assets = if let Some(path_under_caret) = ctx.find_node_at_offset::<ast::Path>()
40-
{
41-
AutoImportAssets::for_regular_path(path_under_caret, &ctx)?
42-
} else {
43-
AutoImportAssets::for_method_call(ctx.find_node_at_offset()?, &ctx)?
44-
};
45-
46-
let proposed_imports = auto_import_assets
47-
.search_for_imports(ctx.db, auto_import_assets.module_with_name_to_import);
39+
let auto_import_assets = AutoImportAssets::new(&ctx)?;
40+
let proposed_imports = auto_import_assets.search_for_imports(ctx.db);
4841
if proposed_imports.is_empty() {
4942
return None;
5043
}
@@ -54,7 +47,6 @@ pub(crate) fn auto_import(ctx: AssistCtx) -> Option<Assist> {
5447
} else {
5548
auto_import_assets.get_import_group_message()
5649
};
57-
5850
let mut group = ctx.add_assist_group(assist_group_name);
5951
for import in proposed_imports {
6052
group.add_assist(AssistId("auto_import"), format!("Import `{}`", &import), |edit| {
@@ -77,6 +69,14 @@ struct AutoImportAssets {
7769
}
7870

7971
impl AutoImportAssets {
72+
fn new(ctx: &AssistCtx) -> Option<Self> {
73+
if let Some(path_under_caret) = ctx.find_node_at_offset::<ast::Path>() {
74+
Self::for_regular_path(path_under_caret, &ctx)
75+
} else {
76+
Self::for_method_call(ctx.find_node_at_offset()?, &ctx)
77+
}
78+
}
79+
8080
fn for_method_call(method_call: ast::MethodCallExpr, ctx: &AssistCtx) -> Option<Self> {
8181
let syntax_under_caret = method_call.syntax().to_owned();
8282
let source_analyzer = ctx.source_analyzer(&syntax_under_caret, None);
@@ -111,94 +111,88 @@ impl AutoImportAssets {
111111
})
112112
}
113113

114-
fn get_search_query(&self) -> String {
114+
fn get_search_query(&self) -> &str {
115115
match &self.import_candidate {
116-
ImportCandidate::UnqualifiedName(name_ref)
117-
| ImportCandidate::QualifierStart(name_ref) => name_ref.syntax().to_string(),
118-
ImportCandidate::TraitFunction(_, trait_function) => {
119-
trait_function.syntax().to_string()
120-
}
121-
ImportCandidate::TraitMethod(_, trait_method) => trait_method.syntax().to_string(),
116+
ImportCandidate::UnqualifiedName(name) => name,
117+
ImportCandidate::QualifierStart(qualifier_start) => qualifier_start,
118+
ImportCandidate::TraitFunction(_, trait_function_name) => trait_function_name,
119+
ImportCandidate::TraitMethod(_, trait_method_name) => trait_method_name,
122120
}
123121
}
124122

125123
fn get_import_group_message(&self) -> String {
126124
match &self.import_candidate {
127-
ImportCandidate::UnqualifiedName(name_ref)
128-
| ImportCandidate::QualifierStart(name_ref) => format!("Import {}", name_ref.syntax()),
129-
ImportCandidate::TraitFunction(_, trait_function) => {
130-
format!("Import a trait for function {}", trait_function.syntax())
125+
ImportCandidate::UnqualifiedName(name) => format!("Import {}", name),
126+
ImportCandidate::QualifierStart(qualifier_start) => {
127+
format!("Import {}", qualifier_start)
131128
}
132-
ImportCandidate::TraitMethod(_, trait_method) => {
133-
format!("Import a trait for method {}", trait_method.syntax())
129+
ImportCandidate::TraitFunction(_, trait_function_name) => {
130+
format!("Import a trait for function {}", trait_function_name)
131+
}
132+
ImportCandidate::TraitMethod(_, trait_method_name) => {
133+
format!("Import a trait for method {}", trait_method_name)
134134
}
135135
}
136136
}
137137

138-
fn search_for_imports(
139-
&self,
140-
db: &RootDatabase,
141-
module_with_name_to_import: Module,
142-
) -> BTreeSet<ModPath> {
138+
fn search_for_imports(&self, db: &RootDatabase) -> BTreeSet<ModPath> {
143139
let _p = profile("auto_import::search_for_imports");
140+
let current_crate = self.module_with_name_to_import.krate();
144141
ImportsLocator::new(db)
145142
.find_imports(&self.get_search_query())
146143
.into_iter()
147144
.map(|module_def| match &self.import_candidate {
148145
ImportCandidate::TraitFunction(function_callee, _) => {
149146
let mut applicable_traits = Vec::new();
150147
if let ModuleDef::Function(located_function) = module_def {
151-
let trait_candidates = Self::get_trait_candidates(
152-
db,
153-
located_function,
154-
module_with_name_to_import.krate(),
155-
)
156-
.into_iter()
157-
.map(|trait_candidate| trait_candidate.into())
158-
.collect();
159-
160-
function_callee.iterate_path_candidates(
161-
db,
162-
module_with_name_to_import.krate(),
163-
&trait_candidates,
164-
None,
165-
|_, assoc| {
166-
if let AssocContainerId::TraitId(trait_id) = assoc.container(db) {
167-
applicable_traits.push(
168-
module_with_name_to_import
169-
.find_use_path(db, ModuleDef::Trait(trait_id.into())),
170-
);
171-
};
172-
None::<()>
173-
},
174-
);
148+
let trait_candidates: FxHashSet<_> =
149+
Self::get_trait_candidates(db, located_function, current_crate)
150+
.into_iter()
151+
.map(|trait_candidate| trait_candidate.into())
152+
.collect();
153+
if !trait_candidates.is_empty() {
154+
function_callee.iterate_path_candidates(
155+
db,
156+
current_crate,
157+
&trait_candidates,
158+
None,
159+
|_, assoc| {
160+
if let AssocContainerId::TraitId(trait_id) = assoc.container(db)
161+
{
162+
applicable_traits.push(
163+
self.module_with_name_to_import.find_use_path(
164+
db,
165+
ModuleDef::Trait(trait_id.into()),
166+
),
167+
);
168+
};
169+
None::<()>
170+
},
171+
);
172+
};
175173
}
176174
applicable_traits
177175
}
178176
ImportCandidate::TraitMethod(function_callee, _) => {
179177
let mut applicable_traits = Vec::new();
180178
if let ModuleDef::Function(located_function) = module_def {
181-
let trait_candidates: FxHashSet<_> = Self::get_trait_candidates(
182-
db,
183-
located_function,
184-
module_with_name_to_import.krate(),
185-
)
186-
.into_iter()
187-
.map(|trait_candidate| trait_candidate.into())
188-
.collect();
189-
179+
let trait_candidates: FxHashSet<_> =
180+
Self::get_trait_candidates(db, located_function, current_crate)
181+
.into_iter()
182+
.map(|trait_candidate| trait_candidate.into())
183+
.collect();
190184
if !trait_candidates.is_empty() {
191185
function_callee.iterate_method_candidates(
192186
db,
193-
module_with_name_to_import.krate(),
187+
current_crate,
194188
&trait_candidates,
195189
None,
196190
|_, funciton| {
197191
if let AssocContainerId::TraitId(trait_id) =
198192
funciton.container(db)
199193
{
200194
applicable_traits.push(
201-
module_with_name_to_import.find_use_path(
195+
self.module_with_name_to_import.find_use_path(
202196
db,
203197
ModuleDef::Trait(trait_id.into()),
204198
),
@@ -211,7 +205,7 @@ impl AutoImportAssets {
211205
}
212206
applicable_traits
213207
}
214-
_ => vec![module_with_name_to_import.find_use_path(db, module_def)],
208+
_ => vec![self.module_with_name_to_import.find_use_path(db, module_def)],
215209
})
216210
.flatten()
217211
.filter_map(std::convert::identity)
@@ -235,22 +229,19 @@ impl AutoImportAssets {
235229
crate_def_map
236230
.modules
237231
.iter()
238-
.map(|(_, module_data)| {
239-
let mut traits = Vec::new();
240-
for module_def_id in module_data.scope.declarations() {
241-
if let ModuleDef::Trait(trait_candidate) = module_def_id.into() {
242-
if trait_candidate
243-
.items(db)
244-
.into_iter()
245-
.any(|item| item == AssocItem::Function(called_function))
246-
{
247-
traits.push(trait_candidate)
248-
}
249-
}
232+
.map(|(_, module_data)| module_data.scope.declarations())
233+
.flatten()
234+
.filter_map(|module_def_id| match module_def_id.into() {
235+
ModuleDef::Trait(trait_candidate)
236+
if trait_candidate
237+
.items(db)
238+
.into_iter()
239+
.any(|item| item == AssocItem::Function(called_function)) =>
240+
{
241+
Some(trait_candidate)
250242
}
251-
traits
243+
_ => None,
252244
})
253-
.flatten()
254245
.collect::<FxHashSet<_>>()
255246
})
256247
.flatten()
@@ -259,12 +250,20 @@ impl AutoImportAssets {
259250
}
260251

261252
#[derive(Debug)]
262-
// TODO kb rustdocs
263253
enum ImportCandidate {
264-
UnqualifiedName(ast::NameRef),
265-
QualifierStart(ast::NameRef),
266-
TraitFunction(Type, ast::PathSegment),
267-
TraitMethod(Type, ast::NameRef),
254+
/// Simple name like 'HashMap'
255+
UnqualifiedName(String),
256+
/// First part of the qualified name.
257+
/// For 'std::collections::HashMap', that will be 'std'.
258+
QualifierStart(String),
259+
/// A trait function that has no self parameter.
260+
/// For 'test_mod::TestEnum::test_function', `Type` is the `test_mod::TestEnum` expression type
261+
/// and `String`is the `test_function`
262+
TraitFunction(Type, String),
263+
/// A trait method with self parameter.
264+
/// For 'test_enum.test_method()', `Type` is the `test_enum` expression type
265+
/// and `String` is the `test_method`
266+
TraitMethod(Type, String),
268267
}
269268

270269
impl ImportCandidate {
@@ -278,7 +277,7 @@ impl ImportCandidate {
278277
}
279278
Some(Self::TraitMethod(
280279
source_analyzer.type_of(db, &method_call.expr()?)?,
281-
method_call.name_ref()?,
280+
method_call.name_ref()?.syntax().to_string(),
282281
))
283282
}
284283

@@ -299,36 +298,34 @@ impl ImportCandidate {
299298
if let Some(qualifier_start_resolution) =
300299
source_analyzer.resolve_path(db, &qualifier_start_path)
301300
{
302-
let qualifier_resolution = if &qualifier_start_path == path_under_caret {
301+
let qualifier_resolution = if qualifier_start_path == qualifier {
303302
qualifier_start_resolution
304303
} else {
305304
source_analyzer.resolve_path(db, &qualifier)?
306305
};
307306
if let PathResolution::Def(ModuleDef::Adt(function_callee)) = qualifier_resolution {
308-
Some(ImportCandidate::TraitFunction(function_callee.ty(db), segment))
307+
Some(ImportCandidate::TraitFunction(
308+
function_callee.ty(db),
309+
segment.syntax().to_string(),
310+
))
309311
} else {
310312
None
311313
}
312314
} else {
313-
Some(ImportCandidate::QualifierStart(qualifier_start))
315+
Some(ImportCandidate::QualifierStart(qualifier_start.syntax().to_string()))
314316
}
315317
} else {
316-
if source_analyzer.resolve_path(db, path_under_caret).is_none() {
317-
Some(ImportCandidate::UnqualifiedName(
318-
segment.syntax().descendants().find_map(ast::NameRef::cast)?,
319-
))
320-
} else {
321-
None
322-
}
318+
Some(ImportCandidate::UnqualifiedName(
319+
segment.syntax().descendants().find_map(ast::NameRef::cast)?.syntax().to_string(),
320+
))
323321
}
324322
}
325323
}
326324

327325
#[cfg(test)]
328326
mod tests {
329-
use crate::helpers::{check_assist, check_assist_not_applicable, check_assist_target};
330-
331327
use super::*;
328+
use crate::helpers::{check_assist, check_assist_not_applicable, check_assist_target};
332329

333330
#[test]
334331
fn applicable_when_found_an_import() {

0 commit comments

Comments
 (0)