4 files changed, 142 insertions, 80 deletions

diff --git a/crates/syntax/src/algo.rs b/crates/syntax/src/algo.rs
index 30adb7217..320c430c9 100644
--- a/crates/syntax/src/algo.rs
+++ b/crates/syntax/src/algo.rs
@@ -111,18 +111,28 @@ pub enum InsertPosition<T> {
 
 type FxIndexMap<K, V> = IndexMap<K, V, BuildHasherDefault<rustc_hash::FxHasher>>;
 
+#[derive(Debug, Hash, PartialEq, Eq)]
+enum TreeDiffInsertPos {
+    After(SyntaxElement),
+    AsFirstChild(SyntaxElement),
+}
+
 #[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>>,
+    insertions: FxIndexMap<TreeDiffInsertPos, 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()));
+            let offset = match anchor {
+                TreeDiffInsertPos::After(it) => it.text_range().end(),
+                TreeDiffInsertPos::AsFirstChild(it) => it.text_range().start(),
+            };
+            to.iter().for_each(|to| builder.insert(offset, to.to_string()));
         }
         for (from, to) in self.replacements.iter() {
             builder.replace(from.text_range(), to.to_string())
@@ -188,19 +198,20 @@ pub fn diff(from: &SyntaxNode, to: &SyntaxNode) -> TreeDiff {
             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;
-                    }
-                },
+                (None, Some(element)) => {
+                    let insert_pos = match last_lhs.clone() {
+                        Some(prev) => {
+                            mark::hit!(diff_insert);
+                            TreeDiffInsertPos::After(prev)
+                        }
+                        // first iteration, insert into out parent as the first child
+                        None => {
+                            mark::hit!(diff_insert_as_first_child);
+                            TreeDiffInsertPos::AsFirstChild(lhs.clone().into())
+                        }
+                    };
+                    diff.insertions.entry(insert_pos).or_insert_with(Vec::new).push(element);
+                }
                 (Some(element), None) => {
                     mark::hit!(diff_delete);
                     diff.deletions.push(element);
@@ -224,8 +235,15 @@ pub fn diff(from: &SyntaxNode, to: &SyntaxNode) -> TreeDiff {
                         }
                     }
                     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);
+                    if insert {
+                        let insert_pos = if let Some(prev) = last_lhs.clone().filter(|_| insert) {
+                            TreeDiffInsertPos::After(prev)
+                        } else {
+                            mark::hit!(insert_first_child);
+                            TreeDiffInsertPos::AsFirstChild(lhs.clone().into())
+                        };
+
+                        diff.insertions.entry(insert_pos).or_insert_with(Vec::new).extend(drain);
                         rhs_children = rhs_children_clone;
                     } else {
                         go(diff, lhs_ele, rhs_ele)
@@ -632,18 +650,19 @@ mod tests {
 
     #[test]
     fn replace_parent() {
-        mark::check!(diff_replace_parent);
+        mark::check!(diff_insert_as_first_child);
         check_diff(
             r#""#,
             r#"use foo::bar;"#,
             expect![[r#"
                 insertions:
 
-
+                Line 0: AsFirstChild(Node([email protected]))
+                -> use foo::bar;
 
                 replacements:
 
-                Line 0: Node([email protected]) -> use foo::bar;
+
 
                 deletions:
 
@@ -666,7 +685,7 @@ use baz;"#,
             expect![[r#"
                 insertions:
 
-                Line 2: Node([email protected])
+                Line 2: After(Node([email protected]))
                 -> "\n"
                 -> use baz;
 
@@ -694,7 +713,7 @@ use baz;"#,
             expect![[r#"
                 insertions:
 
-                Line 2: Token([email protected] "\n")
+                Line 2: After(Token([email protected] "\n"))
                 -> use bar;
                 -> "\n"
 
@@ -722,7 +741,7 @@ use baz;"#,
             expect![[r#"
                 insertions:
 
-                Line 0: Token([email protected] "\n")
+                Line 0: After(Token([email protected] "\n"))
                 -> use foo;
                 -> "\n"
 
@@ -738,6 +757,36 @@ use baz;"#,
     }
 
     #[test]
+    fn first_child_insertion() {
+        mark::check!(insert_first_child);
+        check_diff(
+            r#"fn main() {
+        stdi
+    }"#,
+            r#"use foo::bar;
+
+    fn main() {
+        stdi
+    }"#,
+            expect![[r#"
+                insertions:
+
+                Line 0: AsFirstChild(Node([email protected]))
+                -> use foo::bar;
+                -> "\n\n    "
+
+                replacements:
+
+
+
+                deletions:
+
+
+            "#]],
+        );
+    }
+
+    #[test]
     fn delete_last() {
         mark::check!(diff_delete);
         check_diff(
@@ -779,7 +828,7 @@ use crate::AstNode;
             expect![[r#"
                 insertions:
 
-                Line 1: Node([email protected])
+                Line 1: After(Node([email protected]))
                 -> "\n\n"
                 -> use crate::AstNode;
 
@@ -845,10 +894,10 @@ use std::ops::{self, RangeInclusive};
             expect![[r#"
                 insertions:
 
-                Line 2: Node([email protected])
+                Line 2: After(Node([email protected]))
                 -> ::
                 -> fmt
-                Line 6: Token([email protected] "\n")
+                Line 6: After(Token([email protected] "\n"))
                 -> use std::hash::BuildHasherDefault;
                 -> "\n"
                 -> use std::ops::{self, RangeInclusive};
@@ -892,14 +941,14 @@ fn main() {
             expect![[r#"
                 insertions:
 
-                Line 3: Node([email protected])
+                Line 3: After(Node([email protected]))
                 -> " "
                 -> match Err(92) {
                         Ok(it) => it,
                         _ => return,
                     }
                 -> ;
-                Line 3: Node([email protected])
+                Line 3: After(Node([email protected]))
                 -> "\n    "
                 -> foo(x);
 
@@ -934,14 +983,18 @@ fn main() {
             _ => 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 insertions =
+            diff.insertions.iter().format_with("\n", |(k, v), f| -> Result<(), std::fmt::Error> {
+                f(&format!(
+                    "Line {}: {:?}\n-> {}",
+                    line_number(match k {
+                        super::TreeDiffInsertPos::After(syn) => syn,
+                        super::TreeDiffInsertPos::AsFirstChild(syn) => syn,
+                    }),
+                    k,
+                    v.iter().format_with("\n-> ", |v, f| f(&fmt_syntax(v)))
+                ))
+            });
 
         let replacements = diff
             .replacements
diff --git a/crates/syntax/src/ast.rs b/crates/syntax/src/ast.rs
index 8a0e3d27b..7844f9ed6 100644
--- a/crates/syntax/src/ast.rs
+++ b/crates/syntax/src/ast.rs
@@ -115,10 +115,10 @@ fn test_doc_comment_none() {
 }
 
 #[test]
-fn test_doc_comment_of_items() {
+fn test_outer_doc_comment_of_items() {
     let file = SourceFile::parse(
         r#"
-        //! doc
+        /// doc
         // non-doc
         mod foo {}
         "#,
@@ -130,6 +130,21 @@ fn test_doc_comment_of_items() {
 }
 
 #[test]
+fn test_inner_doc_comment_of_items() {
+    let file = SourceFile::parse(
+        r#"
+        //! doc
+        // non-doc
+        mod foo {}
+        "#,
+    )
+    .ok()
+    .unwrap();
+    let module = file.syntax().descendants().find_map(Module::cast).unwrap();
+    assert!(module.doc_comment_text().is_none());
+}
+
+#[test]
 fn test_doc_comment_of_statics() {
     let file = SourceFile::parse(
         r#"
diff --git a/crates/syntax/src/ast/token_ext.rs b/crates/syntax/src/ast/token_ext.rs
index e4e512f2e..2661c753e 100644
--- a/crates/syntax/src/ast/token_ext.rs
+++ b/crates/syntax/src/ast/token_ext.rs
@@ -14,16 +14,15 @@ use crate::{
 
 impl ast::Comment {
     pub fn kind(&self) -> CommentKind {
-        kind_by_prefix(self.text())
+        CommentKind::from_text(self.text())
     }
 
     pub fn prefix(&self) -> &'static str {
-        for (prefix, k) in COMMENT_PREFIX_TO_KIND.iter() {
-            if *k == self.kind() && self.text().starts_with(prefix) {
-                return prefix;
-            }
-        }
-        unreachable!()
+        let &(prefix, _kind) = CommentKind::BY_PREFIX
+            .iter()
+            .find(|&(prefix, kind)| self.kind() == *kind && self.text().starts_with(prefix))
+            .unwrap();
+        prefix
     }
 }
 
@@ -55,29 +54,25 @@ pub enum CommentPlacement {
     Outer,
 }
 
-const COMMENT_PREFIX_TO_KIND: &[(&str, CommentKind)] = {
-    use {CommentPlacement::*, CommentShape::*};
-    &[
-        ("////", CommentKind { shape: Line, doc: None }),
-        ("///", CommentKind { shape: Line, doc: Some(Outer) }),
-        ("//!", CommentKind { shape: Line, doc: Some(Inner) }),
-        ("/**", CommentKind { shape: Block, doc: Some(Outer) }),
-        ("/*!", CommentKind { shape: Block, doc: Some(Inner) }),
-        ("//", CommentKind { shape: Line, doc: None }),
-        ("/*", CommentKind { shape: Block, doc: None }),
-    ]
-};
+impl CommentKind {
+    const BY_PREFIX: [(&'static str, CommentKind); 8] = [
+        ("/**/", CommentKind { shape: CommentShape::Block, doc: None }),
+        ("////", CommentKind { shape: CommentShape::Line, doc: None }),
+        ("///", CommentKind { shape: CommentShape::Line, doc: Some(CommentPlacement::Outer) }),
+        ("//!", CommentKind { shape: CommentShape::Line, doc: Some(CommentPlacement::Inner) }),
+        ("/**", CommentKind { shape: CommentShape::Block, doc: Some(CommentPlacement::Outer) }),
+        ("/*!", CommentKind { shape: CommentShape::Block, doc: Some(CommentPlacement::Inner) }),
+        ("//", CommentKind { shape: CommentShape::Line, doc: None }),
+        ("/*", CommentKind { shape: CommentShape::Block, doc: None }),
+    ];
 
-fn kind_by_prefix(text: &str) -> CommentKind {
-    if text == "/**/" {
-        return CommentKind { shape: CommentShape::Block, doc: None };
-    }
-    for (prefix, kind) in COMMENT_PREFIX_TO_KIND.iter() {
-        if text.starts_with(prefix) {
-            return *kind;
-        }
+    pub(crate) fn from_text(text: &str) -> CommentKind {
+        let &(_prefix, kind) = CommentKind::BY_PREFIX
+            .iter()
+            .find(|&(prefix, _kind)| text.starts_with(prefix))
+            .unwrap();
+        kind
     }
-    panic!("bad comment text: {:?}", text)
 }
 
 impl ast::Whitespace {
diff --git a/crates/syntax/src/parsing/text_tree_sink.rs b/crates/syntax/src/parsing/text_tree_sink.rs
index c1b5f246d..997bc5d28 100644
--- a/crates/syntax/src/parsing/text_tree_sink.rs
+++ b/crates/syntax/src/parsing/text_tree_sink.rs
@@ -5,6 +5,7 @@ use std::mem;
 use parser::{ParseError, TreeSink};
 
 use crate::{
+    ast,
     parsing::Token,
     syntax_node::GreenNode,
     SmolStr, SyntaxError,
@@ -153,24 +154,22 @@ fn n_attached_trivias<'a>(
 
             while let Some((i, (kind, text))) = trivias.next() {
                 match kind {
-                    WHITESPACE => {
-                        if text.contains("\n\n") {
-                            // we check whether the next token is a doc-comment
-                            // and skip the whitespace in this case
-                            if let Some((peek_kind, peek_text)) =
-                                trivias.peek().map(|(_, pair)| pair)
-                            {
-                                if *peek_kind == COMMENT
-                                    && peek_text.starts_with("///")
-                                    && !peek_text.starts_with("////")
-                                {
-                                    continue;
-                                }
+                    WHITESPACE if text.contains("\n\n") => {
+                        // we check whether the next token is a doc-comment
+                        // and skip the whitespace in this case
+                        if let Some((COMMENT, peek_text)) = trivias.peek().map(|(_, pair)| pair) {
+                            let comment_kind = ast::CommentKind::from_text(peek_text);
+                            if comment_kind.doc == Some(ast::CommentPlacement::Outer) {
+                                continue;
                             }
-                            break;
                         }
+                        break;
                     }
                     COMMENT => {
+                        let comment_kind = ast::CommentKind::from_text(text);
+                        if comment_kind.doc == Some(ast::CommentPlacement::Inner) {
+                            break;
+                        }
                         res = i + 1;
                     }
                     _ => (),