7 files changed, 275 insertions, 79 deletions

diff --git a/crates/syntax/src/algo.rs b/crates/syntax/src/algo.rs
index 4f9a7a6e8..9dc7182bd 100644
--- a/crates/syntax/src/algo.rs
+++ b/crates/syntax/src/algo.rs
@@ -137,7 +137,7 @@ impl TreeDiff {
     }
 }
 
-/// Finds minimal the diff, which, applied to `from`, will result in `to`.
+/// Finds a (potentially minimal) diff, which, applied to `from`, will result in `to`.
 ///
 /// Specifically, returns a structure that consists of a replacements, insertions and deletions
 /// such that applying this map on `from` will result in `to`.
@@ -151,7 +151,6 @@ pub fn diff(from: &SyntaxNode, to: &SyntaxNode) -> TreeDiff {
     };
     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);
     }
@@ -169,6 +168,7 @@ pub fn diff(from: &SyntaxNode, to: &SyntaxNode) -> TreeDiff {
             }
     }
 
+    // FIXME: this is horrible inefficient. I bet there's a cool algorithm to diff trees properly.
     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),
@@ -179,6 +179,8 @@ pub fn diff(from: &SyntaxNode, to: &SyntaxNode) -> TreeDiff {
             }
         };
 
+        let mut look_ahead_scratch = Vec::default();
+
         let mut rhs_children = rhs.children_with_tokens();
         let mut lhs_children = lhs.children_with_tokens();
         let mut last_lhs = None;
@@ -204,7 +206,31 @@ pub fn diff(from: &SyntaxNode, to: &SyntaxNode) -> TreeDiff {
                     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),
+                (Some(lhs_ele), Some(rhs_ele)) => {
+                    // nodes differ, look for lhs_ele in rhs, if its found we can mark everything up
+                    // until that element as insertions. This is important to keep the diff minimal
+                    // in regards to insertions that have been actually done, this is important for
+                    // use insertions as we do not want to replace the entire module node.
+                    look_ahead_scratch.push(rhs_ele.clone());
+                    let mut rhs_children_clone = rhs_children.clone();
+                    let mut insert = false;
+                    while let Some(rhs_child) = rhs_children_clone.next() {
+                        if syntax_element_eq(&lhs_ele, &rhs_child) {
+                            mark::hit!(diff_insertions);
+                            insert = true;
+                            break;
+                        } else {
+                            look_ahead_scratch.push(rhs_child);
+                        }
+                    }
+                    let drain = look_ahead_scratch.drain(..);
+                    if let Some(prev) = last_lhs.clone().filter(|_| insert) {
+                        diff.insertions.entry(prev).or_insert_with(Vec::new).extend(drain);
+                        rhs_children = rhs_children_clone;
+                    } else {
+                        go(diff, lhs_ele, rhs_ele)
+                    }
+                }
             }
             last_lhs = lhs_child.or(last_lhs);
         }
@@ -289,11 +315,18 @@ fn _replace_children(
     with_children(parent, new_children)
 }
 
+#[derive(Debug, PartialEq, Eq, Hash)]
+enum InsertPos {
+    FirstChildOf(SyntaxNode),
+    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<'_> {
@@ -304,13 +337,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());
@@ -330,7 +446,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)
@@ -346,14 +462,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())
@@ -367,9 +491,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);
         }
@@ -383,34 +514,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()));
+                }
+            }
         }
-        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)
+        if let Some(elements) = self.insertions(&InsertPos::After(element.clone())) {
+            acc.extend(elements.map(element_to_green));
+        }
+    }
+}
+
+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)),
+            }
+        }
     }
 }
 
@@ -486,18 +628,44 @@ mod tests {
     }
 
     #[test]
-    fn insert() {
+    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 insert_last() {
         mark::check!(diff_insert);
         check_diff(
-            r#"use foo;"#,
-            r#"use foo;
+            r#"
+use foo;
 use bar;"#,
+            r#"
+use foo;
+use bar;
+use baz;"#,
             expect![[r#"
                 insertions:
 
-                Line 0: Node([email protected])
+                Line 2: Node(USE@10..18)
                 -> "\n"
-                -> use bar;
+                -> use baz;
 
                 replacements:
 
@@ -511,29 +679,63 @@ use bar;"#,
     }
 
     #[test]
-    fn replace_parent() {
-        mark::check!(diff_replace_parent);
+    fn insert_middle() {
         check_diff(
-            r#""#,
-            r#"use foo::bar;"#,
+            r#"
+use foo;
+use baz;"#,
+            r#"
+use foo;
+use bar;
+use baz;"#,
             expect![[r#"
                 insertions:
 
+                Line 2: Token([email protected] "\n")
+                -> use bar;
+                -> "\n"
+
+                replacements:
+
+
+
+                deletions:
+
 
+            "#]],
+        )
+    }
+
+    #[test]
+    fn insert_first() {
+        check_diff(
+            r#"
+use bar;
+use baz;"#,
+            r#"
+use foo;
+use bar;
+use baz;"#,
+            expect![[r#"
+                insertions:
+
+                Line 0: Token([email protected] "\n")
+                -> use foo;
+                -> "\n"
 
                 replacements:
 
-                Line 0: Node([email protected]) -> use foo::bar;
+
 
                 deletions:
 
 
             "#]],
-        );
+        )
     }
 
     #[test]
-    fn delete() {
+    fn delete_last() {
         mark::check!(diff_delete);
         check_diff(
             r#"use foo;
@@ -557,52 +759,50 @@ use bar;"#,
     }
 
     #[test]
-    fn insert_use() {
+    fn delete_middle() {
+        mark::check!(diff_insertions);
         check_diff(
             r#"
 use expect_test::{expect, Expect};
+use text_edit::TextEdit;
 
 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")
+                Line 1: Node([email protected])
+                -> "\n\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:
 
+                deletions:
 
+                Line 2: use text_edit::TextEdit;
+                Line 3: "\n\n"
+                Line 4: use crate::AstNode;
+                Line 5: "\n"
             "#]],
         )
     }
 
     #[test]
-    fn remove_use() {
+    fn delete_first() {
         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#"
@@ -612,15 +812,14 @@ use crate::AstNode;
 
                 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"
+                Line 2: Node([email protected]) -> crate
+                Line 2: Token([email protected] "TextEdit") -> AstNode
+                Line 2: Token([email protected] "\n\n") -> "\n"
 
                 deletions:
 
-                Line 4: use crate::AstNode;
-                Line 5: "\n"
+                Line 3: use crate::AstNode;
+                Line 4: "\n"
             "#]],
         )
     }
@@ -697,17 +896,15 @@ fn main() {
                         _ => return,
                     }
                 -> ;
-                Line 5: Token(R_CURLY@64..65 "}")
-                -> "\n"
-                -> }
+                Line 3: Node(IF_EXPR@17..63)
+                -> "\n    "
+                -> foo(x);
 
                 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:
 
diff --git a/crates/syntax/src/ast.rs b/crates/syntax/src/ast.rs
index d536bb1e7..8a0e3d27b 100644
--- a/crates/syntax/src/ast.rs
+++ b/crates/syntax/src/ast.rs
@@ -17,7 +17,7 @@ use crate::{
 
 pub use self::{
     expr_ext::{ArrayExprKind, BinOp, Effect, ElseBranch, LiteralKind, PrefixOp, RangeOp},
-    generated::*,
+    generated::{nodes::*, tokens::*},
     node_ext::{
         AttrKind, FieldKind, NameOrNameRef, PathSegmentKind, SelfParamKind, SlicePatComponents,
         StructKind, TypeBoundKind, VisibilityKind,
diff --git a/crates/syntax/src/ast/generated.rs b/crates/syntax/src/ast/generated.rs
index 4a6f41ee7..843b43cf0 100644
--- a/crates/syntax/src/ast/generated.rs
+++ b/crates/syntax/src/ast/generated.rs
@@ -1,8 +1,8 @@
 //! This file is actually hand-written, but the submodules are indeed generated.
 #[rustfmt::skip]
-mod nodes;
+pub(crate) mod nodes;
 #[rustfmt::skip]
-mod tokens;
+pub(crate) mod tokens;
 
 use crate::{
     AstNode,
@@ -10,7 +10,7 @@ use crate::{
     SyntaxNode,
 };
 
-pub use {nodes::*, tokens::*};
+pub(crate) use nodes::*;
 
 // Stmt is the only nested enum, so it's easier to just hand-write it
 impl AstNode for Stmt {
diff --git a/crates/syntax/src/lib.rs b/crates/syntax/src/lib.rs
index 849a1cdd6..e753b11bb 100644
--- a/crates/syntax/src/lib.rs
+++ b/crates/syntax/src/lib.rs
@@ -46,16 +46,19 @@ use text_edit::Indel;
 pub use crate::{
     algo::InsertPosition,
     ast::{AstNode, AstToken},
-    parsing::{lex_single_syntax_kind, lex_single_valid_syntax_kind, tokenize, Token},
+    parsing::lexer::{lex_single_syntax_kind, lex_single_valid_syntax_kind, tokenize, Token},
     ptr::{AstPtr, SyntaxNodePtr},
     syntax_error::SyntaxError,
     syntax_node::{
-        Direction, GreenNode, NodeOrToken, SyntaxElement, SyntaxElementChildren, SyntaxNode,
-        SyntaxNodeChildren, SyntaxToken, SyntaxTreeBuilder,
+        SyntaxElement, SyntaxElementChildren, SyntaxNode, SyntaxNodeChildren, SyntaxToken,
+        SyntaxTreeBuilder,
     },
 };
 pub use parser::{SyntaxKind, T};
-pub use rowan::{SmolStr, SyntaxText, TextRange, TextSize, TokenAtOffset, WalkEvent};
+pub use rowan::{
+    Direction, GreenNode, NodeOrToken, SmolStr, SyntaxText, TextRange, TextSize, TokenAtOffset,
+    WalkEvent,
+};
 
 /// `Parse` is the result of the parsing: a syntax tree and a collection of
 /// errors.
diff --git a/crates/syntax/src/parsing.rs b/crates/syntax/src/parsing.rs
index 68a39eb21..333bde54a 100644
--- a/crates/syntax/src/parsing.rs
+++ b/crates/syntax/src/parsing.rs
@@ -1,7 +1,7 @@
 //! Lexing, bridging to parser (which does the actual parsing) and
 //! incremental reparsing.
 
-mod lexer;
+pub(crate) mod lexer;
 mod text_token_source;
 mod text_tree_sink;
 mod reparsing;
@@ -10,7 +10,7 @@ use crate::{syntax_node::GreenNode, AstNode, SyntaxError, SyntaxNode};
 use text_token_source::TextTokenSource;
 use text_tree_sink::TextTreeSink;
 
-pub use lexer::*;
+pub(crate) use lexer::*;
 
 pub(crate) use self::reparsing::incremental_reparse;
 use parser::SyntaxKind;
diff --git a/crates/syntax/src/parsing/text_token_source.rs b/crates/syntax/src/parsing/text_token_source.rs
index df866dc2b..0614194a5 100644
--- a/crates/syntax/src/parsing/text_token_source.rs
+++ b/crates/syntax/src/parsing/text_token_source.rs
@@ -65,7 +65,7 @@ fn mk_token(pos: usize, token_offset_pairs: &[(Token, TextSize)]) -> parser::Tok
 
 impl<'t> TextTokenSource<'t> {
     /// Generate input from tokens(expect comment and whitespace).
-    pub fn new(text: &'t str, raw_tokens: &'t [Token]) -> TextTokenSource<'t> {
+    pub(crate) fn new(text: &'t str, raw_tokens: &'t [Token]) -> TextTokenSource<'t> {
         let token_offset_pairs: Vec<_> = raw_tokens
             .iter()
             .filter_map({
diff --git a/crates/syntax/src/syntax_node.rs b/crates/syntax/src/syntax_node.rs
index b2abcbfbb..cc30138fa 100644
--- a/crates/syntax/src/syntax_node.rs
+++ b/crates/syntax/src/syntax_node.rs
@@ -10,9 +10,7 @@ use rowan::{GreenNodeBuilder, Language};
 
 use crate::{Parse, SmolStr, SyntaxError, SyntaxKind, TextSize};
 
-pub use rowan::GreenNode;
-
-pub(crate) use rowan::GreenToken;
+pub(crate) use rowan::{GreenNode, GreenToken, NodeOrToken};
 
 #[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
 pub enum RustLanguage {}
@@ -34,8 +32,6 @@ pub type SyntaxElement = rowan::SyntaxElement<RustLanguage>;
 pub type SyntaxNodeChildren = rowan::SyntaxNodeChildren<RustLanguage>;
 pub type SyntaxElementChildren = rowan::SyntaxElementChildren<RustLanguage>;
 
-pub use rowan::{Direction, NodeOrToken};
-
 #[derive(Default)]
 pub struct SyntaxTreeBuilder {
     errors: Vec<SyntaxError>,