use hir::{self, db::HirDatabase}; use ra_text_edit::TextEditBuilder; use crate::{ assist_ctx::{Assist, AssistCtx}, AssistId, }; use ra_syntax::{ ast::{self, NameOwner}, AstNode, Direction, SmolStr, SyntaxKind::{PATH, PATH_SEGMENT}, SyntaxNode, TextRange, T, }; fn collect_path_segments_raw<'a>( segments: &mut Vec<&'a ast::PathSegment>, mut path: &'a ast::Path, ) -> Option { let oldlen = segments.len(); loop { let mut children = path.syntax().children_with_tokens(); let (first, second, third) = ( children.next().map(|n| (n, n.kind())), children.next().map(|n| (n, n.kind())), children.next().map(|n| (n, n.kind())), ); match (first, second, third) { (Some((subpath, PATH)), Some((_, T![::])), Some((segment, PATH_SEGMENT))) => { path = ast::Path::cast(subpath.as_node()?)?; segments.push(ast::PathSegment::cast(segment.as_node()?)?); } (Some((segment, PATH_SEGMENT)), _, _) => { segments.push(ast::PathSegment::cast(segment.as_node()?)?); break; } (_, _, _) => return None, } } // We need to reverse only the new added segments let only_new_segments = segments.split_at_mut(oldlen).1; only_new_segments.reverse(); return Some(segments.len() - oldlen); } fn fmt_segments(segments: &[SmolStr]) -> String { let mut buf = String::new(); fmt_segments_raw(segments, &mut buf); return buf; } fn fmt_segments_raw(segments: &[SmolStr], buf: &mut String) { let mut iter = segments.iter(); if let Some(s) = iter.next() { buf.push_str(s); } for s in iter { buf.push_str("::"); buf.push_str(s); } } // Returns the numeber of common segments. fn compare_path_segments(left: &[SmolStr], right: &[&ast::PathSegment]) -> usize { return left.iter().zip(right).filter(|(l, r)| compare_path_segment(l, r)).count(); } fn compare_path_segment(a: &SmolStr, b: &ast::PathSegment) -> bool { if let Some(kb) = b.kind() { match kb { ast::PathSegmentKind::Name(nameref_b) => a == nameref_b.text(), ast::PathSegmentKind::SelfKw => a == "self", ast::PathSegmentKind::SuperKw => a == "super", ast::PathSegmentKind::CrateKw => a == "crate", } } else { false } } fn compare_path_segment_with_name(a: &SmolStr, b: &ast::Name) -> bool { a == b.text() } #[derive(Copy, Clone)] enum ImportAction<'a> { Nothing, // Add a brand new use statement. AddNewUse { anchor: Option<&'a SyntaxNode>, // anchor node add_after_anchor: bool, }, // To split an existing use statement creating a nested import. AddNestedImport { // how may segments matched with the target path common_segments: usize, path_to_split: &'a ast::Path, // the first segment of path_to_split we want to add into the new nested list first_segment_to_split: Option<&'a ast::PathSegment>, // Wether to add 'self' in addition to the target path add_self: bool, }, // To add the target path to an existing nested import tree list. AddInTreeList { common_segments: usize, // The UseTreeList where to add the target path tree_list: &'a ast::UseTreeList, add_self: bool, }, } impl<'a> ImportAction<'a> { fn add_new_use(anchor: Option<&'a SyntaxNode>, add_after_anchor: bool) -> Self { ImportAction::AddNewUse { anchor, add_after_anchor } } fn add_nested_import( common_segments: usize, path_to_split: &'a ast::Path, first_segment_to_split: Option<&'a ast::PathSegment>, add_self: bool, ) -> Self { ImportAction::AddNestedImport { common_segments, path_to_split, first_segment_to_split, add_self, } } fn add_in_tree_list( common_segments: usize, tree_list: &'a ast::UseTreeList, add_self: bool, ) -> Self { ImportAction::AddInTreeList { common_segments, tree_list, add_self } } fn better<'b>(left: &'b ImportAction<'a>, right: &'b ImportAction<'a>) -> &'b ImportAction<'a> { if left.is_better(right) { left } else { right } } fn is_better(&self, other: &ImportAction) -> bool { match (self, other) { (ImportAction::Nothing, _) => true, (ImportAction::AddInTreeList { .. }, ImportAction::Nothing) => false, ( ImportAction::AddNestedImport { common_segments: n, .. }, ImportAction::AddInTreeList { common_segments: m, .. }, ) => n > m, ( ImportAction::AddInTreeList { common_segments: n, .. }, ImportAction::AddNestedImport { common_segments: m, .. }, ) => n > m, (ImportAction::AddInTreeList { .. }, _) => true, (ImportAction::AddNestedImport { .. }, ImportAction::Nothing) => false, (ImportAction::AddNestedImport { .. }, _) => true, (ImportAction::AddNewUse { .. }, _) => false, } } } // Find out the best ImportAction to import target path against current_use_tree. // If current_use_tree has a nested import the function gets called recursively on every UseTree inside a UseTreeList. fn walk_use_tree_for_best_action<'a>( current_path_segments: &mut Vec<&'a ast::PathSegment>, // buffer containing path segments current_parent_use_tree_list: Option<&'a ast::UseTreeList>, // will be Some value if we are in a nested import current_use_tree: &'a ast::UseTree, // the use tree we are currently examinating target: &[SmolStr], // the path we want to import ) -> ImportAction<'a> { // We save the number of segments in the buffer so we can restore the correct segments // before returning. Recursive call will add segments so we need to delete them. let prev_len = current_path_segments.len(); let tree_list = current_use_tree.use_tree_list(); let alias = current_use_tree.alias(); let path = match current_use_tree.path() { Some(path) => path, None => { // If the use item don't have a path, it means it's broken (syntax error) return ImportAction::add_new_use( current_use_tree .syntax() .ancestors() .find_map(ast::UseItem::cast) .map(AstNode::syntax), true, ); } }; // This can happen only if current_use_tree is a direct child of a UseItem if let Some(name) = alias.and_then(ast::NameOwner::name) { if compare_path_segment_with_name(&target[0], name) { return ImportAction::Nothing; } } collect_path_segments_raw(current_path_segments, path); // We compare only the new segments added in the line just above. // The first prev_len segments were already compared in 'parent' recursive calls. let left = target.split_at(prev_len).1; let right = current_path_segments.split_at(prev_len).1; let common = compare_path_segments(left, right); let mut action = match common { 0 => ImportAction::add_new_use( // e.g: target is std::fmt and we can have // use foo::bar // We add a brand new use statement current_use_tree.syntax().ancestors().find_map(ast::UseItem::cast).map(AstNode::syntax), true, ), common if common == left.len() && left.len() == right.len() => { // e.g: target is std::fmt and we can have // 1- use std::fmt; // 2- use std::fmt:{ ... } if let Some(list) = tree_list { // In case 2 we need to add self to the nested list // unless it's already there let has_self = list.use_trees().map(ast::UseTree::path).any(|p| { p.and_then(ast::Path::segment) .and_then(ast::PathSegment::kind) .filter(|k| *k == ast::PathSegmentKind::SelfKw) .is_some() }); if has_self { ImportAction::Nothing } else { ImportAction::add_in_tree_list(current_path_segments.len(), list, true) } } else { // Case 1 ImportAction::Nothing } } common if common != left.len() && left.len() == right.len() => { // e.g: target is std::fmt and we have // use std::io; // We need to split. let segments_to_split = current_path_segments.split_at(prev_len + common).1; ImportAction::add_nested_import( prev_len + common, path, Some(segments_to_split[0]), false, ) } common if common == right.len() && left.len() > right.len() => { // e.g: target is std::fmt and we can have // 1- use std; // 2- use std::{ ... }; // fallback action let mut better_action = ImportAction::add_new_use( current_use_tree .syntax() .ancestors() .find_map(ast::UseItem::cast) .map(AstNode::syntax), true, ); if let Some(list) = tree_list { // Case 2, check recursively if the path is already imported in the nested list for u in list.use_trees() { let child_action = walk_use_tree_for_best_action(current_path_segments, Some(list), u, target); if child_action.is_better(&better_action) { better_action = child_action; if let ImportAction::Nothing = better_action { return better_action; } } } } else { // Case 1, split adding self better_action = ImportAction::add_nested_import(prev_len + common, path, None, true) } better_action } common if common == left.len() && left.len() < right.len() => { // e.g: target is std::fmt and we can have // use std::fmt::Debug; let segments_to_split = current_path_segments.split_at(prev_len + common).1; ImportAction::add_nested_import( prev_len + common, path, Some(segments_to_split[0]), true, ) } common if common < left.len() && common < right.len() => { // e.g: target is std::fmt::nested::Debug // use std::fmt::Display let segments_to_split = current_path_segments.split_at(prev_len + common).1; ImportAction::add_nested_import( prev_len + common, path, Some(segments_to_split[0]), false, ) } _ => unreachable!(), }; // If we are inside a UseTreeList adding a use statement become adding to the existing // tree list. action = match (current_parent_use_tree_list, action) { (Some(use_tree_list), ImportAction::AddNewUse { .. }) => { ImportAction::add_in_tree_list(prev_len, use_tree_list, false) } (_, _) => action, }; // We remove the segments added current_path_segments.truncate(prev_len); return action; } fn best_action_for_target<'b, 'a: 'b>( container: &'a SyntaxNode, anchor: &'a SyntaxNode, target: &'b [SmolStr], ) -> ImportAction<'a> { let mut storage = Vec::with_capacity(16); // this should be the only allocation let best_action = container .children() .filter_map(ast::UseItem::cast) .filter_map(ast::UseItem::use_tree) .map(|u| walk_use_tree_for_best_action(&mut storage, None, u, target)) .fold(None, |best, a| { best.and_then(|best| Some(*ImportAction::better(&best, &a))).or_else(|| Some(a)) }); match best_action { Some(action) => return action, None => { // We have no action and no UseItem was found in container so we find // another item and we use it as anchor. // If there are no items above, we choose the target path itself as anchor. // todo: we should include even whitespace blocks as anchor candidates let anchor = container .children() .find(|n| n.range().start() < anchor.range().start()) .or_else(|| Some(anchor)); return ImportAction::add_new_use(anchor, false); } } } fn make_assist(action: &ImportAction, target: &[SmolStr], edit: &mut TextEditBuilder) { match action { ImportAction::AddNewUse { anchor, add_after_anchor } => { make_assist_add_new_use(anchor, *add_after_anchor, target, edit) } ImportAction::AddInTreeList { common_segments, tree_list, add_self } => { // We know that the fist n segments already exists in the use statement we want // to modify, so we want to add only the last target.len() - n segments. let segments_to_add = target.split_at(*common_segments).1; make_assist_add_in_tree_list(tree_list, segments_to_add, *add_self, edit) } ImportAction::AddNestedImport { common_segments, path_to_split, first_segment_to_split, add_self, } => { let segments_to_add = target.split_at(*common_segments).1; make_assist_add_nested_import( path_to_split, first_segment_to_split, segments_to_add, *add_self, edit, ) } _ => {} } } fn make_assist_add_new_use( anchor: &Option<&SyntaxNode>, after: bool, target: &[SmolStr], edit: &mut TextEditBuilder, ) { if let Some(anchor) = anchor { let indent = ra_fmt::leading_indent(anchor); let mut buf = String::new(); if after { buf.push_str("\n"); if let Some(spaces) = indent { buf.push_str(spaces); } } buf.push_str("use "); fmt_segments_raw(target, &mut buf); buf.push_str(";"); if !after { buf.push_str("\n\n"); if let Some(spaces) = indent { buf.push_str(spaces); } } let position = if after { anchor.range().end() } else { anchor.range().start() }; edit.insert(position, buf); } } fn make_assist_add_in_tree_list( tree_list: &ast::UseTreeList, target: &[SmolStr], add_self: bool, edit: &mut TextEditBuilder, ) { let last = tree_list.use_trees().last(); if let Some(last) = last { let mut buf = String::new(); let comma = last.syntax().siblings(Direction::Next).find(|n| n.kind() == T![,]); let offset = if let Some(comma) = comma { comma.range().end() } else { buf.push_str(","); last.syntax().range().end() }; if add_self { buf.push_str(" self") } else { buf.push_str(" "); } fmt_segments_raw(target, &mut buf); edit.insert(offset, buf); } else { } } fn make_assist_add_nested_import( path: &ast::Path, first_segment_to_split: &Option<&ast::PathSegment>, target: &[SmolStr], add_self: bool, edit: &mut TextEditBuilder, ) { let use_tree = path.syntax().ancestors().find_map(ast::UseTree::cast); if let Some(use_tree) = use_tree { let (start, add_colon_colon) = if let Some(first_segment_to_split) = first_segment_to_split { (first_segment_to_split.syntax().range().start(), false) } else { (use_tree.syntax().range().end(), true) }; let end = use_tree.syntax().range().end(); let mut buf = String::new(); if add_colon_colon { buf.push_str("::"); } buf.push_str("{ "); if add_self { buf.push_str("self, "); } fmt_segments_raw(target, &mut buf); if !target.is_empty() { buf.push_str(", "); } edit.insert(start, buf); edit.insert(end, "}".to_string()); } } fn apply_auto_import( container: &SyntaxNode, path: &ast::Path, target: &[SmolStr], edit: &mut TextEditBuilder, ) { let action = best_action_for_target(container, path.syntax(), target); make_assist(&action, target, edit); if let Some(last) = path.segment() { // Here we are assuming the assist will provide a correct use statement // so we can delete the path qualifier edit.delete(TextRange::from_to( path.syntax().range().start(), last.syntax().range().start(), )); } } pub fn collect_hir_path_segments(path: &hir::Path) -> Vec { let mut ps = Vec::::with_capacity(10); match path.kind { hir::PathKind::Abs => ps.push("".into()), hir::PathKind::Crate => ps.push("crate".into()), hir::PathKind::Plain => {} hir::PathKind::Self_ => ps.push("self".into()), hir::PathKind::Super => ps.push("super".into()), } for s in path.segments.iter() { ps.push(s.name.to_string().into()); } ps } // This function produces sequence of text edits into edit // to import the target path in the most appropriate scope given // the cursor position pub fn auto_import_text_edit( // Ideally the position of the cursor, used to position: &SyntaxNode, // The statement to use as anchor (last resort) anchor: &SyntaxNode, // The path to import as a sequence of strings target: &[SmolStr], edit: &mut TextEditBuilder, ) { let container = position.ancestors().find_map(|n| { if let Some(module) = ast::Module::cast(n) { return module.item_list().map(ast::AstNode::syntax); } ast::SourceFile::cast(n).map(ast::AstNode::syntax) }); if let Some(container) = container { let action = best_action_for_target(container, anchor, target); make_assist(&action, target, edit); } } pub(crate) fn auto_import(mut ctx: AssistCtx) -> Option { let path: &ast::Path = ctx.node_at_offset()?; // We don't want to mess with use statements if path.syntax().ancestors().find_map(ast::UseItem::cast).is_some() { return None; } let hir_path = hir::Path::from_ast(path)?; let segments = collect_hir_path_segments(&hir_path); if segments.len() < 2 { return None; } if let Some(module) = path.syntax().ancestors().find_map(ast::Module::cast) { if let (Some(item_list), Some(name)) = (module.item_list(), module.name()) { ctx.add_action( AssistId("auto_import"), format!("import {} in mod {}", fmt_segments(&segments), name.text()), |edit| { let mut text_edit = TextEditBuilder::default(); apply_auto_import(item_list.syntax(), path, &segments, &mut text_edit); edit.set_edit_builder(text_edit); }, ); } } else { let current_file = path.syntax().ancestors().find_map(ast::SourceFile::cast)?; ctx.add_action( AssistId("auto_import"), format!("import {} in the current file", fmt_segments(&segments)), |edit| { let mut text_edit = TextEditBuilder::default(); apply_auto_import(current_file.syntax(), path, &segments, &mut text_edit); edit.set_edit_builder(text_edit); }, ); } ctx.build() } #[cfg(test)] mod tests { use super::*; use crate::helpers::{check_assist, check_assist_not_applicable}; #[test] fn test_auto_import_add_use_no_anchor() { check_assist( auto_import, " std::fmt::Debug<|> ", " use std::fmt::Debug; Debug<|> ", ); } #[test] fn test_auto_import_add_use_no_anchor_with_item_below() { check_assist( auto_import, " std::fmt::Debug<|> fn main() { } ", " use std::fmt::Debug; Debug<|> fn main() { } ", ); } #[test] fn test_auto_import_add_use_no_anchor_with_item_above() { check_assist( auto_import, " fn main() { } std::fmt::Debug<|> ", " use std::fmt::Debug; fn main() { } Debug<|> ", ); } #[test] fn test_auto_import_add_use_no_anchor_2seg() { check_assist( auto_import, " std::fmt<|>::Debug ", " use std::fmt; fmt<|>::Debug ", ); } #[test] fn test_auto_import_add_use() { check_assist( auto_import, " use stdx; impl std::fmt::Debug<|> for Foo { } ", " use stdx; use std::fmt::Debug; impl Debug<|> for Foo { } ", ); } #[test] fn test_auto_import_file_use_other_anchor() { check_assist( auto_import, " impl std::fmt::Debug<|> for Foo { } ", " use std::fmt::Debug; impl Debug<|> for Foo { } ", ); } #[test] fn test_auto_import_add_use_other_anchor_indent() { check_assist( auto_import, " impl std::fmt::Debug<|> for Foo { } ", " use std::fmt::Debug; impl Debug<|> for Foo { } ", ); } #[test] fn test_auto_import_split_different() { check_assist( auto_import, " use std::fmt; impl std::io<|> for Foo { } ", " use std::{ io, fmt}; impl io<|> for Foo { } ", ); } #[test] fn test_auto_import_split_self_for_use() { check_assist( auto_import, " use std::fmt; impl std::fmt::Debug<|> for Foo { } ", " use std::fmt::{ self, Debug, }; impl Debug<|> for Foo { } ", ); } #[test] fn test_auto_import_split_self_for_target() { check_assist( auto_import, " use std::fmt::Debug; impl std::fmt<|> for Foo { } ", " use std::fmt::{ self, Debug}; impl fmt<|> for Foo { } ", ); } #[test] fn test_auto_import_add_to_nested_self_nested() { check_assist( auto_import, " use std::fmt::{Debug, nested::{Display}}; impl std::fmt::nested<|> for Foo { } ", " use std::fmt::{Debug, nested::{Display, self}}; impl nested<|> for Foo { } ", ); } #[test] fn test_auto_import_add_to_nested_self_already_included() { check_assist( auto_import, " use std::fmt::{Debug, nested::{self, Display}}; impl std::fmt::nested<|> for Foo { } ", " use std::fmt::{Debug, nested::{self, Display}}; impl nested<|> for Foo { } ", ); } #[test] fn test_auto_import_add_to_nested_nested() { check_assist( auto_import, " use std::fmt::{Debug, nested::{Display}}; impl std::fmt::nested::Debug<|> for Foo { } ", " use std::fmt::{Debug, nested::{Display, Debug}}; impl Debug<|> for Foo { } ", ); } #[test] fn test_auto_import_split_common_target_longer() { check_assist( auto_import, " use std::fmt::Debug; impl std::fmt::nested::Display<|> for Foo { } ", " use std::fmt::{ nested::Display, Debug}; impl Display<|> for Foo { } ", ); } #[test] fn test_auto_import_split_common_use_longer() { check_assist( auto_import, " use std::fmt::nested::Debug; impl std::fmt::Display<|> for Foo { } ", " use std::fmt::{ Display, nested::Debug}; impl Display<|> for Foo { } ", ); } #[test] fn test_auto_import_alias() { check_assist( auto_import, " use std::fmt as foo; impl foo::Debug<|> for Foo { } ", " use std::fmt as foo; impl Debug<|> for Foo { } ", ); } #[test] fn test_auto_import_not_applicable_one_segment() { check_assist_not_applicable( auto_import, " impl foo<|> for Foo { } ", ); } #[test] fn test_auto_import_not_applicable_in_use() { check_assist_not_applicable( auto_import, " use std::fmt<|>; ", ); } #[test] fn test_auto_import_add_use_no_anchor_in_mod_mod() { check_assist( auto_import, " mod foo { mod bar { std::fmt::Debug<|> } } ", " mod foo { mod bar { use std::fmt::Debug; Debug<|> } } ", ); } }