move add_missing_members to structured editing API

Currently, this is more code, and we also loose auto-indenting of
bodies, but, long-term, this is the right approach

1 files changed, 140 insertions, 18 deletions

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 {