crates/ra_assists/src/auto_import.rs

   1 use hir::db::HirDatabase;
   2 use ra_syntax::{
   3     ast, AstNode, SyntaxNode, Direction, TextRange,
   4     SyntaxKind::{ PATH, PATH_SEGMENT, COLONCOLON, COMMA }
   5 };
   6 use crate::assist_ctx::{AssistCtx, Assist, AssistBuilder};
   7
   8 fn collect_path_segments(path: &ast::Path) -> Option<Vec<&ast::PathSegment>> {
   9     let mut v = Vec::new();
  10     collect_path_segments_raw(&mut v, path)?;
  11     return Some(v);
  12 }
  13
  14 fn collect_path_segments_raw<'a>(
  15     segments: &mut Vec<&'a ast::PathSegment>,
  16     mut path: &'a ast::Path,
  17 ) -> Option<usize> {
  18     let oldlen = segments.len();
  19     loop {
  20         let mut children = path.syntax().children();
  21         let (first, second, third) = (
  22             children.next().map(|n| (n, n.kind())),
  23             children.next().map(|n| (n, n.kind())),
  24             children.next().map(|n| (n, n.kind())),
  25         );
  26         match (first, second, third) {
  27             (Some((subpath, PATH)), Some((_, COLONCOLON)), Some((segment, PATH_SEGMENT))) => {
  28                 path = ast::Path::cast(subpath)?;
  29                 segments.push(ast::PathSegment::cast(segment)?);
  30             }
  31             (Some((segment, PATH_SEGMENT)), _, _) => {
  32                 segments.push(ast::PathSegment::cast(segment)?);
  33                 break;
  34             }
  35             (_, _, _) => return None,
  36         }
  37     }
  38     // We need to reverse only the new added segments
  39     let only_new_segments = segments.split_at_mut(oldlen).1;
  40     only_new_segments.reverse();
  41     return Some(segments.len() - oldlen);
  42 }
  43
  44 fn fmt_segments(segments: &[&ast::PathSegment]) -> String {
  45     let mut buf = String::new();
  46     fmt_segments_raw(segments, &mut buf);
  47     return buf;
  48 }
  49
  50 fn fmt_segments_raw(segments: &[&ast::PathSegment], buf: &mut String) {
  51     let mut first = true;
  52     for s in segments {
  53         if !first {
  54             buf.push_str("::");
  55         }
  56         match s.kind() {
  57             Some(ast::PathSegmentKind::Name(nameref)) => buf.push_str(nameref.text()),
  58             Some(ast::PathSegmentKind::SelfKw) => buf.push_str("self"),
  59             Some(ast::PathSegmentKind::SuperKw) => buf.push_str("super"),
  60             Some(ast::PathSegmentKind::CrateKw) => buf.push_str("crate"),
  61             None => {}
  62         }
  63         first = false;
  64     }
  65 }
  66
  67 // Returns the numeber of common segments.
  68 fn compare_path_segments(left: &[&ast::PathSegment], right: &[&ast::PathSegment]) -> usize {
  69     return left.iter().zip(right).filter(|(l, r)| compare_path_segment(l, r)).count();
  70 }
  71
  72 fn compare_path_segment(a: &ast::PathSegment, b: &ast::PathSegment) -> bool {
  73     if let (Some(ka), Some(kb)) = (a.kind(), b.kind()) {
  74         match (ka, kb) {
  75             (ast::PathSegmentKind::Name(nameref_a), ast::PathSegmentKind::Name(nameref_b)) => {
  76                 nameref_a.text() == nameref_b.text()
  77             }
  78             (ast::PathSegmentKind::SelfKw, ast::PathSegmentKind::SelfKw) => true,
  79             (ast::PathSegmentKind::SuperKw, ast::PathSegmentKind::SuperKw) => true,
  80             (ast::PathSegmentKind::CrateKw, ast::PathSegmentKind::CrateKw) => true,
  81             (_, _) => false,
  82         }
  83     } else {
  84         false
  85     }
  86 }
  87
  88 fn compare_path_segment_with_name(a: &ast::PathSegment, b: &ast::Name) -> bool {
  89     if let Some(ka) = a.kind() {
  90         return match (ka, b) {
  91             (ast::PathSegmentKind::Name(nameref_a), _) => nameref_a.text() == b.text(),
  92             (_, _) => false,
  93         };
  94     } else {
  95         false
  96     }
  97 }
  98
  99 #[derive(Copy, Clone)]
 100 enum ImportAction<'a> {
 101     Nothing,
 102     // Add a brand new use statement.
 103     AddNewUse {
 104         anchor: Option<&'a SyntaxNode>, // anchor node
 105         add_after_anchor: bool,
 106     },
 107
 108     // To split an existing use statement creating a nested import.
 109     AddNestedImport {
 110         // how may segments matched with the target path
 111         common_segments: usize,
 112         path_to_split: &'a ast::Path,
 113         // the first segment of path_to_split we want to add into the new nested list
 114         first_segment_to_split: Option<&'a ast::PathSegment>,
 115         // Wether to add 'self' in addition to the target path
 116         add_self: bool,
 117     },
 118     // To add the target path to an existing nested import tree list.
 119     AddInTreeList {
 120         common_segments: usize,
 121         // The UseTreeList where to add the target path
 122         tree_list: &'a ast::UseTreeList,
 123         add_self: bool,
 124     },
 125 }
 126
 127 impl<'a> ImportAction<'a> {
 128     fn add_new_use(anchor: Option<&'a SyntaxNode>, add_after_anchor: bool) -> Self {
 129         ImportAction::AddNewUse { anchor, add_after_anchor }
 130     }
 131
 132     fn add_nested_import(
 133         common_segments: usize,
 134         path_to_split: &'a ast::Path,
 135         first_segment_to_split: Option<&'a ast::PathSegment>,
 136         add_self: bool,
 137     ) -> Self {
 138         ImportAction::AddNestedImport {
 139             common_segments,
 140             path_to_split,
 141             first_segment_to_split,
 142             add_self,
 143         }
 144     }
 145
 146     fn add_in_tree_list(
 147         common_segments: usize,
 148         tree_list: &'a ast::UseTreeList,
 149         add_self: bool,
 150     ) -> Self {
 151         ImportAction::AddInTreeList { common_segments, tree_list, add_self }
 152     }
 153
 154     fn better<'b>(left: &'b ImportAction<'a>, right: &'b ImportAction<'a>) -> &'b ImportAction<'a> {
 155         if left.is_better(right) {
 156             left
 157         } else {
 158             right
 159         }
 160     }
 161
 162     fn is_better(&self, other: &ImportAction) -> bool {
 163         match (self, other) {
 164             (ImportAction::Nothing, _) => true,
 165             (ImportAction::AddInTreeList { .. }, ImportAction::Nothing) => false,
 166             (
 167                 ImportAction::AddNestedImport { common_segments: n, .. },
 168                 ImportAction::AddInTreeList { common_segments: m, .. },
 169             ) => n > m,
 170             (
 171                 ImportAction::AddInTreeList { common_segments: n, .. },
 172                 ImportAction::AddNestedImport { common_segments: m, .. },
 173             ) => n > m,
 174             (ImportAction::AddInTreeList { .. }, _) => true,
 175             (ImportAction::AddNestedImport { .. }, ImportAction::Nothing) => false,
 176             (ImportAction::AddNestedImport { .. }, _) => true,
 177             (ImportAction::AddNewUse { .. }, _) => false,
 178         }
 179     }
 180 }
 181
 182 // Find out the best ImportAction to import target path against current_use_tree.
 183 // If current_use_tree has a nested import the function gets called recursively on every UseTree inside a UseTreeList.
 184 fn walk_use_tree_for_best_action<'a>(
 185     current_path_segments: &mut Vec<&'a ast::PathSegment>, // buffer containing path segments
 186     current_parent_use_tree_list: Option<&'a ast::UseTreeList>, // will be Some value if we are in a nested import
 187     current_use_tree: &'a ast::UseTree, // the use tree we are currently examinating
 188     target: &[&'a ast::PathSegment],    // the path we want to import
 189 ) -> ImportAction<'a> {
 190     // We save the number of segments in the buffer so we can restore the correct segments
 191     // before returning. Recursive call will add segments so we need to delete them.
 192     let prev_len = current_path_segments.len();
 193
 194     let tree_list = current_use_tree.use_tree_list();
 195     let alias = current_use_tree.alias();
 196
 197     let path = match current_use_tree.path() {
 198         Some(path) => path,
 199         None => {
 200             // If the use item don't have a path, it means it's broken (syntax error)
 201             return ImportAction::add_new_use(
 202                 current_use_tree
 203                     .syntax()
 204                     .ancestors()
 205                     .find_map(ast::UseItem::cast)
 206                     .map(AstNode::syntax),
 207                 true,
 208             );
 209         }
 210     };
 211
 212     // This can happen only if current_use_tree is a direct child of a UseItem
 213     if let Some(name) = alias.and_then(ast::NameOwner::name) {
 214         if compare_path_segment_with_name(target[0], name) {
 215             return ImportAction::Nothing;
 216         }
 217     }
 218
 219     collect_path_segments_raw(current_path_segments, path);
 220
 221     // We compare only the new segments added in the line just above.
 222     // The first prev_len segments were already compared in 'parent' recursive calls.
 223     let left = target.split_at(prev_len).1;
 224     let right = current_path_segments.split_at(prev_len).1;
 225     let common = compare_path_segments(left, right);
 226     let mut action = match common {
 227         0 => ImportAction::add_new_use(
 228             // e.g: target is std::fmt and we can have
 229             // use foo::bar
 230             // We add a brand new use statement
 231             current_use_tree.syntax().ancestors().find_map(ast::UseItem::cast).map(AstNode::syntax),
 232             true,
 233         ),
 234         common if common == left.len() && left.len() == right.len() => {
 235             // e.g: target is std::fmt and we can have
 236             // 1- use std::fmt;
 237             // 2- use std::fmt:{ ... }
 238             if let Some(list) = tree_list {
 239                 // In case 2 we need to add self to the nested list
 240                 // unless it's already there
 241                 let has_self = list.use_trees().map(ast::UseTree::path).any(|p| {
 242                     p.and_then(ast::Path::segment)
 243                         .and_then(ast::PathSegment::kind)
 244                         .filter(|k| *k == ast::PathSegmentKind::SelfKw)
 245                         .is_some()
 246                 });
 247
 248                 if has_self {
 249                     ImportAction::Nothing
 250                 } else {
 251                     ImportAction::add_in_tree_list(current_path_segments.len(), list, true)
 252                 }
 253             } else {
 254                 // Case 1
 255                 ImportAction::Nothing
 256             }
 257         }
 258         common if common != left.len() && left.len() == right.len() => {
 259             // e.g: target is std::fmt and we have
 260             // use std::io;
 261             // We need to split.
 262             let segments_to_split = current_path_segments.split_at(prev_len + common).1;
 263             ImportAction::add_nested_import(
 264                 prev_len + common,
 265                 path,
 266                 Some(segments_to_split[0]),
 267                 false,
 268             )
 269         }
 270         common if left.len() > right.len() => {
 271             // e.g: target is std::fmt and we can have
 272             // 1- use std;
 273             // 2- use std::{ ... };
 274
 275             // fallback action
 276             let mut better_action = ImportAction::add_new_use(
 277                 current_use_tree
 278                     .syntax()
 279                     .ancestors()
 280                     .find_map(ast::UseItem::cast)
 281                     .map(AstNode::syntax),
 282                 true,
 283             );
 284             if let Some(list) = tree_list {
 285                 // Case 2, check recursively if the path is already imported in the nested list
 286                 for u in list.use_trees() {
 287                     let child_action =
 288                         walk_use_tree_for_best_action(current_path_segments, Some(list), u, target);
 289                     if child_action.is_better(&better_action) {
 290                         better_action = child_action;
 291                         if let ImportAction::Nothing = better_action {
 292                             return better_action;
 293                         }
 294                     }
 295                 }
 296             } else {
 297                 // Case 1, split
 298                 better_action = ImportAction::add_nested_import(prev_len + common, path, None, true)
 299             }
 300             better_action
 301         }
 302         common if left.len() < right.len() => {
 303             // e.g: target is std::fmt and we can have
 304             // use std::fmt::Debug;
 305             let segments_to_split = current_path_segments.split_at(prev_len + common).1;
 306             ImportAction::add_nested_import(
 307                 prev_len + common,
 308                 path,
 309                 Some(segments_to_split[0]),
 310                 true,
 311             )
 312         }
 313         _ => unreachable!(),
 314     };
 315
 316     // If we are inside a UseTreeList adding a use statement become adding to the existing
 317     // tree list.
 318     action = match (current_parent_use_tree_list, action) {
 319         (Some(use_tree_list), ImportAction::AddNewUse { .. }) => {
 320             ImportAction::add_in_tree_list(prev_len, use_tree_list, false)
 321         }
 322         (_, _) => action,
 323     };
 324
 325     // We remove the segments added
 326     current_path_segments.truncate(prev_len);
 327     return action;
 328 }
 329
 330 fn best_action_for_target<'b, 'a: 'b>(
 331     container: &'a SyntaxNode,
 332     path: &'a ast::Path,
 333     target: &'b [&'a ast::PathSegment],
 334 ) -> ImportAction<'a> {
 335     let mut storage = Vec::with_capacity(16); // this should be the only allocation
 336     let best_action = container
 337         .children()
 338         .filter_map(ast::UseItem::cast)
 339         .filter_map(ast::UseItem::use_tree)
 340         .map(|u| walk_use_tree_for_best_action(&mut storage, None, u, target))
 341         .fold(None, |best, a| {
 342             best.and_then(|best| Some(*ImportAction::better(&best, &a))).or(Some(a))
 343         });
 344
 345     match best_action {
 346         Some(action) => return action,
 347         None => {
 348             // We have no action we no use item was found in container so we find
 349             // another item and we use it as anchor.
 350             // If there are not items, we choose the target path itself as anchor.
 351             let anchor = container
 352                 .children()
 353                 .find_map(ast::ModuleItem::cast)
 354                 .map(AstNode::syntax)
 355                 .or(Some(path.syntax()));
 356
 357             return ImportAction::add_new_use(anchor, false);
 358         }
 359     }
 360 }
 361
 362 fn make_assist(action: &ImportAction, target: &[&ast::PathSegment], edit: &mut AssistBuilder) {
 363     match action {
 364         ImportAction::AddNewUse { anchor, add_after_anchor } => {
 365             make_assist_add_new_use(anchor, *add_after_anchor, target, edit)
 366         }
 367         ImportAction::AddInTreeList { common_segments, tree_list, add_self } => {
 368             // We know that the fist n segments already exists in the use statement we want
 369             // to modify, so we want to add only the last target.len() - n segments.
 370             let segments_to_add = target.split_at(*common_segments).1;
 371             make_assist_add_in_tree_list(tree_list, segments_to_add, *add_self, edit)
 372         }
 373         ImportAction::AddNestedImport {
 374             common_segments,
 375             path_to_split,
 376             first_segment_to_split,
 377             add_self,
 378         } => {
 379             let segments_to_add = target.split_at(*common_segments).1;
 380             make_assist_add_nested_import(
 381                 path_to_split,
 382                 first_segment_to_split,
 383                 segments_to_add,
 384                 *add_self,
 385                 edit,
 386             )
 387         }
 388         _ => {}
 389     }
 390 }
 391
 392 fn make_assist_add_new_use(
 393     anchor: &Option<&SyntaxNode>,
 394     after: bool,
 395     target: &[&ast::PathSegment],
 396     edit: &mut AssistBuilder,
 397 ) {
 398     if let Some(anchor) = anchor {
 399         let indent = ra_fmt::leading_indent(anchor);
 400         let mut buf = String::new();
 401         if after {
 402             buf.push_str("\n");
 403             if let Some(spaces) = indent {
 404                 buf.push_str(spaces);
 405             }
 406         }
 407         buf.push_str("use ");
 408         fmt_segments_raw(target, &mut buf);
 409         buf.push_str(";");
 410         if !after {
 411             buf.push_str("\n\n");
 412             if let Some(spaces) = indent {
 413                 buf.push_str(spaces);
 414             }
 415         }
 416         let position = if after { anchor.range().end() } else { anchor.range().start() };
 417         edit.insert(position, buf);
 418     }
 419 }
 420
 421 fn make_assist_add_in_tree_list(
 422     tree_list: &ast::UseTreeList,
 423     target: &[&ast::PathSegment],
 424     add_self: bool,
 425     edit: &mut AssistBuilder,
 426 ) {
 427     let last = tree_list.use_trees().last();
 428     if let Some(last) = last {
 429         let mut buf = String::new();
 430         let comma = last.syntax().siblings(Direction::Next).find(|n| n.kind() == COMMA);
 431         let offset = if let Some(comma) = comma {
 432             comma.range().end()
 433         } else {
 434             buf.push_str(",");
 435             last.syntax().range().end()
 436         };
 437         if add_self {
 438             buf.push_str(" self")
 439         } else {
 440             buf.push_str(" ");
 441         }
 442         fmt_segments_raw(target, &mut buf);
 443         edit.insert(offset, buf);
 444     } else {
 445
 446     }
 447 }
 448
 449 fn make_assist_add_nested_import(
 450     path: &ast::Path,
 451     first_segment_to_split: &Option<&ast::PathSegment>,
 452     target: &[&ast::PathSegment],
 453     add_self: bool,
 454     edit: &mut AssistBuilder,
 455 ) {
 456     let use_tree = path.syntax().ancestors().find_map(ast::UseTree::cast);
 457     if let Some(use_tree) = use_tree {
 458         let (start, add_colon_colon) = if let Some(first_segment_to_split) = first_segment_to_split
 459         {
 460             (first_segment_to_split.syntax().range().start(), false)
 461         } else {
 462             (use_tree.syntax().range().end(), true)
 463         };
 464         let end = use_tree.syntax().range().end();
 465
 466         let mut buf = String::new();
 467         if add_colon_colon {
 468             buf.push_str("::");
 469         }
 470         buf.push_str("{ ");
 471         if add_self {
 472             buf.push_str("self, ");
 473         }
 474         fmt_segments_raw(target, &mut buf);
 475         if !target.is_empty() {
 476             buf.push_str(", ");
 477         }
 478         edit.insert(start, buf);
 479         edit.insert(end, "}");
 480     }
 481 }
 482
 483 pub(crate) fn auto_import(ctx: AssistCtx<impl HirDatabase>) -> Option<Assist> {
 484     let node = ctx.covering_node();
 485     let current_file = node.ancestors().find_map(ast::SourceFile::cast)?;
 486
 487     let path = node.ancestors().find_map(ast::Path::cast)?;
 488     // We don't want to mess with use statements
 489     if path.syntax().ancestors().find_map(ast::UseItem::cast).is_some() {
 490         return None;
 491     }
 492
 493     let segments = collect_path_segments(path)?;
 494     if segments.len() < 2 {
 495         return None;
 496     }
 497
 498     ctx.build(format!("import {} in the current file", fmt_segments(&segments)), |edit| {
 499         let action = best_action_for_target(current_file.syntax(), path, &segments);
 500         make_assist(&action, segments.as_slice(), edit);
 501         if let Some(last_segment) = path.segment() {
 502             // Here we are assuming the assist will provide a  correct use statement
 503             // so we can delete the path qualifier
 504             edit.delete(TextRange::from_to(
 505                 path.syntax().range().start(),
 506                 last_segment.syntax().range().start(),
 507             ));
 508         }
 509     })
 510 }
 511
 512 #[cfg(test)]
 513 mod tests {
 514     use super::*;
 515     use crate::helpers::{ check_assist, check_assist_not_applicable };
 516
 517     #[test]
 518     fn test_auto_import_file_add_use_no_anchor() {
 519         check_assist(
 520             auto_import,
 521             "
 522 std::fmt::Debug<|>
 523     ",
 524             "
 525 use std::fmt::Debug;
 526
 527 Debug<|>
 528     ",
 529         );
 530     }
 531
 532     #[test]
 533     fn test_auto_import_file_add_use() {
 534         check_assist(
 535             auto_import,
 536             "
 537 use stdx;
 538
 539 impl std::fmt::Debug<|> for Foo {
 540 }
 541     ",
 542             "
 543 use stdx;
 544 use std::fmt::Debug;
 545
 546 impl Debug<|> for Foo {
 547 }
 548     ",
 549         );
 550     }
 551
 552     #[test]
 553     fn test_auto_import_file_add_use_other_anchor() {
 554         check_assist(
 555             auto_import,
 556             "
 557 impl std::fmt::Debug<|> for Foo {
 558 }
 559     ",
 560             "
 561 use std::fmt::Debug;
 562
 563 impl Debug<|> for Foo {
 564 }
 565     ",
 566         );
 567     }
 568
 569     #[test]
 570     fn test_auto_import_file_add_use_other_anchor_indent() {
 571         check_assist(
 572             auto_import,
 573             "
 574     impl std::fmt::Debug<|> for Foo {
 575     }
 576     ",
 577             "
 578     use std::fmt::Debug;
 579
 580     impl Debug<|> for Foo {
 581     }
 582     ",
 583         );
 584     }
 585
 586     #[test]
 587     fn test_auto_import_file_split_different() {
 588         check_assist(
 589             auto_import,
 590             "
 591 use std::fmt;
 592
 593 impl std::io<|> for Foo {
 594 }
 595     ",
 596             "
 597 use std::{ io, fmt};
 598
 599 impl io<|> for Foo {
 600 }
 601     ",
 602         );
 603     }
 604
 605     #[test]
 606     fn test_auto_import_file_split_self_for_use() {
 607         check_assist(
 608             auto_import,
 609             "
 610 use std::fmt;
 611
 612 impl std::fmt::Debug<|> for Foo {
 613 }
 614     ",
 615             "
 616 use std::fmt::{ self, Debug, };
 617
 618 impl Debug<|> for Foo {
 619 }
 620     ",
 621         );
 622     }
 623
 624     #[test]
 625     fn test_auto_import_file_split_self_for_target() {
 626         check_assist(
 627             auto_import,
 628             "
 629 use std::fmt::Debug;
 630
 631 impl std::fmt<|> for Foo {
 632 }
 633     ",
 634             "
 635 use std::fmt::{ self, Debug};
 636
 637 impl fmt<|> for Foo {
 638 }
 639     ",
 640         );
 641     }
 642
 643     #[test]
 644     fn test_auto_import_file_add_to_nested_self_nested() {
 645         check_assist(
 646             auto_import,
 647             "
 648 use std::fmt::{Debug, nested::{Display}};
 649
 650 impl std::fmt::nested<|> for Foo {
 651 }
 652 ",
 653             "
 654 use std::fmt::{Debug, nested::{Display, self}};
 655
 656 impl nested<|> for Foo {
 657 }
 658 ",
 659         );
 660     }
 661
 662     #[test]
 663     fn test_auto_import_file_add_to_nested_self_already_included() {
 664         check_assist(
 665             auto_import,
 666             "
 667 use std::fmt::{Debug, nested::{self, Display}};
 668
 669 impl std::fmt::nested<|> for Foo {
 670 }
 671 ",
 672             "
 673 use std::fmt::{Debug, nested::{self, Display}};
 674
 675 impl nested<|> for Foo {
 676 }
 677 ",
 678         );
 679     }
 680
 681     #[test]
 682     fn test_auto_import_file_add_to_nested_nested() {
 683         check_assist(
 684             auto_import,
 685             "
 686 use std::fmt::{Debug, nested::{Display}};
 687
 688 impl std::fmt::nested::Debug<|> for Foo {
 689 }
 690 ",
 691             "
 692 use std::fmt::{Debug, nested::{Display, Debug}};
 693
 694 impl Debug<|> for Foo {
 695 }
 696 ",
 697         );
 698     }
 699
 700     #[test]
 701     fn test_auto_import_file_alias() {
 702         check_assist(
 703             auto_import,
 704             "
 705 use std::fmt as foo;
 706
 707 impl foo::Debug<|> for Foo {
 708 }
 709 ",
 710             "
 711 use std::fmt as foo;
 712
 713 impl Debug<|> for Foo {
 714 }
 715 ",
 716         );
 717     }
 718
 719     #[test]
 720     fn test_auto_import_not_applicable_one_segment() {
 721         check_assist_not_applicable(
 722             auto_import,
 723             "
 724 impl foo<|> for Foo {
 725 }
 726 ",
 727         );
 728     }
 729 }