crates/ra_assists/src/handlers/auto_import.rs

   1 use std::collections::BTreeSet;
   2
   3 use either::Either;
   4 use hir::{
   5     AsAssocItem, AssocItemContainer, ModPath, Module, ModuleDef, PathResolution, Semantics, Trait,
   6     Type,
   7 };
   8 use ra_ide_db::{imports_locator::ImportsLocator, RootDatabase};
   9 use ra_prof::profile;
  10 use ra_syntax::{
  11     ast::{self, AstNode},
  12     SyntaxNode,
  13 };
  14 use rustc_hash::FxHashSet;
  15
  16 use crate::{utils::insert_use_statement, AssistContext, AssistId, Assists, GroupLabel};
  17
  18 // Assist: auto_import
  19 //
  20 // If the name is unresolved, provides all possible imports for it.
  21 //
  22 // ```
  23 // fn main() {
  24 //     let map = HashMap<|>::new();
  25 // }
  26 // # pub mod std { pub mod collections { pub struct HashMap { } } }
  27 // ```
  28 // ->
  29 // ```
  30 // use std::collections::HashMap;
  31 //
  32 // fn main() {
  33 //     let map = HashMap::new();
  34 // }
  35 // # pub mod std { pub mod collections { pub struct HashMap { } } }
  36 // ```
  37 pub(crate) fn auto_import(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
  38     let auto_import_assets = AutoImportAssets::new(&ctx)?;
  39     let proposed_imports = auto_import_assets.search_for_imports(ctx.db);
  40     if proposed_imports.is_empty() {
  41         return None;
  42     }
  43
  44     let range = ctx.sema.original_range(&auto_import_assets.syntax_under_caret).range;
  45     let group = auto_import_assets.get_import_group_message();
  46     for import in proposed_imports {
  47         acc.add_group(
  48             &group,
  49             AssistId("auto_import"),
  50             format!("Import `{}`", &import),
  51             range,
  52             |builder| {
  53                 insert_use_statement(&auto_import_assets.syntax_under_caret, &import, ctx, builder);
  54             },
  55         );
  56     }
  57     Some(())
  58 }
  59
  60 #[derive(Debug)]
  61 struct AutoImportAssets {
  62     import_candidate: ImportCandidate,
  63     module_with_name_to_import: Module,
  64     syntax_under_caret: SyntaxNode,
  65 }
  66
  67 impl AutoImportAssets {
  68     fn new(ctx: &AssistContext) -> Option<Self> {
  69         if let Some(path_under_caret) = ctx.find_node_at_offset_with_descend::<ast::Path>() {
  70             Self::for_regular_path(path_under_caret, &ctx)
  71         } else {
  72             Self::for_method_call(ctx.find_node_at_offset_with_descend()?, &ctx)
  73         }
  74     }
  75
  76     fn for_method_call(method_call: ast::MethodCallExpr, ctx: &AssistContext) -> Option<Self> {
  77         let syntax_under_caret = method_call.syntax().to_owned();
  78         let module_with_name_to_import = ctx.sema.scope(&syntax_under_caret).module()?;
  79         Some(Self {
  80             import_candidate: ImportCandidate::for_method_call(&ctx.sema, &method_call)?,
  81             module_with_name_to_import,
  82             syntax_under_caret,
  83         })
  84     }
  85
  86     fn for_regular_path(path_under_caret: ast::Path, ctx: &AssistContext) -> Option<Self> {
  87         let syntax_under_caret = path_under_caret.syntax().to_owned();
  88         if syntax_under_caret.ancestors().find_map(ast::UseItem::cast).is_some() {
  89             return None;
  90         }
  91
  92         let module_with_name_to_import = ctx.sema.scope(&syntax_under_caret).module()?;
  93         Some(Self {
  94             import_candidate: ImportCandidate::for_regular_path(&ctx.sema, &path_under_caret)?,
  95             module_with_name_to_import,
  96             syntax_under_caret,
  97         })
  98     }
  99
 100     fn get_search_query(&self) -> &str {
 101         match &self.import_candidate {
 102             ImportCandidate::UnqualifiedName(name) => name,
 103             ImportCandidate::QualifierStart(qualifier_start) => qualifier_start,
 104             ImportCandidate::TraitAssocItem(_, trait_assoc_item_name) => trait_assoc_item_name,
 105             ImportCandidate::TraitMethod(_, trait_method_name) => trait_method_name,
 106         }
 107     }
 108
 109     fn get_import_group_message(&self) -> GroupLabel {
 110         let name = match &self.import_candidate {
 111             ImportCandidate::UnqualifiedName(name) => format!("Import {}", name),
 112             ImportCandidate::QualifierStart(qualifier_start) => {
 113                 format!("Import {}", qualifier_start)
 114             }
 115             ImportCandidate::TraitAssocItem(_, trait_assoc_item_name) => {
 116                 format!("Import a trait for item {}", trait_assoc_item_name)
 117             }
 118             ImportCandidate::TraitMethod(_, trait_method_name) => {
 119                 format!("Import a trait for method {}", trait_method_name)
 120             }
 121         };
 122         GroupLabel(name)
 123     }
 124
 125     fn search_for_imports(&self, db: &RootDatabase) -> BTreeSet<ModPath> {
 126         let _p = profile("auto_import::search_for_imports");
 127         let current_crate = self.module_with_name_to_import.krate();
 128         ImportsLocator::new(db)
 129             .find_imports(&self.get_search_query())
 130             .into_iter()
 131             .filter_map(|candidate| match &self.import_candidate {
 132                 ImportCandidate::TraitAssocItem(assoc_item_type, _) => {
 133                     let located_assoc_item = match candidate {
 134                         Either::Left(ModuleDef::Function(located_function)) => located_function
 135                             .as_assoc_item(db)
 136                             .map(|assoc| assoc.container(db))
 137                             .and_then(Self::assoc_to_trait),
 138                         Either::Left(ModuleDef::Const(located_const)) => located_const
 139                             .as_assoc_item(db)
 140                             .map(|assoc| assoc.container(db))
 141                             .and_then(Self::assoc_to_trait),
 142                         _ => None,
 143                     }?;
 144
 145                     let mut trait_candidates = FxHashSet::default();
 146                     trait_candidates.insert(located_assoc_item.into());
 147
 148                     assoc_item_type
 149                         .iterate_path_candidates(
 150                             db,
 151                             current_crate,
 152                             &trait_candidates,
 153                             None,
 154                             |_, assoc| Self::assoc_to_trait(assoc.container(db)),
 155                         )
 156                         .map(ModuleDef::from)
 157                         .map(Either::Left)
 158                 }
 159                 ImportCandidate::TraitMethod(function_callee, _) => {
 160                     let located_assoc_item =
 161                         if let Either::Left(ModuleDef::Function(located_function)) = candidate {
 162                             located_function
 163                                 .as_assoc_item(db)
 164                                 .map(|assoc| assoc.container(db))
 165                                 .and_then(Self::assoc_to_trait)
 166                         } else {
 167                             None
 168                         }?;
 169
 170                     let mut trait_candidates = FxHashSet::default();
 171                     trait_candidates.insert(located_assoc_item.into());
 172
 173                     function_callee
 174                         .iterate_method_candidates(
 175                             db,
 176                             current_crate,
 177                             &trait_candidates,
 178                             None,
 179                             |_, function| {
 180                                 Self::assoc_to_trait(function.as_assoc_item(db)?.container(db))
 181                             },
 182                         )
 183                         .map(ModuleDef::from)
 184                         .map(Either::Left)
 185                 }
 186                 _ => Some(candidate),
 187             })
 188             .filter_map(|candidate| match candidate {
 189                 Either::Left(module_def) => {
 190                     self.module_with_name_to_import.find_use_path(db, module_def)
 191                 }
 192                 Either::Right(macro_def) => {
 193                     self.module_with_name_to_import.find_use_path(db, macro_def)
 194                 }
 195             })
 196             .filter(|use_path| !use_path.segments.is_empty())
 197             .take(20)
 198             .collect::<BTreeSet<_>>()
 199     }
 200
 201     fn assoc_to_trait(assoc: AssocItemContainer) -> Option<Trait> {
 202         if let AssocItemContainer::Trait(extracted_trait) = assoc {
 203             Some(extracted_trait)
 204         } else {
 205             None
 206         }
 207     }
 208 }
 209
 210 #[derive(Debug)]
 211 enum ImportCandidate {
 212     /// Simple name like 'HashMap'
 213     UnqualifiedName(String),
 214     /// First part of the qualified name.
 215     /// For 'std::collections::HashMap', that will be 'std'.
 216     QualifierStart(String),
 217     /// A trait associated function (with no self parameter) or associated constant.
 218     /// For 'test_mod::TestEnum::test_function', `Type` is the `test_mod::TestEnum` expression type
 219     /// and `String` is the `test_function`
 220     TraitAssocItem(Type, String),
 221     /// A trait method with self parameter.
 222     /// For 'test_enum.test_method()', `Type` is the `test_enum` expression type
 223     /// and `String` is the `test_method`
 224     TraitMethod(Type, String),
 225 }
 226
 227 impl ImportCandidate {
 228     fn for_method_call(
 229         sema: &Semantics<RootDatabase>,
 230         method_call: &ast::MethodCallExpr,
 231     ) -> Option<Self> {
 232         if sema.resolve_method_call(method_call).is_some() {
 233             return None;
 234         }
 235         Some(Self::TraitMethod(
 236             sema.type_of_expr(&method_call.expr()?)?,
 237             method_call.name_ref()?.syntax().to_string(),
 238         ))
 239     }
 240
 241     fn for_regular_path(
 242         sema: &Semantics<RootDatabase>,
 243         path_under_caret: &ast::Path,
 244     ) -> Option<Self> {
 245         if sema.resolve_path(path_under_caret).is_some() {
 246             return None;
 247         }
 248
 249         let segment = path_under_caret.segment()?;
 250         if let Some(qualifier) = path_under_caret.qualifier() {
 251             let qualifier_start = qualifier.syntax().descendants().find_map(ast::NameRef::cast)?;
 252             let qualifier_start_path =
 253                 qualifier_start.syntax().ancestors().find_map(ast::Path::cast)?;
 254             if let Some(qualifier_start_resolution) = sema.resolve_path(&qualifier_start_path) {
 255                 let qualifier_resolution = if qualifier_start_path == qualifier {
 256                     qualifier_start_resolution
 257                 } else {
 258                     sema.resolve_path(&qualifier)?
 259                 };
 260                 if let PathResolution::Def(ModuleDef::Adt(assoc_item_path)) = qualifier_resolution {
 261                     Some(ImportCandidate::TraitAssocItem(
 262                         assoc_item_path.ty(sema.db),
 263                         segment.syntax().to_string(),
 264                     ))
 265                 } else {
 266                     None
 267                 }
 268             } else {
 269                 Some(ImportCandidate::QualifierStart(qualifier_start.syntax().to_string()))
 270             }
 271         } else {
 272             Some(ImportCandidate::UnqualifiedName(
 273                 segment.syntax().descendants().find_map(ast::NameRef::cast)?.syntax().to_string(),
 274             ))
 275         }
 276     }
 277 }
 278
 279 #[cfg(test)]
 280 mod tests {
 281     use super::*;
 282     use crate::tests::{check_assist, check_assist_not_applicable, check_assist_target};
 283
 284     #[test]
 285     fn applicable_when_found_an_import() {
 286         check_assist(
 287             auto_import,
 288             r"
 289             <|>PubStruct
 290
 291             pub mod PubMod {
 292                 pub struct PubStruct;
 293             }
 294             ",
 295             r"
 296             <|>use PubMod::PubStruct;
 297
 298             PubStruct
 299
 300             pub mod PubMod {
 301                 pub struct PubStruct;
 302             }
 303             ",
 304         );
 305     }
 306
 307     #[test]
 308     fn applicable_when_found_an_import_in_macros() {
 309         check_assist(
 310             auto_import,
 311             r"
 312             macro_rules! foo {
 313                 ($i:ident) => { fn foo(a: $i) {} }
 314             }
 315             foo!(Pub<|>Struct);
 316
 317             pub mod PubMod {
 318                 pub struct PubStruct;
 319             }
 320             ",
 321             r"
 322             use PubMod::PubStruct;
 323
 324             macro_rules! foo {
 325                 ($i:ident) => { fn foo(a: $i) {} }
 326             }
 327             foo!(Pub<|>Struct);
 328
 329             pub mod PubMod {
 330                 pub struct PubStruct;
 331             }
 332             ",
 333         );
 334     }
 335
 336     #[test]
 337     fn auto_imports_are_merged() {
 338         check_assist(
 339             auto_import,
 340             r"
 341             use PubMod::PubStruct1;
 342
 343             struct Test {
 344                 test: Pub<|>Struct2<u8>,
 345             }
 346
 347             pub mod PubMod {
 348                 pub struct PubStruct1;
 349                 pub struct PubStruct2<T> {
 350                     _t: T,
 351                 }
 352             }
 353             ",
 354             r"
 355             use PubMod::{PubStruct2, PubStruct1};
 356
 357             struct Test {
 358                 test: Pub<|>Struct2<u8>,
 359             }
 360
 361             pub mod PubMod {
 362                 pub struct PubStruct1;
 363                 pub struct PubStruct2<T> {
 364                     _t: T,
 365                 }
 366             }
 367             ",
 368         );
 369     }
 370
 371     #[test]
 372     fn applicable_when_found_multiple_imports() {
 373         check_assist(
 374             auto_import,
 375             r"
 376             PubSt<|>ruct
 377
 378             pub mod PubMod1 {
 379                 pub struct PubStruct;
 380             }
 381             pub mod PubMod2 {
 382                 pub struct PubStruct;
 383             }
 384             pub mod PubMod3 {
 385                 pub struct PubStruct;
 386             }
 387             ",
 388             r"
 389             use PubMod3::PubStruct;
 390
 391             PubSt<|>ruct
 392
 393             pub mod PubMod1 {
 394                 pub struct PubStruct;
 395             }
 396             pub mod PubMod2 {
 397                 pub struct PubStruct;
 398             }
 399             pub mod PubMod3 {
 400                 pub struct PubStruct;
 401             }
 402             ",
 403         );
 404     }
 405
 406     #[test]
 407     fn not_applicable_for_already_imported_types() {
 408         check_assist_not_applicable(
 409             auto_import,
 410             r"
 411             use PubMod::PubStruct;
 412
 413             PubStruct<|>
 414
 415             pub mod PubMod {
 416                 pub struct PubStruct;
 417             }
 418             ",
 419         );
 420     }
 421
 422     #[test]
 423     fn not_applicable_for_types_with_private_paths() {
 424         check_assist_not_applicable(
 425             auto_import,
 426             r"
 427             PrivateStruct<|>
 428
 429             pub mod PubMod {
 430                 struct PrivateStruct;
 431             }
 432             ",
 433         );
 434     }
 435
 436     #[test]
 437     fn not_applicable_when_no_imports_found() {
 438         check_assist_not_applicable(
 439             auto_import,
 440             "
 441             PubStruct<|>",
 442         );
 443     }
 444
 445     #[test]
 446     fn not_applicable_in_import_statements() {
 447         check_assist_not_applicable(
 448             auto_import,
 449             r"
 450             use PubStruct<|>;
 451
 452             pub mod PubMod {
 453                 pub struct PubStruct;
 454             }",
 455         );
 456     }
 457
 458     #[test]
 459     fn function_import() {
 460         check_assist(
 461             auto_import,
 462             r"
 463             test_function<|>
 464
 465             pub mod PubMod {
 466                 pub fn test_function() {};
 467             }
 468             ",
 469             r"
 470             use PubMod::test_function;
 471
 472             test_function<|>
 473
 474             pub mod PubMod {
 475                 pub fn test_function() {};
 476             }
 477             ",
 478         );
 479     }
 480
 481     #[test]
 482     fn macro_import() {
 483         check_assist(
 484             auto_import,
 485             r"
 486                     //- /lib.rs crate:crate_with_macro
 487                     #[macro_export]
 488                     macro_rules! foo {
 489                         () => ()
 490                     }
 491
 492                     //- /main.rs crate:main deps:crate_with_macro
 493                     fn main() {
 494                         foo<|>
 495                     }",
 496             r"use crate_with_macro::foo;
 497
 498 fn main() {
 499     foo<|>
 500 }
 501 ",
 502         );
 503     }
 504
 505     #[test]
 506     fn auto_import_target() {
 507         check_assist_target(
 508             auto_import,
 509             r"
 510             struct AssistInfo {
 511                 group_label: Option<<|>GroupLabel>,
 512             }
 513
 514             mod m { pub struct GroupLabel; }
 515             ",
 516             "GroupLabel",
 517         )
 518     }
 519
 520     #[test]
 521     fn not_applicable_when_path_start_is_imported() {
 522         check_assist_not_applicable(
 523             auto_import,
 524             r"
 525             pub mod mod1 {
 526                 pub mod mod2 {
 527                     pub mod mod3 {
 528                         pub struct TestStruct;
 529                     }
 530                 }
 531             }
 532
 533             use mod1::mod2;
 534             fn main() {
 535                 mod2::mod3::TestStruct<|>
 536             }
 537             ",
 538         );
 539     }
 540
 541     #[test]
 542     fn not_applicable_for_imported_function() {
 543         check_assist_not_applicable(
 544             auto_import,
 545             r"
 546             pub mod test_mod {
 547                 pub fn test_function() {}
 548             }
 549
 550             use test_mod::test_function;
 551             fn main() {
 552                 test_function<|>
 553             }
 554             ",
 555         );
 556     }
 557
 558     #[test]
 559     fn associated_struct_function() {
 560         check_assist(
 561             auto_import,
 562             r"
 563             mod test_mod {
 564                 pub struct TestStruct {}
 565                 impl TestStruct {
 566                     pub fn test_function() {}
 567                 }
 568             }
 569
 570             fn main() {
 571                 TestStruct::test_function<|>
 572             }
 573             ",
 574             r"
 575             use test_mod::TestStruct;
 576
 577             mod test_mod {
 578                 pub struct TestStruct {}
 579                 impl TestStruct {
 580                     pub fn test_function() {}
 581                 }
 582             }
 583
 584             fn main() {
 585                 TestStruct::test_function<|>
 586             }
 587             ",
 588         );
 589     }
 590
 591     #[test]
 592     fn associated_struct_const() {
 593         check_assist(
 594             auto_import,
 595             r"
 596             mod test_mod {
 597                 pub struct TestStruct {}
 598                 impl TestStruct {
 599                     const TEST_CONST: u8 = 42;
 600                 }
 601             }
 602
 603             fn main() {
 604                 TestStruct::TEST_CONST<|>
 605             }
 606             ",
 607             r"
 608             use test_mod::TestStruct;
 609
 610             mod test_mod {
 611                 pub struct TestStruct {}
 612                 impl TestStruct {
 613                     const TEST_CONST: u8 = 42;
 614                 }
 615             }
 616
 617             fn main() {
 618                 TestStruct::TEST_CONST<|>
 619             }
 620             ",
 621         );
 622     }
 623
 624     #[test]
 625     fn associated_trait_function() {
 626         check_assist(
 627             auto_import,
 628             r"
 629             mod test_mod {
 630                 pub trait TestTrait {
 631                     fn test_function();
 632                 }
 633                 pub struct TestStruct {}
 634                 impl TestTrait for TestStruct {
 635                     fn test_function() {}
 636                 }
 637             }
 638
 639             fn main() {
 640                 test_mod::TestStruct::test_function<|>
 641             }
 642             ",
 643             r"
 644             use test_mod::TestTrait;
 645
 646             mod test_mod {
 647                 pub trait TestTrait {
 648                     fn test_function();
 649                 }
 650                 pub struct TestStruct {}
 651                 impl TestTrait for TestStruct {
 652                     fn test_function() {}
 653                 }
 654             }
 655
 656             fn main() {
 657                 test_mod::TestStruct::test_function<|>
 658             }
 659             ",
 660         );
 661     }
 662
 663     #[test]
 664     fn not_applicable_for_imported_trait_for_function() {
 665         check_assist_not_applicable(
 666             auto_import,
 667             r"
 668             mod test_mod {
 669                 pub trait TestTrait {
 670                     fn test_function();
 671                 }
 672                 pub trait TestTrait2 {
 673                     fn test_function();
 674                 }
 675                 pub enum TestEnum {
 676                     One,
 677                     Two,
 678                 }
 679                 impl TestTrait2 for TestEnum {
 680                     fn test_function() {}
 681                 }
 682                 impl TestTrait for TestEnum {
 683                     fn test_function() {}
 684                 }
 685             }
 686
 687             use test_mod::TestTrait2;
 688             fn main() {
 689                 test_mod::TestEnum::test_function<|>;
 690             }
 691             ",
 692         )
 693     }
 694
 695     #[test]
 696     fn associated_trait_const() {
 697         check_assist(
 698             auto_import,
 699             r"
 700             mod test_mod {
 701                 pub trait TestTrait {
 702                     const TEST_CONST: u8;
 703                 }
 704                 pub struct TestStruct {}
 705                 impl TestTrait for TestStruct {
 706                     const TEST_CONST: u8 = 42;
 707                 }
 708             }
 709
 710             fn main() {
 711                 test_mod::TestStruct::TEST_CONST<|>
 712             }
 713             ",
 714             r"
 715             use test_mod::TestTrait;
 716
 717             mod test_mod {
 718                 pub trait TestTrait {
 719                     const TEST_CONST: u8;
 720                 }
 721                 pub struct TestStruct {}
 722                 impl TestTrait for TestStruct {
 723                     const TEST_CONST: u8 = 42;
 724                 }
 725             }
 726
 727             fn main() {
 728                 test_mod::TestStruct::TEST_CONST<|>
 729             }
 730             ",
 731         );
 732     }
 733
 734     #[test]
 735     fn not_applicable_for_imported_trait_for_const() {
 736         check_assist_not_applicable(
 737             auto_import,
 738             r"
 739             mod test_mod {
 740                 pub trait TestTrait {
 741                     const TEST_CONST: u8;
 742                 }
 743                 pub trait TestTrait2 {
 744                     const TEST_CONST: f64;
 745                 }
 746                 pub enum TestEnum {
 747                     One,
 748                     Two,
 749                 }
 750                 impl TestTrait2 for TestEnum {
 751                     const TEST_CONST: f64 = 42.0;
 752                 }
 753                 impl TestTrait for TestEnum {
 754                     const TEST_CONST: u8 = 42;
 755                 }
 756             }
 757
 758             use test_mod::TestTrait2;
 759             fn main() {
 760                 test_mod::TestEnum::TEST_CONST<|>;
 761             }
 762             ",
 763         )
 764     }
 765
 766     #[test]
 767     fn trait_method() {
 768         check_assist(
 769             auto_import,
 770             r"
 771             mod test_mod {
 772                 pub trait TestTrait {
 773                     fn test_method(&self);
 774                 }
 775                 pub struct TestStruct {}
 776                 impl TestTrait for TestStruct {
 777                     fn test_method(&self) {}
 778                 }
 779             }
 780
 781             fn main() {
 782                 let test_struct = test_mod::TestStruct {};
 783                 test_struct.test_meth<|>od()
 784             }
 785             ",
 786             r"
 787             use test_mod::TestTrait;
 788
 789             mod test_mod {
 790                 pub trait TestTrait {
 791                     fn test_method(&self);
 792                 }
 793                 pub struct TestStruct {}
 794                 impl TestTrait for TestStruct {
 795                     fn test_method(&self) {}
 796                 }
 797             }
 798
 799             fn main() {
 800                 let test_struct = test_mod::TestStruct {};
 801                 test_struct.test_meth<|>od()
 802             }
 803             ",
 804         );
 805     }
 806
 807     #[test]
 808     fn not_applicable_for_imported_trait_for_method() {
 809         check_assist_not_applicable(
 810             auto_import,
 811             r"
 812             mod test_mod {
 813                 pub trait TestTrait {
 814                     fn test_method(&self);
 815                 }
 816                 pub trait TestTrait2 {
 817                     fn test_method(&self);
 818                 }
 819                 pub enum TestEnum {
 820                     One,
 821                     Two,
 822                 }
 823                 impl TestTrait2 for TestEnum {
 824                     fn test_method(&self) {}
 825                 }
 826                 impl TestTrait for TestEnum {
 827                     fn test_method(&self) {}
 828                 }
 829             }
 830
 831             use test_mod::TestTrait2;
 832             fn main() {
 833                 let one = test_mod::TestEnum::One;
 834                 one.test<|>_method();
 835             }
 836             ",
 837         )
 838     }
 839 }