2 files changed, 543 insertions, 53 deletions

diff --git a/crates/syntax/src/algo.rs b/crates/syntax/src/algo.rs
index ea199f9b8..065035fe6 100644
--- a/crates/syntax/src/algo.rs
+++ b/crates/syntax/src/algo.rs
@@ -2,11 +2,14 @@
 
 use std::{
     fmt,
+    hash::BuildHasherDefault,
     ops::{self, RangeInclusive},
 };
 
+use indexmap::IndexMap;
 use itertools::Itertools;
 use rustc_hash::FxHashMap;
+use test_utils::mark;
 use text_edit::TextEditBuilder;
 
 use crate::{
@@ -106,42 +109,56 @@ pub enum InsertPosition<T> {
     After(T),
 }
 
+type FxIndexMap<K, V> = IndexMap<K, V, BuildHasherDefault<rustc_hash::FxHasher>>;
+
+#[derive(Debug)]
 pub struct TreeDiff {
     replacements: FxHashMap<SyntaxElement, SyntaxElement>,
+    deletions: Vec<SyntaxElement>,
+    // the vec as well as the indexmap are both here to preserve order
+    insertions: FxIndexMap<SyntaxElement, Vec<SyntaxElement>>,
 }
 
 impl TreeDiff {
     pub fn into_text_edit(&self, builder: &mut TextEditBuilder) {
+        for (anchor, to) in self.insertions.iter() {
+            to.iter().for_each(|to| builder.insert(anchor.text_range().end(), to.to_string()));
+        }
         for (from, to) in self.replacements.iter() {
             builder.replace(from.text_range(), to.to_string())
         }
+        for text_range in self.deletions.iter().map(SyntaxElement::text_range) {
+            builder.delete(text_range);
+        }
     }
 
     pub fn is_empty(&self) -> bool {
-        self.replacements.is_empty()
+        self.replacements.is_empty() && self.deletions.is_empty() && self.insertions.is_empty()
     }
 }
 
 /// Finds minimal the diff, which, applied to `from`, will result in `to`.
 ///
-/// Specifically, returns a map whose keys are descendants of `from` and values
-/// are descendants of `to`, such that  `replace_descendants(from, map) == to`.
+/// Specifically, returns a structure that consists of a replacements, insertions and deletions
+/// such that applying this map on `from` will result in `to`.
 ///
-/// A trivial solution is a singleton map `{ from: to }`, but this function
-/// tries to find a more fine-grained diff.
+/// This function tries to find a fine-grained diff.
 pub fn diff(from: &SyntaxNode, to: &SyntaxNode) -> TreeDiff {
-    let mut buf = FxHashMap::default();
-    // FIXME: this is both horrible inefficient and gives larger than
-    // necessary diff. I bet there's a cool algorithm to diff trees properly.
-    go(&mut buf, from.clone().into(), to.clone().into());
-    return TreeDiff { replacements: buf };
-
-    fn go(
-        buf: &mut FxHashMap<SyntaxElement, SyntaxElement>,
-        lhs: SyntaxElement,
-        rhs: SyntaxElement,
-    ) {
-        if lhs.kind() == rhs.kind()
+    let mut diff = TreeDiff {
+        replacements: FxHashMap::default(),
+        insertions: FxIndexMap::default(),
+        deletions: Vec::new(),
+    };
+    let (from, to) = (from.clone().into(), to.clone().into());
+
+    // FIXME: this is horrible inefficient. I bet there's a cool algorithm to diff trees properly.
+    if !syntax_element_eq(&from, &to) {
+        go(&mut diff, from, to);
+    }
+    return diff;
+
+    fn syntax_element_eq(lhs: &SyntaxElement, rhs: &SyntaxElement) -> bool {
+        lhs.kind() == rhs.kind()
             && lhs.text_range().len() == rhs.text_range().len()
             && match (&lhs, &rhs) {
                 (NodeOrToken::Node(lhs), NodeOrToken::Node(rhs)) => {
@@ -150,18 +167,47 @@ pub fn diff(from: &SyntaxNode, to: &SyntaxNode) -> TreeDiff {
                 (NodeOrToken::Token(lhs), NodeOrToken::Token(rhs)) => lhs.text() == rhs.text(),
                 _ => false,
             }
-        {
-            return;
-        }
-        if let (Some(lhs), Some(rhs)) = (lhs.as_node(), rhs.as_node()) {
-            if lhs.children_with_tokens().count() == rhs.children_with_tokens().count() {
-                for (lhs, rhs) in lhs.children_with_tokens().zip(rhs.children_with_tokens()) {
-                    go(buf, lhs, rhs)
-                }
+    }
+
+    fn go(diff: &mut TreeDiff, lhs: SyntaxElement, rhs: SyntaxElement) {
+        let (lhs, rhs) = match lhs.as_node().zip(rhs.as_node()) {
+            Some((lhs, rhs)) => (lhs, rhs),
+            _ => {
+                mark::hit!(diff_node_token_replace);
+                diff.replacements.insert(lhs, rhs);
                 return;
             }
+        };
+
+        let mut rhs_children = rhs.children_with_tokens();
+        let mut lhs_children = lhs.children_with_tokens();
+        let mut last_lhs = None;
+        loop {
+            let lhs_child = lhs_children.next();
+            match (lhs_child.clone(), rhs_children.next()) {
+                (None, None) => break,
+                (None, Some(element)) => match last_lhs.clone() {
+                    Some(prev) => {
+                        mark::hit!(diff_insert);
+                        diff.insertions.entry(prev).or_insert_with(Vec::new).push(element);
+                    }
+                    // first iteration, this means we got no anchor element to insert after
+                    // therefor replace the parent node instead
+                    None => {
+                        mark::hit!(diff_replace_parent);
+                        diff.replacements.insert(lhs.clone().into(), rhs.clone().into());
+                        break;
+                    }
+                },
+                (Some(element), None) => {
+                    mark::hit!(diff_delete);
+                    diff.deletions.push(element);
+                }
+                (Some(ref lhs_ele), Some(ref rhs_ele)) if syntax_element_eq(lhs_ele, rhs_ele) => {}
+                (Some(lhs_ele), Some(rhs_ele)) => go(diff, lhs_ele, rhs_ele),
+            }
+            last_lhs = lhs_child.or(last_lhs);
         }
-        buf.insert(lhs, rhs);
     }
 }
 
@@ -243,11 +289,19 @@ fn _replace_children(
     with_children(parent, new_children)
 }
 
+#[derive(Debug, PartialEq, Eq, Hash)]
+enum InsertPos {
+    FirstChildOf(SyntaxNode),
+    // Before(SyntaxElement),
+    After(SyntaxElement),
+}
+
 #[derive(Default)]
 pub struct SyntaxRewriter<'a> {
     f: Option<Box<dyn Fn(&SyntaxElement) -> Option<SyntaxElement> + 'a>>,
     //FIXME: add debug_assertions that all elements are in fact from the same file.
     replacements: FxHashMap<SyntaxElement, Replacement>,
+    insertions: IndexMap<InsertPos, Vec<SyntaxElement>>,
 }
 
 impl fmt::Debug for SyntaxRewriter<'_> {
@@ -258,13 +312,96 @@ impl fmt::Debug for SyntaxRewriter<'_> {
 
 impl<'a> SyntaxRewriter<'a> {
     pub fn from_fn(f: impl Fn(&SyntaxElement) -> Option<SyntaxElement> + 'a) -> SyntaxRewriter<'a> {
-        SyntaxRewriter { f: Some(Box::new(f)), replacements: FxHashMap::default() }
+        SyntaxRewriter {
+            f: Some(Box::new(f)),
+            replacements: FxHashMap::default(),
+            insertions: IndexMap::default(),
+        }
     }
     pub fn delete<T: Clone + Into<SyntaxElement>>(&mut self, what: &T) {
         let what = what.clone().into();
         let replacement = Replacement::Delete;
         self.replacements.insert(what, replacement);
     }
+    pub fn insert_before<T: Clone + Into<SyntaxElement>, U: Clone + Into<SyntaxElement>>(
+        &mut self,
+        before: &T,
+        what: &U,
+    ) {
+        let before = before.clone().into();
+        let pos = match before.prev_sibling_or_token() {
+            Some(sibling) => InsertPos::After(sibling),
+            None => match before.parent() {
+                Some(parent) => InsertPos::FirstChildOf(parent),
+                None => return,
+            },
+        };
+        self.insertions.entry(pos).or_insert_with(Vec::new).push(what.clone().into());
+    }
+    pub fn insert_after<T: Clone + Into<SyntaxElement>, U: Clone + Into<SyntaxElement>>(
+        &mut self,
+        after: &T,
+        what: &U,
+    ) {
+        self.insertions
+            .entry(InsertPos::After(after.clone().into()))
+            .or_insert_with(Vec::new)
+            .push(what.clone().into());
+    }
+    pub fn insert_as_first_child<T: Clone + Into<SyntaxNode>, U: Clone + Into<SyntaxElement>>(
+        &mut self,
+        parent: &T,
+        what: &U,
+    ) {
+        self.insertions
+            .entry(InsertPos::FirstChildOf(parent.clone().into()))
+            .or_insert_with(Vec::new)
+            .push(what.clone().into());
+    }
+    pub fn insert_many_before<
+        T: Clone + Into<SyntaxElement>,
+        U: IntoIterator<Item = SyntaxElement>,
+    >(
+        &mut self,
+        before: &T,
+        what: U,
+    ) {
+        let before = before.clone().into();
+        let pos = match before.prev_sibling_or_token() {
+            Some(sibling) => InsertPos::After(sibling),
+            None => match before.parent() {
+                Some(parent) => InsertPos::FirstChildOf(parent),
+                None => return,
+            },
+        };
+        self.insertions.entry(pos).or_insert_with(Vec::new).extend(what);
+    }
+    pub fn insert_many_after<
+        T: Clone + Into<SyntaxElement>,
+        U: IntoIterator<Item = SyntaxElement>,
+    >(
+        &mut self,
+        after: &T,
+        what: U,
+    ) {
+        self.insertions
+            .entry(InsertPos::After(after.clone().into()))
+            .or_insert_with(Vec::new)
+            .extend(what);
+    }
+    pub fn insert_many_as_first_children<
+        T: Clone + Into<SyntaxNode>,
+        U: IntoIterator<Item = SyntaxElement>,
+    >(
+        &mut self,
+        parent: &T,
+        what: U,
+    ) {
+        self.insertions
+            .entry(InsertPos::FirstChildOf(parent.clone().into()))
+            .or_insert_with(Vec::new)
+            .extend(what)
+    }
     pub fn replace<T: Clone + Into<SyntaxElement>>(&mut self, what: &T, with: &T) {
         let what = what.clone().into();
         let replacement = Replacement::Single(with.clone().into());
@@ -284,7 +421,7 @@ impl<'a> SyntaxRewriter<'a> {
     }
 
     pub fn rewrite(&self, node: &SyntaxNode) -> SyntaxNode {
-        if self.f.is_none() && self.replacements.is_empty() {
+        if self.f.is_none() && self.replacements.is_empty() && self.insertions.is_empty() {
             return node.clone();
         }
         self.rewrite_children(node)
@@ -300,14 +437,22 @@ impl<'a> SyntaxRewriter<'a> {
     ///
     /// Returns `None` when there are no replacements.
     pub fn rewrite_root(&self) -> Option<SyntaxNode> {
+        fn element_to_node_or_parent(element: &SyntaxElement) -> SyntaxNode {
+            match element {
+                SyntaxElement::Node(it) => it.clone(),
+                SyntaxElement::Token(it) => it.parent(),
+            }
+        }
+
         assert!(self.f.is_none());
         self.replacements
             .keys()
-            .map(|element| match element {
-                SyntaxElement::Node(it) => it.clone(),
-                SyntaxElement::Token(it) => it.parent(),
-            })
-            // If we only have one replacement, we must return its parent node, since `rewrite` does
+            .map(element_to_node_or_parent)
+            .chain(self.insertions.keys().map(|pos| match pos {
+                InsertPos::FirstChildOf(it) => it.clone(),
+                InsertPos::After(it) => element_to_node_or_parent(it),
+            }))
+            // If we only have one replacement/insertion, we must return its parent node, since `rewrite` does
             // not replace the node passed to it.
             .map(|it| it.parent().unwrap_or(it))
             .fold1(|a, b| least_common_ancestor(&a, &b).unwrap())
@@ -321,9 +466,16 @@ impl<'a> SyntaxRewriter<'a> {
         self.replacements.get(element).cloned()
     }
 
+    fn insertions(&self, pos: &InsertPos) -> Option<impl Iterator<Item = SyntaxElement> + '_> {
+        self.insertions.get(pos).map(|insertions| insertions.iter().cloned())
+    }
+
     fn rewrite_children(&self, node: &SyntaxNode) -> SyntaxNode {
         //  FIXME: this could be made much faster.
         let mut new_children = Vec::new();
+        if let Some(elements) = self.insertions(&InsertPos::FirstChildOf(node.clone())) {
+            new_children.extend(elements.map(element_to_green));
+        }
         for child in node.children_with_tokens() {
             self.rewrite_self(&mut new_children, &child);
         }
@@ -337,34 +489,45 @@ impl<'a> SyntaxRewriter<'a> {
     ) {
         if let Some(replacement) = self.replacement(&element) {
             match replacement {
-                Replacement::Single(NodeOrToken::Node(it)) => {
-                    acc.push(NodeOrToken::Node(it.green().clone()))
-                }
-                Replacement::Single(NodeOrToken::Token(it)) => {
-                    acc.push(NodeOrToken::Token(it.green().clone()))
-                }
+                Replacement::Single(element) => acc.push(element_to_green(element)),
                 Replacement::Many(replacements) => {
-                    acc.extend(replacements.iter().map(|it| match it {
-                        NodeOrToken::Node(it) => NodeOrToken::Node(it.green().clone()),
-                        NodeOrToken::Token(it) => NodeOrToken::Token(it.green().clone()),
-                    }))
+                    acc.extend(replacements.into_iter().map(element_to_green))
                 }
                 Replacement::Delete => (),
             };
-            return;
+        } else {
+            match element {
+                NodeOrToken::Token(it) => acc.push(NodeOrToken::Token(it.green().clone())),
+                NodeOrToken::Node(it) => {
+                    acc.push(NodeOrToken::Node(self.rewrite_children(it).green().clone()));
+                }
+            }
+        }
+        if let Some(elements) = self.insertions(&InsertPos::After(element.clone())) {
+            acc.extend(elements.map(element_to_green));
         }
-        let res = match element {
-            NodeOrToken::Token(it) => NodeOrToken::Token(it.green().clone()),
-            NodeOrToken::Node(it) => NodeOrToken::Node(self.rewrite_children(it).green().clone()),
-        };
-        acc.push(res)
+    }
+}
+
+fn element_to_green(element: SyntaxElement) -> NodeOrToken<rowan::GreenNode, rowan::GreenToken> {
+    match element {
+        NodeOrToken::Node(it) => NodeOrToken::Node(it.green().clone()),
+        NodeOrToken::Token(it) => NodeOrToken::Token(it.green().clone()),
     }
 }
 
 impl ops::AddAssign for SyntaxRewriter<'_> {
     fn add_assign(&mut self, rhs: SyntaxRewriter) {
         assert!(rhs.f.is_none());
-        self.replacements.extend(rhs.replacements)
+        self.replacements.extend(rhs.replacements);
+        for (pos, insertions) in rhs.insertions.into_iter() {
+            match self.insertions.entry(pos) {
+                indexmap::map::Entry::Occupied(mut occupied) => {
+                    occupied.get_mut().extend(insertions)
+                }
+                indexmap::map::Entry::Vacant(vacant) => drop(vacant.insert(insertions)),
+            }
+        }
     }
 }
 
@@ -404,3 +567,322 @@ fn to_green_element(element: SyntaxElement) -> NodeOrToken<rowan::GreenNode, row
         NodeOrToken::Token(it) => it.green().clone().into(),
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use expect_test::{expect, Expect};
+    use itertools::Itertools;
+    use parser::SyntaxKind;
+    use test_utils::mark;
+    use text_edit::TextEdit;
+
+    use crate::{AstNode, SyntaxElement};
+
+    #[test]
+    fn replace_node_token() {
+        mark::check!(diff_node_token_replace);
+        check_diff(
+            r#"use node;"#,
+            r#"ident"#,
+            expect![[r#"
+                insertions:
+
+
+
+                replacements:
+
+                Line 0: Token([email protected] "use") -> ident
+
+                deletions:
+
+                Line 1: " "
+                Line 1: node
+                Line 1: ;
+            "#]],
+        );
+    }
+
+    #[test]
+    fn insert() {
+        mark::check!(diff_insert);
+        check_diff(
+            r#"use foo;"#,
+            r#"use foo;
+use bar;"#,
+            expect![[r#"
+                insertions:
+
+                Line 0: Node([email protected])
+                -> "\n"
+                -> use bar;
+
+                replacements:
+
+
+
+                deletions:
+
+
+            "#]],
+        );
+    }
+
+    #[test]
+    fn replace_parent() {
+        mark::check!(diff_replace_parent);
+        check_diff(
+            r#""#,
+            r#"use foo::bar;"#,
+            expect![[r#"
+                insertions:
+
+
+
+                replacements:
+
+                Line 0: Node([email protected]) -> use foo::bar;
+
+                deletions:
+
+
+            "#]],
+        );
+    }
+
+    #[test]
+    fn delete() {
+        mark::check!(diff_delete);
+        check_diff(
+            r#"use foo;
+            use bar;"#,
+            r#"use foo;"#,
+            expect![[r#"
+                insertions:
+
+
+
+                replacements:
+
+
+
+                deletions:
+
+                Line 1: "\n            "
+                Line 2: use bar;
+            "#]],
+        );
+    }
+
+    #[test]
+    fn insert_use() {
+        check_diff(
+            r#"
+use expect_test::{expect, Expect};
+
+use crate::AstNode;
+"#,
+            r#"
+use expect_test::{expect, Expect};
+use text_edit::TextEdit;
+
+use crate::AstNode;
+"#,
+            expect![[r#"
+                insertions:
+
+                Line 4: Token([email protected] "\n")
+                -> use crate::AstNode;
+                -> "\n"
+
+                replacements:
+
+                Line 2: Token([email protected] "\n\n") -> "\n"
+                Line 4: Token([email protected] "crate") -> text_edit
+                Line 4: Token([email protected] "AstNode") -> TextEdit
+                Line 4: Token([email protected] "\n") -> "\n\n"
+
+                deletions:
+
+
+            "#]],
+        )
+    }
+
+    #[test]
+    fn remove_use() {
+        check_diff(
+            r#"
+use expect_test::{expect, Expect};
+use text_edit::TextEdit;
+
+use crate::AstNode;
+"#,
+            r#"
+use expect_test::{expect, Expect};
+
+use crate::AstNode;
+"#,
+            expect![[r#"
+                insertions:
+
+
+
+                replacements:
+
+                Line 2: Token([email protected] "\n") -> "\n\n"
+                Line 3: Node([email protected]) -> crate
+                Line 3: Token([email protected] "TextEdit") -> AstNode
+                Line 3: Token([email protected] "\n\n") -> "\n"
+
+                deletions:
+
+                Line 4: use crate::AstNode;
+                Line 5: "\n"
+            "#]],
+        )
+    }
+
+    #[test]
+    fn merge_use() {
+        check_diff(
+            r#"
+use std::{
+    fmt,
+    hash::BuildHasherDefault,
+    ops::{self, RangeInclusive},
+};
+"#,
+            r#"
+use std::fmt;
+use std::hash::BuildHasherDefault;
+use std::ops::{self, RangeInclusive};
+"#,
+            expect![[r#"
+                insertions:
+
+                Line 2: Node([email protected])
+                -> ::
+                -> fmt
+                Line 6: Token([email protected] "\n")
+                -> use std::hash::BuildHasherDefault;
+                -> "\n"
+                -> use std::ops::{self, RangeInclusive};
+                -> "\n"
+
+                replacements:
+
+                Line 2: Token([email protected] "std") -> std
+
+                deletions:
+
+                Line 2: ::
+                Line 2: {
+                    fmt,
+                    hash::BuildHasherDefault,
+                    ops::{self, RangeInclusive},
+                }
+            "#]],
+        )
+    }
+
+    #[test]
+    fn early_return_assist() {
+        check_diff(
+            r#"
+fn main() {
+    if let Ok(x) = Err(92) {
+        foo(x);
+    }
+}
+            "#,
+            r#"
+fn main() {
+    let x = match Err(92) {
+        Ok(it) => it,
+        _ => return,
+    };
+    foo(x);
+}
+            "#,
+            expect![[r#"
+                insertions:
+
+                Line 3: Node([email protected])
+                -> " "
+                -> match Err(92) {
+                        Ok(it) => it,
+                        _ => return,
+                    }
+                -> ;
+                Line 5: Token([email protected] "}")
+                -> "\n"
+                -> }
+
+                replacements:
+
+                Line 3: Token([email protected] "if") -> let
+                Line 3: Token([email protected] "let") -> x
+                Line 3: Node([email protected]) -> =
+                Line 5: Token([email protected] "\n") -> "\n    "
+                Line 5: Token([email protected] "}") -> foo(x);
+
+                deletions:
+
+                Line 3: " "
+                Line 3: Ok(x)
+                Line 3: " "
+                Line 3: =
+                Line 3: " "
+                Line 3: Err(92)
+            "#]],
+        )
+    }
+
+    fn check_diff(from: &str, to: &str, expected_diff: Expect) {
+        let from_node = crate::SourceFile::parse(from).tree().syntax().clone();
+        let to_node = crate::SourceFile::parse(to).tree().syntax().clone();
+        let diff = super::diff(&from_node, &to_node);
+
+        let line_number =
+            |syn: &SyntaxElement| from[..syn.text_range().start().into()].lines().count();
+
+        let fmt_syntax = |syn: &SyntaxElement| match syn.kind() {
+            SyntaxKind::WHITESPACE => format!("{:?}", syn.to_string()),
+            _ => format!("{}", syn),
+        };
+
+        let insertions = diff.insertions.iter().format_with("\n", |(k, v), f| {
+            f(&format!(
+                "Line {}: {:?}\n-> {}",
+                line_number(k),
+                k,
+                v.iter().format_with("\n-> ", |v, f| f(&fmt_syntax(v)))
+            ))
+        });
+
+        let replacements = diff
+            .replacements
+            .iter()
+            .sorted_by_key(|(syntax, _)| syntax.text_range().start())
+            .format_with("\n", |(k, v), f| {
+                f(&format!("Line {}: {:?} -> {}", line_number(k), k, fmt_syntax(v)))
+            });
+
+        let deletions = diff
+            .deletions
+            .iter()
+            .format_with("\n", |v, f| f(&format!("Line {}: {}", line_number(v), &fmt_syntax(v))));
+
+        let actual = format!(
+            "insertions:\n\n{}\n\nreplacements:\n\n{}\n\ndeletions:\n\n{}\n",
+            insertions, replacements, deletions
+        );
+        expected_diff.assert_eq(&actual);
+
+        let mut from = from.to_owned();
+        let mut text_edit = TextEdit::builder();
+        diff.into_text_edit(&mut text_edit);
+        text_edit.finish().apply(&mut from);
+        assert_eq!(&*from, to, "diff did not turn `from` to `to`");
+    }
+}
diff --git a/crates/syntax/src/ast/node_ext.rs b/crates/syntax/src/ast/node_ext.rs
index 50c1c157d..c5cd1c504 100644
--- a/crates/syntax/src/ast/node_ext.rs
+++ b/crates/syntax/src/ast/node_ext.rs
@@ -7,7 +7,7 @@ use itertools::Itertools;
 use parser::SyntaxKind;
 
 use crate::{
-    ast::{self, support, AstNode, NameOwner, SyntaxNode},
+    ast::{self, support, token_ext::HasStringValue, AstNode, AstToken, NameOwner, SyntaxNode},
     SmolStr, SyntaxElement, SyntaxToken, T,
 };
 
@@ -53,8 +53,16 @@ impl ast::Attr {
     pub fn as_simple_key_value(&self) -> Option<(SmolStr, SmolStr)> {
         let lit = self.literal()?;
         let key = self.simple_name()?;
-        // FIXME: escape? raw string?
-        let value = lit.syntax().first_token()?.text().trim_matches('"').into();
+        let value_token = lit.syntax().first_token()?;
+
+        let value: SmolStr = if let Some(s) = ast::String::cast(value_token.clone()) {
+            s.value()?.into()
+        } else if let Some(s) = ast::RawString::cast(value_token) {
+            s.value()?.into()
+        } else {
+            return None;
+        };
+
         Some((key, value))
     }