3 files changed, 222 insertions, 118 deletions

diff --git a/crates/assists/src/handlers/merge_imports.rs b/crates/assists/src/handlers/merge_imports.rs
index 35b884206..0bd679260 100644
--- a/crates/assists/src/handlers/merge_imports.rs
+++ b/crates/assists/src/handlers/merge_imports.rs
@@ -1,14 +1,14 @@
-use std::iter::successors;
-
 use syntax::{
-    algo::{neighbor, skip_trivia_token, SyntaxRewriter},
-    ast::{self, edit::AstNodeEdit, make},
-    AstNode, Direction, InsertPosition, SyntaxElement, T,
+    algo::{neighbor, SyntaxRewriter},
+    ast, AstNode,
 };
 
 use crate::{
     assist_context::{AssistContext, Assists},
-    utils::next_prev,
+    utils::{
+        insert_use::{try_merge_imports, try_merge_trees},
+        next_prev, MergeBehaviour,
+    },
     AssistId, AssistKind,
 };
 
@@ -30,23 +30,22 @@ pub(crate) fn merge_imports(acc: &mut Assists, ctx: &AssistContext) -> Option<()
     let mut offset = ctx.offset();
 
     if let Some(use_item) = tree.syntax().parent().and_then(ast::Use::cast) {
-        let (merged, to_delete) = next_prev()
-            .filter_map(|dir| neighbor(&use_item, dir))
-            .filter_map(|it| Some((it.clone(), it.use_tree()?)))
-            .find_map(|(use_item, use_tree)| {
-                Some((try_merge_trees(&tree, &use_tree)?, use_item))
+        let (merged, to_delete) =
+            next_prev().filter_map(|dir| neighbor(&use_item, dir)).find_map(|use_item2| {
+                try_merge_imports(&use_item, &use_item2, MergeBehaviour::Full).zip(Some(use_item2))
             })?;
 
-        rewriter.replace_ast(&tree, &merged);
+        rewriter.replace_ast(&use_item, &merged);
         rewriter += to_delete.remove();
 
         if to_delete.syntax().text_range().end() < offset {
             offset -= to_delete.syntax().text_range().len();
         }
     } else {
-        let (merged, to_delete) = next_prev()
-            .filter_map(|dir| neighbor(&tree, dir))
-            .find_map(|use_tree| Some((try_merge_trees(&tree, &use_tree)?, use_tree.clone())))?;
+        let (merged, to_delete) =
+            next_prev().filter_map(|dir| neighbor(&tree, dir)).find_map(|use_tree| {
+                try_merge_trees(&tree, &use_tree, MergeBehaviour::Full).zip(Some(use_tree))
+            })?;
 
         rewriter.replace_ast(&tree, &merged);
         rewriter += to_delete.remove();
@@ -67,66 +66,6 @@ pub(crate) fn merge_imports(acc: &mut Assists, ctx: &AssistContext) -> Option<()
     )
 }
 
-fn try_merge_trees(old: &ast::UseTree, new: &ast::UseTree) -> Option<ast::UseTree> {
-    let lhs_path = old.path()?;
-    let rhs_path = new.path()?;
-
-    let (lhs_prefix, rhs_prefix) = common_prefix(&lhs_path, &rhs_path)?;
-
-    let lhs = old.split_prefix(&lhs_prefix);
-    let rhs = new.split_prefix(&rhs_prefix);
-
-    let should_insert_comma = lhs
-        .use_tree_list()?
-        .r_curly_token()
-        .and_then(|it| skip_trivia_token(it.prev_token()?, Direction::Prev))
-        .map(|it| it.kind() != T![,])
-        .unwrap_or(true);
-
-    let mut to_insert: Vec<SyntaxElement> = Vec::new();
-    if should_insert_comma {
-        to_insert.push(make::token(T![,]).into());
-        to_insert.push(make::tokens::single_space().into());
-    }
-    to_insert.extend(
-        rhs.use_tree_list()?
-            .syntax()
-            .children_with_tokens()
-            .filter(|it| it.kind() != T!['{'] && it.kind() != T!['}']),
-    );
-    let use_tree_list = lhs.use_tree_list()?;
-    let pos = InsertPosition::Before(use_tree_list.r_curly_token()?.into());
-    let use_tree_list = use_tree_list.insert_children(pos, to_insert);
-    Some(lhs.with_use_tree_list(use_tree_list))
-}
-
-fn common_prefix(lhs: &ast::Path, rhs: &ast::Path) -> Option<(ast::Path, ast::Path)> {
-    let mut res = None;
-    let mut lhs_curr = first_path(&lhs);
-    let mut rhs_curr = first_path(&rhs);
-    loop {
-        match (lhs_curr.segment(), rhs_curr.segment()) {
-            (Some(lhs), Some(rhs)) if lhs.syntax().text() == rhs.syntax().text() => (),
-            _ => break,
-        }
-        res = Some((lhs_curr.clone(), rhs_curr.clone()));
-
-        match (lhs_curr.parent_path(), rhs_curr.parent_path()) {
-            (Some(lhs), Some(rhs)) => {
-                lhs_curr = lhs;
-                rhs_curr = rhs;
-            }
-            _ => break,
-        }
-    }
-
-    res
-}
-
-fn first_path(path: &ast::Path) -> ast::Path {
-    successors(Some(path.clone()), |it| it.qualifier()).last().unwrap()
-}
-
 #[cfg(test)]
 mod tests {
     use crate::tests::{check_assist, check_assist_not_applicable};
@@ -189,6 +128,78 @@ use std::{fmt::{Display, self}};
     }
 
     #[test]
+    fn skip_pub1() {
+        check_assist_not_applicable(
+            merge_imports,
+            r"
+pub use std::fmt<|>::Debug;
+use std::fmt::Display;
+",
+        );
+    }
+
+    #[test]
+    fn skip_pub_last() {
+        check_assist_not_applicable(
+            merge_imports,
+            r"
+use std::fmt<|>::Debug;
+pub use std::fmt::Display;
+",
+        );
+    }
+
+    #[test]
+    fn skip_pub_crate_pub() {
+        check_assist_not_applicable(
+            merge_imports,
+            r"
+pub(crate) use std::fmt<|>::Debug;
+pub use std::fmt::Display;
+",
+        );
+    }
+
+    #[test]
+    fn skip_pub_pub_crate() {
+        check_assist_not_applicable(
+            merge_imports,
+            r"
+pub use std::fmt<|>::Debug;
+pub(crate) use std::fmt::Display;
+",
+        );
+    }
+
+    #[test]
+    fn merge_pub() {
+        check_assist(
+            merge_imports,
+            r"
+pub use std::fmt<|>::Debug;
+pub use std::fmt::Display;
+",
+            r"
+pub use std::fmt::{Debug, Display};
+",
+        )
+    }
+
+    #[test]
+    fn merge_pub_crate() {
+        check_assist(
+            merge_imports,
+            r"
+pub(crate) use std::fmt<|>::Debug;
+pub(crate) use std::fmt::Display;
+",
+            r"
+pub(crate) use std::fmt::{Debug, Display};
+",
+        )
+    }
+
+    #[test]
     fn test_merge_nested() {
         check_assist(
             merge_imports,
diff --git a/crates/assists/src/handlers/remove_dbg.rs b/crates/assists/src/handlers/remove_dbg.rs
index 4e252edf0..a8ab2aecc 100644
--- a/crates/assists/src/handlers/remove_dbg.rs
+++ b/crates/assists/src/handlers/remove_dbg.rs
@@ -1,6 +1,6 @@
 use syntax::{
     ast::{self, AstNode},
-    TextRange, TextSize, T,
+    SyntaxElement, TextRange, TextSize, T,
 };
 
 use crate::{AssistContext, AssistId, AssistKind, Assists};
@@ -22,62 +22,108 @@ use crate::{AssistContext, AssistId, AssistKind, Assists};
 // ```
 pub(crate) fn remove_dbg(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
     let macro_call = ctx.find_node_at_offset::<ast::MacroCall>()?;
+    let new_contents = adjusted_macro_contents(&macro_call)?;
 
-    if !is_valid_macrocall(&macro_call, "dbg")? {
-        return None;
-    }
-
-    let is_leaf = macro_call.syntax().next_sibling().is_none();
-
+    let macro_text_range = macro_call.syntax().text_range();
     let macro_end = if macro_call.semicolon_token().is_some() {
-        macro_call.syntax().text_range().end() - TextSize::of(';')
+        macro_text_range.end() - TextSize::of(';')
     } else {
-        macro_call.syntax().text_range().end()
+        macro_text_range.end()
     };
 
-    // macro_range determines what will be deleted and replaced with macro_content
-    let macro_range = TextRange::new(macro_call.syntax().text_range().start(), macro_end);
-    let paste_instead_of_dbg = {
-        let text = macro_call.token_tree()?.syntax().text();
-
-        // leafiness determines if we should include the parenthesis or not
-        let slice_index: TextRange = if is_leaf {
-            // leaf means - we can extract the contents of the dbg! in text
-            TextRange::new(TextSize::of('('), text.len() - TextSize::of(')'))
-        } else {
-            // not leaf - means we should keep the parens
-            TextRange::up_to(text.len())
-        };
-        text.slice(slice_index).to_string()
-    };
+    acc.add(
+        AssistId("remove_dbg", AssistKind::Refactor),
+        "Remove dbg!()",
+        macro_text_range,
+        |builder| {
+            builder.replace(TextRange::new(macro_text_range.start(), macro_end), new_contents);
+        },
+    )
+}
 
-    let target = macro_call.syntax().text_range();
-    acc.add(AssistId("remove_dbg", AssistKind::Refactor), "Remove dbg!()", target, |builder| {
-        builder.replace(macro_range, paste_instead_of_dbg);
+fn adjusted_macro_contents(macro_call: &ast::MacroCall) -> Option<String> {
+    let contents = get_valid_macrocall_contents(&macro_call, "dbg")?;
+    let macro_text_with_brackets = macro_call.token_tree()?.syntax().text();
+    let macro_text_in_brackets = macro_text_with_brackets.slice(TextRange::new(
+        TextSize::of('('),
+        macro_text_with_brackets.len() - TextSize::of(')'),
+    ));
+
+    let is_leaf = macro_call.syntax().next_sibling().is_none();
+    Some(if !is_leaf && needs_parentheses_around_macro_contents(contents) {
+        format!("({})", macro_text_in_brackets)
+    } else {
+        macro_text_in_brackets.to_string()
     })
 }
 
 /// Verifies that the given macro_call actually matches the given name
-/// and contains proper ending tokens
-fn is_valid_macrocall(macro_call: &ast::MacroCall, macro_name: &str) -> Option<bool> {
+/// and contains proper ending tokens, then returns the contents between the ending tokens
+fn get_valid_macrocall_contents(
+    macro_call: &ast::MacroCall,
+    macro_name: &str,
+) -> Option<Vec<SyntaxElement>> {
     let path = macro_call.path()?;
     let name_ref = path.segment()?.name_ref()?;
 
     // Make sure it is actually a dbg-macro call, dbg followed by !
     let excl = path.syntax().next_sibling_or_token()?;
-
     if name_ref.text() != macro_name || excl.kind() != T![!] {
         return None;
     }
 
-    let node = macro_call.token_tree()?.syntax().clone();
-    let first_child = node.first_child_or_token()?;
-    let last_child = node.last_child_or_token()?;
+    let mut children_with_tokens = macro_call.token_tree()?.syntax().children_with_tokens();
+    let first_child = children_with_tokens.next()?;
+    let mut contents_between_brackets = children_with_tokens.collect::<Vec<_>>();
+    let last_child = contents_between_brackets.pop()?;
+
+    if contents_between_brackets.is_empty() {
+        None
+    } else {
+        match (first_child.kind(), last_child.kind()) {
+            (T!['('], T![')']) | (T!['['], T![']']) | (T!['{'], T!['}']) => {
+                Some(contents_between_brackets)
+            }
+            _ => None,
+        }
+    }
+}
 
-    match (first_child.kind(), last_child.kind()) {
-        (T!['('], T![')']) | (T!['['], T![']']) | (T!['{'], T!['}']) => Some(true),
-        _ => Some(false),
+fn needs_parentheses_around_macro_contents(macro_contents: Vec<SyntaxElement>) -> bool {
+    if macro_contents.len() < 2 {
+        return false;
     }
+    let mut unpaired_brackets_in_contents = Vec::new();
+    for element in macro_contents {
+        match element.kind() {
+            T!['('] | T!['['] | T!['{'] => unpaired_brackets_in_contents.push(element),
+            T![')'] => {
+                if !matches!(unpaired_brackets_in_contents.pop(), Some(correct_bracket) if correct_bracket.kind() == T!['('])
+                {
+                    return true;
+                }
+            }
+            T![']'] => {
+                if !matches!(unpaired_brackets_in_contents.pop(), Some(correct_bracket) if correct_bracket.kind() == T!['['])
+                {
+                    return true;
+                }
+            }
+            T!['}'] => {
+                if !matches!(unpaired_brackets_in_contents.pop(), Some(correct_bracket) if correct_bracket.kind() == T!['{'])
+                {
+                    return true;
+                }
+            }
+            symbol_kind => {
+                let symbol_not_in_bracket = unpaired_brackets_in_contents.is_empty();
+                if symbol_not_in_bracket && symbol_kind.is_punct() {
+                    return true;
+                }
+            }
+        }
+    }
+    !unpaired_brackets_in_contents.is_empty()
 }
 
 #[cfg(test)]
@@ -157,12 +203,38 @@ fn foo(n: usize) {
     }
 
     #[test]
+    fn remove_dbg_from_non_leaf_simple_expression() {
+        check_assist(
+            remove_dbg,
+            "
+fn main() {
+    let mut a = 1;
+    while dbg!<|>(a) < 10000 {
+        a += 1;
+    }
+}
+",
+            "
+fn main() {
+    let mut a = 1;
+    while a < 10000 {
+        a += 1;
+    }
+}
+",
+        );
+    }
+
+    #[test]
     fn test_remove_dbg_keep_expression() {
         check_assist(
             remove_dbg,
             r#"let res = <|>dbg!(a + b).foo();"#,
             r#"let res = (a + b).foo();"#,
         );
+
+        check_assist(remove_dbg, r#"let res = <|>dbg!(2 + 2) * 5"#, r#"let res = (2 + 2) * 5"#);
+        check_assist(remove_dbg, r#"let res = <|>dbg![2 + 2] * 5"#, r#"let res = (2 + 2) * 5"#);
     }
 
     #[test]
diff --git a/crates/assists/src/utils/insert_use.rs b/crates/assists/src/utils/insert_use.rs
index 8a4c8520d..98553b2e0 100644
--- a/crates/assists/src/utils/insert_use.rs
+++ b/crates/assists/src/utils/insert_use.rs
@@ -138,13 +138,23 @@ pub(crate) fn insert_use(
     algo::insert_children(scope.as_syntax_node(), insert_position, to_insert)
 }
 
-fn try_merge_imports(
+fn eq_visibility(vis0: Option<ast::Visibility>, vis1: Option<ast::Visibility>) -> bool {
+    match (vis0, vis1) {
+        (None, None) => true,
+        // FIXME: Don't use the string representation to check for equality
+        // spaces inside of the node would break this comparison
+        (Some(vis0), Some(vis1)) => vis0.to_string() == vis1.to_string(),
+        _ => false,
+    }
+}
+
+pub(crate) fn try_merge_imports(
     old: &ast::Use,
     new: &ast::Use,
     merge_behaviour: MergeBehaviour,
 ) -> Option<ast::Use> {
-    // don't merge into re-exports
-    if old.visibility().and_then(|vis| vis.pub_token()).is_some() {
+    // don't merge imports with different visibilities
+    if !eq_visibility(old.visibility(), new.visibility()) {
         return None;
     }
     let old_tree = old.use_tree()?;
@@ -161,7 +171,7 @@ fn use_tree_list_is_nested(tl: &ast::UseTreeList) -> bool {
 }
 
 // FIXME: currently this merely prepends the new tree into old, ideally it would insert the items in a sorted fashion
-pub fn try_merge_trees(
+pub(crate) fn try_merge_trees(
     old: &ast::UseTree,
     new: &ast::UseTree,
     merge_behaviour: MergeBehaviour,
@@ -278,7 +288,8 @@ fn first_path(path: &ast::Path) -> ast::Path {
 }
 
 fn segment_iter(path: &ast::Path) -> impl Iterator<Item = ast::PathSegment> + Clone {
-    path.syntax().children().flat_map(ast::PathSegment::cast)
+    // cant make use of SyntaxNode::siblings, because the returned Iterator is not clone
+    successors(first_segment(path), |p| p.parent_path().parent_path().and_then(|p| p.segment()))
 }
 
 #[derive(PartialEq, Eq)]
@@ -684,8 +695,18 @@ use std::io;",
         check_last(
             "foo::bar",
             r"use foo::bar::baz::Qux;",
-            r"use foo::bar::baz::Qux;
-use foo::bar;",
+            r"use foo::bar;
+use foo::bar::baz::Qux;",
+        );
+    }
+
+    #[test]
+    fn insert_short_before_long() {
+        check_none(
+            "foo::bar",
+            r"use foo::bar::baz::Qux;",
+            r"use foo::bar;
+use foo::bar::baz::Qux;",
         );
     }