crates/ra_assists/src/handlers/merge_imports.rs

   1 use std::iter::successors;
   2
   3 use ra_syntax::{
   4     algo::{neighbor, skip_trivia_token, SyntaxRewriter},
   5     ast::{self, edit::AstNodeEdit, make},
   6     AstNode, Direction, InsertPosition, SyntaxElement, T,
   7 };
   8
   9 use crate::{Assist, AssistCtx, AssistId};
  10
  11 // Assist: merge_imports
  12 //
  13 // Merges two imports with a common prefix.
  14 //
  15 // ```
  16 // use std::<|>fmt::Formatter;
  17 // use std::io;
  18 // ```
  19 // ->
  20 // ```
  21 // use std::{fmt::Formatter, io};
  22 // ```
  23 pub(crate) fn merge_imports(ctx: AssistCtx) -> Option<Assist> {
  24     let tree: ast::UseTree = ctx.find_node_at_offset()?;
  25     let mut rewriter = SyntaxRewriter::default();
  26     let mut offset = ctx.frange.range.start();
  27
  28     if let Some(use_item) = tree.syntax().parent().and_then(ast::UseItem::cast) {
  29         let (merged, to_delete) = next_prev()
  30             .filter_map(|dir| neighbor(&use_item, dir))
  31             .filter_map(|it| Some((it.clone(), it.use_tree()?)))
  32             .find_map(|(use_item, use_tree)| {
  33                 Some((try_merge_trees(&tree, &use_tree)?, use_item))
  34             })?;
  35
  36         rewriter.replace_ast(&tree, &merged);
  37         rewriter += to_delete.remove();
  38
  39         if to_delete.syntax().text_range().end() < offset {
  40             offset -= to_delete.syntax().text_range().len();
  41         }
  42     } else {
  43         let (merged, to_delete) = next_prev()
  44             .filter_map(|dir| neighbor(&tree, dir))
  45             .find_map(|use_tree| Some((try_merge_trees(&tree, &use_tree)?, use_tree.clone())))?;
  46
  47         rewriter.replace_ast(&tree, &merged);
  48         rewriter += to_delete.remove();
  49
  50         if to_delete.syntax().text_range().end() < offset {
  51             offset -= to_delete.syntax().text_range().len();
  52         }
  53     };
  54
  55     ctx.add_assist(AssistId("merge_imports"), "Merge imports", |edit| {
  56         edit.rewrite(rewriter);
  57         // FIXME: we only need because our diff is imprecise
  58         edit.set_cursor(offset);
  59     })
  60 }
  61
  62 fn next_prev() -> impl Iterator<Item = Direction> {
  63     [Direction::Next, Direction::Prev].iter().copied()
  64 }
  65
  66 fn try_merge_trees(old: &ast::UseTree, new: &ast::UseTree) -> Option<ast::UseTree> {
  67     let lhs_path = old.path()?;
  68     let rhs_path = new.path()?;
  69
  70     let (lhs_prefix, rhs_prefix) = common_prefix(&lhs_path, &rhs_path)?;
  71
  72     let lhs = old.split_prefix(&lhs_prefix);
  73     let rhs = new.split_prefix(&rhs_prefix);
  74
  75     let should_insert_comma = lhs
  76         .use_tree_list()?
  77         .r_curly_token()
  78         .and_then(|it| skip_trivia_token(it.prev_token()?, Direction::Prev))
  79         .map(|it| it.kind() != T![,])
  80         .unwrap_or(true);
  81
  82     let mut to_insert: Vec<SyntaxElement> = Vec::new();
  83     if should_insert_comma {
  84         to_insert.push(make::token(T![,]).into());
  85         to_insert.push(make::tokens::single_space().into());
  86     }
  87     to_insert.extend(
  88         rhs.use_tree_list()?
  89             .syntax()
  90             .children_with_tokens()
  91             .filter(|it| it.kind() != T!['{'] && it.kind() != T!['}']),
  92     );
  93     let use_tree_list = lhs.use_tree_list()?;
  94     let pos = InsertPosition::Before(use_tree_list.r_curly_token()?.into());
  95     let use_tree_list = use_tree_list.insert_children(pos, to_insert);
  96     Some(lhs.with_use_tree_list(use_tree_list))
  97 }
  98
  99 fn common_prefix(lhs: &ast::Path, rhs: &ast::Path) -> Option<(ast::Path, ast::Path)> {
 100     let mut res = None;
 101     let mut lhs_curr = first_path(&lhs);
 102     let mut rhs_curr = first_path(&rhs);
 103     loop {
 104         match (lhs_curr.segment(), rhs_curr.segment()) {
 105             (Some(lhs), Some(rhs)) if lhs.syntax().text() == rhs.syntax().text() => (),
 106             _ => break,
 107         }
 108         res = Some((lhs_curr.clone(), rhs_curr.clone()));
 109
 110         match (lhs_curr.parent_path(), rhs_curr.parent_path()) {
 111             (Some(lhs), Some(rhs)) => {
 112                 lhs_curr = lhs;
 113                 rhs_curr = rhs;
 114             }
 115             _ => break,
 116         }
 117     }
 118
 119     res
 120 }
 121
 122 fn first_path(path: &ast::Path) -> ast::Path {
 123     successors(Some(path.clone()), |it| it.qualifier()).last().unwrap()
 124 }
 125
 126 #[cfg(test)]
 127 mod tests {
 128     use crate::tests::check_assist;
 129
 130     use super::*;
 131
 132     #[test]
 133     fn test_merge_first() {
 134         check_assist(
 135             merge_imports,
 136             r"
 137 use std::fmt<|>::Debug;
 138 use std::fmt::Display;
 139 ",
 140             r"
 141 use std::fmt<|>::{Debug, Display};
 142 ",
 143         )
 144     }
 145
 146     #[test]
 147     fn test_merge_second() {
 148         check_assist(
 149             merge_imports,
 150             r"
 151 use std::fmt::Debug;
 152 use std::fmt<|>::Display;
 153 ",
 154             r"
 155 use std::fmt:<|>:{Display, Debug};
 156 ",
 157         );
 158     }
 159
 160     #[test]
 161     fn test_merge_nested() {
 162         check_assist(
 163             merge_imports,
 164             r"
 165 use std::{fmt<|>::Debug, fmt::Display};
 166 ",
 167             r"
 168 use std::{fmt<|>::{Debug, Display}};
 169 ",
 170         );
 171         check_assist(
 172             merge_imports,
 173             r"
 174 use std::{fmt::Debug, fmt<|>::Display};
 175 ",
 176             r"
 177 use std::{fmt::<|>{Display, Debug}};
 178 ",
 179         );
 180     }
 181
 182     #[test]
 183     fn test_merge_single_wildcard_diff_prefixes() {
 184         check_assist(
 185             merge_imports,
 186             r"
 187 use std<|>::cell::*;
 188 use std::str;
 189 ",
 190             r"
 191 use std<|>::{cell::*, str};
 192 ",
 193         )
 194     }
 195
 196     #[test]
 197     fn test_merge_both_wildcard_diff_prefixes() {
 198         check_assist(
 199             merge_imports,
 200             r"
 201 use std<|>::cell::*;
 202 use std::str::*;
 203 ",
 204             r"
 205 use std<|>::{cell::*, str::*};
 206 ",
 207         )
 208     }
 209
 210     #[test]
 211     fn removes_just_enough_whitespace() {
 212         check_assist(
 213             merge_imports,
 214             r"
 215 use foo<|>::bar;
 216 use foo::baz;
 217
 218 /// Doc comment
 219 ",
 220             r"
 221 use foo<|>::{bar, baz};
 222
 223 /// Doc comment
 224 ",
 225         );
 226     }
 227
 228     #[test]
 229     fn works_with_trailing_comma() {
 230         check_assist(
 231             merge_imports,
 232             r"
 233 use {
 234     foo<|>::bar,
 235     foo::baz,
 236 };
 237 ",
 238             r"
 239 use {
 240     foo<|>::{bar, baz},
 241 };
 242 ",
 243         );
 244         check_assist(
 245             merge_imports,
 246             r"
 247 use {
 248     foo::baz,
 249     foo<|>::bar,
 250 };
 251 ",
 252             r"
 253 use {
 254     foo::{bar<|>, baz},
 255 };
 256 ",
 257         );
 258     }
 259
 260     #[test]
 261     fn test_double_comma() {
 262         check_assist(
 263             merge_imports,
 264             r"
 265 use foo::bar::baz;
 266 use foo::<|>{
 267     FooBar,
 268 };
 269 ",
 270             r"
 271 use foo::{<|>
 272     FooBar,
 273 bar::baz};
 274 ",
 275         )
 276     }
 277 }