4 files changed, 231 insertions, 146 deletions

diff --git a/crates/ra_assists/src/add_missing_impl_members.rs b/crates/ra_assists/src/add_missing_impl_members.rs
index c82447b84..17c2af899 100644
--- a/crates/ra_assists/src/add_missing_impl_members.rs
+++ b/crates/ra_assists/src/add_missing_impl_members.rs
@@ -1,14 +1,9 @@
-use std::fmt::Write;
-
-use crate::{Assist, AssistId, AssistCtx};
+use crate::{Assist, AssistId, AssistCtx, ast_editor::{AstEditor, AstBuilder}};
 
 use hir::db::HirDatabase;
-use ra_syntax::{SmolStr, SyntaxKind, TextRange, TextUnit, TreeArc};
-use ra_syntax::ast::{self, AstNode, AstToken, FnDef, ImplItem, ImplItemKind, NameOwner};
+use ra_syntax::{SmolStr, TreeArc};
+use ra_syntax::ast::{self, AstNode, FnDef, ImplItem, ImplItemKind, NameOwner};
 use ra_db::FilePosition;
-use ra_fmt::{leading_indent, reindent};
-
-use itertools::Itertools;
 
 enum AddMissingImplMembersMode {
     DefaultMethodsOnly,
@@ -76,48 +71,35 @@ fn add_missing_impl_members_inner(
     }
 
     ctx.add_action(AssistId(assist_id), label, |edit| {
-        let (parent_indent, indent) = {
-            // FIXME: Find a way to get the indent already used in the file.
-            // Now, we copy the indent of first item or indent with 4 spaces relative to impl block
-            const DEFAULT_INDENT: &str = "    ";
-            let first_item = impl_item_list.impl_items().next();
-            let first_item_indent =
-                first_item.and_then(|i| leading_indent(i.syntax())).map(ToOwned::to_owned);
-            let impl_block_indent = leading_indent(impl_node.syntax()).unwrap_or_default();
-
-            (
-                impl_block_indent.to_owned(),
-                first_item_indent.unwrap_or_else(|| impl_block_indent.to_owned() + DEFAULT_INDENT),
-            )
-        };
-
-        let changed_range = {
-            let children = impl_item_list.syntax().children_with_tokens();
-            let last_whitespace =
-                children.filter_map(|it| ast::Whitespace::cast(it.as_token()?)).last();
-
-            last_whitespace.map(|w| w.syntax().range()).unwrap_or_else(|| {
-                let in_brackets = impl_item_list.syntax().range().end() - TextUnit::of_str("}");
-                TextRange::from_to(in_brackets, in_brackets)
-            })
-        };
-
-        let func_bodies = format!("\n{}", missing_fns.into_iter().map(build_func_body).join("\n"));
-        let trailing_whitespace = format!("\n{}", parent_indent);
-        let func_bodies = reindent(&func_bodies, &indent) + &trailing_whitespace;
-
-        let replaced_text_range = TextUnit::of_str(&func_bodies);
-
-        edit.replace(changed_range, func_bodies);
-        // FIXME: place the cursor on the first unimplemented?
-        edit.set_cursor(
-            changed_range.start() + replaced_text_range - TextUnit::of_str(&trailing_whitespace),
-        );
+        let n_existing_items = impl_item_list.impl_items().count();
+        let fns = missing_fns.into_iter().map(add_body_and_strip_docstring).collect::<Vec<_>>();
+
+        let mut ast_editor = AstEditor::new(impl_item_list);
+        if n_existing_items == 0 {
+            ast_editor.make_multiline();
+        }
+        ast_editor.append_functions(fns.iter().map(|it| &**it));
+        let first_new_item = ast_editor.ast().impl_items().nth(n_existing_items).unwrap();
+        let cursor_poisition = first_new_item.syntax().range().start();
+        ast_editor.into_text_edit(edit.text_edit_builder());
+
+        edit.set_cursor(cursor_poisition);
     });
 
     ctx.build()
 }
 
+fn add_body_and_strip_docstring(fn_def: &ast::FnDef) -> TreeArc<ast::FnDef> {
+    let mut ast_editor = AstEditor::new(fn_def);
+    if fn_def.body().is_none() {
+        ast_editor.set_body(&AstBuilder::<ast::Block>::single_expr(
+            &AstBuilder::<ast::Expr>::unimplemented(),
+        ));
+    }
+    ast_editor.strip_attrs_and_docs();
+    ast_editor.ast().to_owned()
+}
+
 /// Given an `ast::ImplBlock`, resolves the target trait (the one being
 /// implemented) to a `ast::TraitDef`.
 fn resolve_target_trait_def(
@@ -134,22 +116,6 @@ fn resolve_target_trait_def(
     }
 }
 
-fn build_func_body(def: &ast::FnDef) -> String {
-    let mut buf = String::new();
-
-    for child in def.syntax().children_with_tokens() {
-        match (child.prev_sibling_or_token().map(|c| c.kind()), child.kind()) {
-            (_, SyntaxKind::SEMI) => buf.push_str(" {\n    unimplemented!()\n}"),
-            (_, SyntaxKind::ATTR) | (_, SyntaxKind::COMMENT) => {}
-            (Some(SyntaxKind::ATTR), SyntaxKind::WHITESPACE)
-            | (Some(SyntaxKind::COMMENT), SyntaxKind::WHITESPACE) => {}
-            _ => write!(buf, "{}", child).unwrap(),
-        };
-    }
-
-    buf.trim_end().to_string()
-}
-
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -170,7 +136,7 @@ struct S;
 
 impl Foo for S {
     fn bar(&self) {}
-    <|>
+<|>
 }",
             "
 trait Foo {
@@ -183,12 +149,9 @@ struct S;
 
 impl Foo for S {
     fn bar(&self) {}
-    fn foo(&self) {
-        unimplemented!()
-    }
-    fn baz(&self) {
-        unimplemented!()
-    }<|>
+    <|>fn foo(&self) { unimplemented!() }
+    fn baz(&self) { unimplemented!() }
+
 }",
         );
     }
@@ -208,7 +171,7 @@ struct S;
 
 impl Foo for S {
     fn bar(&self) {}
-    <|>
+<|>
 }",
             "
 trait Foo {
@@ -221,9 +184,8 @@ struct S;
 
 impl Foo for S {
     fn bar(&self) {}
-    fn foo(&self) {
-        unimplemented!()
-    }<|>
+    <|>fn foo(&self) { unimplemented!() }
+
 }",
         );
     }
@@ -240,9 +202,7 @@ impl Foo for S { <|> }",
 trait Foo { fn foo(&self); }
 struct S;
 impl Foo for S {
-    fn foo(&self) {
-        unimplemented!()
-    }<|>
+    <|>fn foo(&self) { unimplemented!() }
 }",
         );
     }
@@ -259,9 +219,7 @@ impl Foo for S {}<|>",
 trait Foo { fn foo(&self); }
 struct S;
 impl Foo for S {
-    fn foo(&self) {
-        unimplemented!()
-    }<|>
+    <|>fn foo(&self) { unimplemented!() }
 }",
         )
     }
@@ -292,35 +250,6 @@ impl Foo for S { <|> }",
     }
 
     #[test]
-    fn test_indented_impl_block() {
-        check_assist(
-            add_missing_impl_members,
-            "
-trait Foo {
-    fn valid(some: u32) -> bool;
-}
-struct S;
-
-mod my_mod {
-    impl crate::Foo for S { <|> }
-}",
-            "
-trait Foo {
-    fn valid(some: u32) -> bool;
-}
-struct S;
-
-mod my_mod {
-    impl crate::Foo for S {
-        fn valid(some: u32) -> bool {
-            unimplemented!()
-        }<|>
-    }
-}",
-        )
-    }
-
-    #[test]
     fn test_with_docstring_and_attrs() {
         check_assist(
             add_missing_impl_members,
@@ -342,9 +271,7 @@ trait Foo {
 }
 struct S;
 impl Foo for S {
-    fn foo(&self) {
-        unimplemented!()
-    }<|>
+    <|>fn foo(&self) { unimplemented!() }
 }"#,
         )
     }
@@ -367,7 +294,7 @@ trait Foo {
 }
 struct S;
 impl Foo for S {
-    fn valid(some: u32) -> bool { false }<|>
+    <|>fn valid(some: u32) -> bool { false }
 }",
         )
     }
diff --git a/crates/ra_assists/src/ast_editor.rs b/crates/ra_assists/src/ast_editor.rs
index 13ee82879..283b280b6 100644
--- a/crates/ra_assists/src/ast_editor.rs
+++ b/crates/ra_assists/src/ast_editor.rs
@@ -1,4 +1,4 @@
-use std::iter;
+use std::{iter, ops::RangeInclusive};
 
 use arrayvec::ArrayVec;
 use ra_text_edit::TextEditBuilder;
@@ -26,6 +26,7 @@ impl<N: AstNode> AstEditor<N> {
         &*self.ast
     }
 
+    #[must_use]
     fn insert_children<'a>(
         &self,
         position: InsertPosition<SyntaxElement<'_>>,
@@ -34,31 +35,55 @@ impl<N: AstNode> AstEditor<N> {
         let new_syntax = self.ast().syntax().insert_children(position, to_insert);
         N::cast(&new_syntax).unwrap().to_owned()
     }
-}
 
-impl AstEditor<ast::NamedFieldList> {
-    pub fn append_field(&mut self, field: &ast::NamedField) {
-        self.insert_field(InsertPosition::Last, field)
+    #[must_use]
+    fn replace_children<'a>(
+        &self,
+        to_delete: RangeInclusive<SyntaxElement<'_>>,
+        to_insert: impl Iterator<Item = SyntaxElement<'a>>,
+    ) -> TreeArc<N> {
+        let new_syntax = self.ast().syntax().replace_children(to_delete, to_insert);
+        N::cast(&new_syntax).unwrap().to_owned()
     }
 
-    pub fn make_multiline(&mut self) {
-        let l_curly = match self.l_curly() {
-            Some(it) => it,
-            None => return,
-        };
+    fn do_make_multiline(&mut self) {
+        let l_curly =
+            match self.ast().syntax().children_with_tokens().find(|it| it.kind() == L_CURLY) {
+                Some(it) => it,
+                None => return,
+            };
         let sibling = match l_curly.next_sibling_or_token() {
             Some(it) => it,
             None => return,
         };
-        if sibling.as_token().map(|it| it.text().contains('\n')) == Some(true) {
-            return;
-        }
+        let existing_ws = match sibling.as_token() {
+            None => None,
+            Some(tok) if tok.kind() != WHITESPACE => None,
+            Some(ws) => {
+                if ws.text().contains('\n') {
+                    return;
+                }
+                Some(ws)
+            }
+        };
+
+        let indent = leading_indent(self.ast().syntax()).unwrap_or("");
+        let ws = tokens::WsBuilder::new(&format!("\n{}", indent));
+        let to_insert = iter::once(ws.ws().into());
+        self.ast = match existing_ws {
+            None => self.insert_children(InsertPosition::After(l_curly), to_insert),
+            Some(ws) => self.replace_children(RangeInclusive::new(ws.into(), ws.into()), to_insert),
+        };
+    }
+}
+
+impl AstEditor<ast::NamedFieldList> {
+    pub fn append_field(&mut self, field: &ast::NamedField) {
+        self.insert_field(InsertPosition::Last, field)
+    }
 
-        let ws = tokens::WsBuilder::new(&format!(
-            "\n{}",
-            leading_indent(self.ast().syntax()).unwrap_or("")
-        ));
-        self.ast = self.insert_children(InsertPosition::After(l_curly), iter::once(ws.ws().into()));
+    pub fn make_multiline(&mut self) {
+        self.do_make_multiline()
     }
 
     pub fn insert_field(
@@ -132,6 +157,79 @@ impl AstEditor<ast::NamedFieldList> {
     }
 }
 
+impl AstEditor<ast::ItemList> {
+    pub fn make_multiline(&mut self) {
+        self.do_make_multiline()
+    }
+
+    pub fn append_functions<'a>(&mut self, fns: impl Iterator<Item = &'a ast::FnDef>) {
+        fns.for_each(|it| self.append_function(it))
+    }
+
+    pub fn append_function(&mut self, fn_def: &ast::FnDef) {
+        let (indent, position) = match self.ast().impl_items().last() {
+            Some(it) => (
+                leading_indent(it.syntax()).unwrap_or("").to_string(),
+                InsertPosition::After(it.syntax().into()),
+            ),
+            None => match self.l_curly() {
+                Some(it) => (
+                    "    ".to_string() + leading_indent(self.ast().syntax()).unwrap_or(""),
+                    InsertPosition::After(it),
+                ),
+                None => return,
+            },
+        };
+        let ws = tokens::WsBuilder::new(&format!("\n{}", indent));
+        let to_insert: ArrayVec<[SyntaxElement; 2]> =
+            [ws.ws().into(), fn_def.syntax().into()].into();
+        self.ast = self.insert_children(position, to_insert.into_iter());
+    }
+
+    fn l_curly(&self) -> Option<SyntaxElement> {
+        self.ast().syntax().children_with_tokens().find(|it| it.kind() == L_CURLY)
+    }
+}
+
+impl AstEditor<ast::FnDef> {
+    pub fn set_body(&mut self, body: &ast::Block) {
+        let mut to_insert: ArrayVec<[SyntaxElement; 2]> = ArrayVec::new();
+        let old_body_or_semi: SyntaxElement = if let Some(old_body) = self.ast().body() {
+            old_body.syntax().into()
+        } else if let Some(semi) = self.ast().semicolon_token() {
+            to_insert.push(tokens::single_space().into());
+            semi.into()
+        } else {
+            to_insert.push(tokens::single_space().into());
+            to_insert.push(body.syntax().into());
+            self.ast = self.insert_children(InsertPosition::Last, to_insert.into_iter());
+            return;
+        };
+        to_insert.push(body.syntax().into());
+        let replace_range = RangeInclusive::new(old_body_or_semi, old_body_or_semi);
+        self.ast = self.replace_children(replace_range, to_insert.into_iter())
+    }
+
+    pub fn strip_attrs_and_docs(&mut self) {
+        loop {
+            if let Some(start) = self
+                .ast()
+                .syntax()
+                .children_with_tokens()
+                .find(|it| it.kind() == ATTR || it.kind() == COMMENT)
+            {
+                let end = match start.next_sibling_or_token() {
+                    Some(el) if el.kind() == WHITESPACE => el,
+                    Some(_) | None => start,
+                };
+                self.ast = self.replace_children(RangeInclusive::new(start, end), iter::empty());
+            } else {
+                break;
+            }
+        }
+    }
+}
+
 pub struct AstBuilder<N: AstNode> {
     _phantom: std::marker::PhantomData<N>,
 }
@@ -149,6 +247,16 @@ impl AstBuilder<ast::NamedField> {
     }
 }
 
+impl AstBuilder<ast::Block> {
+    fn from_text(text: &str) -> TreeArc<ast::Block> {
+        ast_node_from_file_text(&format!("fn f() {}", text))
+    }
+
+    pub fn single_expr(e: &ast::Expr) -> TreeArc<ast::Block> {
+        Self::from_text(&format!("{{ {} }}", e.syntax()))
+    }
+}
+
 impl AstBuilder<ast::Expr> {
     fn from_text(text: &str) -> TreeArc<ast::Expr> {
         ast_node_from_file_text(&format!("fn f() {{ {}; }}", text))
@@ -157,6 +265,10 @@ impl AstBuilder<ast::Expr> {
     pub fn unit() -> TreeArc<ast::Expr> {
         Self::from_text("()")
     }
+
+    pub fn unimplemented() -> TreeArc<ast::Expr> {
+        Self::from_text("unimplemented!()")
+    }
 }
 
 impl AstBuilder<ast::NameRef> {
@@ -197,6 +309,16 @@ mod tokens {
             .unwrap()
     }
 
+    #[allow(unused)]
+    pub(crate) fn single_newline() -> SyntaxToken<'static> {
+        SOURCE_FILE
+            .syntax()
+            .descendants_with_tokens()
+            .filter_map(|it| it.as_token())
+            .find(|it| it.kind() == WHITESPACE && it.text().as_str() == "\n")
+            .unwrap()
+    }
+
     pub(crate) struct WsBuilder(TreeArc<SourceFile>);
 
     impl WsBuilder {
diff --git a/crates/ra_syntax/src/ast/extensions.rs b/crates/ra_syntax/src/ast/extensions.rs
index 5c4c0ffc1..9cbd2c6b8 100644
--- a/crates/ra_syntax/src/ast/extensions.rs
+++ b/crates/ra_syntax/src/ast/extensions.rs
@@ -210,6 +210,15 @@ impl ast::EnumVariant {
     }
 }
 
+impl ast::FnDef {
+    pub fn semicolon_token(&self) -> Option<SyntaxToken<'_>> {
+        self.syntax()
+            .last_child_or_token()
+            .and_then(|it| it.as_token())
+            .filter(|it| it.kind() == SEMI)
+    }
+}
+
 impl ast::LetStmt {
     pub fn has_semi(&self) -> bool {
         match self.syntax().last_child_or_token() {
diff --git a/crates/ra_syntax/src/syntax_node.rs b/crates/ra_syntax/src/syntax_node.rs
index 628cabc29..92c15234e 100644
--- a/crates/ra_syntax/src/syntax_node.rs
+++ b/crates/ra_syntax/src/syntax_node.rs
@@ -7,6 +7,7 @@
 //! modules just wraps its API.
 
 use std::{
+    ops::RangeInclusive,
     fmt::{self, Write},
     any::Any,
     borrow::Borrow,
@@ -323,8 +324,6 @@ impl SyntaxNode {
     ///
     /// This is a type-unsafe low-level editing API, if you need to use it,
     /// prefer to create a type-safe abstraction on top of it instead.
-    ///
-    ///
     pub fn insert_children<'a>(
         &self,
         position: InsertPosition<SyntaxElement<'_>>,
@@ -338,12 +337,6 @@ impl SyntaxNode {
 
         let old_children = self.0.green().children();
 
-        let get_anchor_pos = |anchor: SyntaxElement| -> usize {
-            self.children_with_tokens()
-                .position(|it| it == anchor)
-                .expect("anchor is not a child of current element")
-        };
-
         let new_children = match position {
             InsertPosition::First => {
                 to_insert.chain(old_children.iter().cloned()).collect::<Box<[_]>>()
@@ -353,7 +346,8 @@ impl SyntaxNode {
             }
             InsertPosition::Before(anchor) | InsertPosition::After(anchor) => {
                 let take_anchor = if let InsertPosition::After(_) = position { 1 } else { 0 };
-                let (before, after) = old_children.split_at(get_anchor_pos(anchor) + take_anchor);
+                let split_at = self.position_of_child(anchor) + take_anchor;
+                let (before, after) = old_children.split_at(split_at);
                 before
                     .iter()
                     .cloned()
@@ -363,6 +357,33 @@ impl SyntaxNode {
             }
         };
 
+        self.with_children(new_children)
+    }
+
+    /// Replaces all nodes in `to_delete` with nodes from `to_insert`
+    ///
+    /// This is a type-unsafe low-level editing API, if you need to use it,
+    /// prefer to create a type-safe abstraction on top of it instead.
+    pub fn replace_children<'a>(
+        &self,
+        to_delete: RangeInclusive<SyntaxElement<'_>>,
+        to_insert: impl Iterator<Item = SyntaxElement<'a>>,
+    ) -> TreeArc<SyntaxNode> {
+        let start = self.position_of_child(*to_delete.start());
+        let end = self.position_of_child(*to_delete.end());
+        let old_children = self.0.green().children();
+
+        let new_children = old_children[..start]
+            .iter()
+            .cloned()
+            .chain(to_insert.map(to_green_element))
+            .chain(old_children[end + 1..].iter().cloned())
+            .collect::<Box<[_]>>();
+        self.with_children(new_children)
+    }
+
+    fn with_children(&self, new_children: Box<[rowan::GreenElement]>) -> TreeArc<SyntaxNode> {
+        let len = new_children.iter().map(|it| it.text_len()).sum::<TextUnit>();
         let new_node = GreenNode::new(rowan::SyntaxKind(self.kind() as u16), new_children);
         let new_file_node = self.replace_with(new_node);
         let file = SourceFile::new(new_file_node, Vec::new());
@@ -370,16 +391,22 @@ impl SyntaxNode {
         // FIXME: use a more elegant way to re-fetch the node (#1185), make
         // `range` private afterwards
         let mut ptr = SyntaxNodePtr::new(self);
-        ptr.range = TextRange::from_to(ptr.range().start(), ptr.range().end() + delta);
+        ptr.range = TextRange::offset_len(ptr.range().start(), len);
         return ptr.to_node(&file).to_owned();
+    }
 
-        fn to_green_element(element: SyntaxElement) -> rowan::GreenElement {
-            match element {
-                SyntaxElement::Node(node) => node.0.green().clone().into(),
-                SyntaxElement::Token(tok) => {
-                    GreenToken::new(rowan::SyntaxKind(tok.kind() as u16), tok.text().clone()).into()
-                }
-            }
+    fn position_of_child(&self, child: SyntaxElement) -> usize {
+        self.children_with_tokens()
+            .position(|it| it == child)
+            .expect("elemetn is not a child of current element")
+    }
+}
+
+fn to_green_element(element: SyntaxElement) -> rowan::GreenElement {
+    match element {
+        SyntaxElement::Node(node) => node.0.green().clone().into(),
+        SyntaxElement::Token(tok) => {
+            GreenToken::new(rowan::SyntaxKind(tok.kind() as u16), tok.text().clone()).into()
         }
     }
 }