4 files changed, 103 insertions, 0 deletions

diff --git a/crates/ra_syntax/Cargo.toml b/crates/ra_syntax/Cargo.toml
index 724c38e17..9bc85404a 100644
--- a/crates/ra_syntax/Cargo.toml
+++ b/crates/ra_syntax/Cargo.toml
@@ -12,6 +12,8 @@ itertools = "0.8.0"
 rowan = "0.6.1"
 rustc_lexer = "0.1.0"
 rustc-hash = "1.0.1"
+arrayvec = "0.4.10"
+once_cell = "1.2.0"
 
 # ideally, `serde` should be enabled by `ra_lsp_server`, but we enable it here
 # to reduce number of compilations
diff --git a/crates/ra_syntax/src/ast.rs b/crates/ra_syntax/src/ast.rs
index f464d6534..fdffd8cb1 100644
--- a/crates/ra_syntax/src/ast.rs
+++ b/crates/ra_syntax/src/ast.rs
@@ -5,6 +5,7 @@ mod traits;
 mod tokens;
 mod extensions;
 mod expr_extensions;
+mod edit;
 pub mod make;
 
 use std::marker::PhantomData;
diff --git a/crates/ra_syntax/src/ast/edit.rs b/crates/ra_syntax/src/ast/edit.rs
new file mode 100644
index 000000000..c65899812
--- /dev/null
+++ b/crates/ra_syntax/src/ast/edit.rs
@@ -0,0 +1,52 @@
+//! This module contains functions for editing syntax trees. As the trees are
+//! immutable, all function here return a fresh copy of the tree, instead of
+//! doing an in-place modification.
+
+use arrayvec::ArrayVec;
+use std::ops::RangeInclusive;
+
+use crate::{
+    algo,
+    ast::{self, make, AstNode},
+    InsertPosition, SyntaxElement,
+};
+
+impl ast::FnDef {
+    #[must_use]
+    pub fn with_body(&self, body: ast::Block) -> ast::FnDef {
+        let mut to_insert: ArrayVec<[SyntaxElement; 2]> = ArrayVec::new();
+        let old_body_or_semi: SyntaxElement = if let Some(old_body) = self.body() {
+            old_body.syntax().clone().into()
+        } else if let Some(semi) = self.semicolon_token() {
+            to_insert.push(make::tokens::single_space().into());
+            semi.into()
+        } else {
+            to_insert.push(make::tokens::single_space().into());
+            to_insert.push(body.syntax().clone().into());
+            return insert_children(self, InsertPosition::Last, to_insert.into_iter());
+        };
+        to_insert.push(body.syntax().clone().into());
+        let replace_range = RangeInclusive::new(old_body_or_semi.clone(), old_body_or_semi);
+        replace_children(self, replace_range, to_insert.into_iter())
+    }
+}
+
+#[must_use]
+fn insert_children<N: AstNode>(
+    parent: &N,
+    position: InsertPosition<SyntaxElement>,
+    mut to_insert: impl Iterator<Item = SyntaxElement>,
+) -> N {
+    let new_syntax = algo::insert_children(parent.syntax(), position, &mut to_insert);
+    N::cast(new_syntax).unwrap()
+}
+
+#[must_use]
+fn replace_children<N: AstNode>(
+    parent: &N,
+    to_replace: RangeInclusive<SyntaxElement>,
+    mut to_insert: impl Iterator<Item = SyntaxElement>,
+) -> N {
+    let new_syntax = algo::replace_children(parent.syntax(), to_replace, &mut to_insert);
+    N::cast(new_syntax).unwrap()
+}
diff --git a/crates/ra_syntax/src/ast/make.rs b/crates/ra_syntax/src/ast/make.rs
index c06c62b3b..287a40bee 100644
--- a/crates/ra_syntax/src/ast/make.rs
+++ b/crates/ra_syntax/src/ast/make.rs
@@ -133,3 +133,51 @@ fn ast_from_text<N: AstNode>(text: &str) -> N {
     let res = parse.tree().syntax().descendants().find_map(N::cast).unwrap();
     res
 }
+
+pub mod tokens {
+    use crate::{AstNode, Parse, SourceFile, SyntaxKind::*, SyntaxToken, T};
+    use once_cell::sync::Lazy;
+
+    static SOURCE_FILE: Lazy<Parse<SourceFile>> = Lazy::new(|| SourceFile::parse(",\n; ;"));
+
+    pub fn comma() -> SyntaxToken {
+        SOURCE_FILE
+            .tree()
+            .syntax()
+            .descendants_with_tokens()
+            .filter_map(|it| it.into_token())
+            .find(|it| it.kind() == T![,])
+            .unwrap()
+    }
+
+    pub fn single_space() -> SyntaxToken {
+        SOURCE_FILE
+            .tree()
+            .syntax()
+            .descendants_with_tokens()
+            .filter_map(|it| it.into_token())
+            .find(|it| it.kind() == WHITESPACE && it.text().as_str() == " ")
+            .unwrap()
+    }
+
+    pub fn single_newline() -> SyntaxToken {
+        SOURCE_FILE
+            .tree()
+            .syntax()
+            .descendants_with_tokens()
+            .filter_map(|it| it.into_token())
+            .find(|it| it.kind() == WHITESPACE && it.text().as_str() == "\n")
+            .unwrap()
+    }
+
+    pub struct WsBuilder(SourceFile);
+
+    impl WsBuilder {
+        pub fn new(text: &str) -> WsBuilder {
+            WsBuilder(SourceFile::parse(text).ok().unwrap())
+        }
+        pub fn ws(&self) -> SyntaxToken {
+            self.0.syntax().first_child_or_token().unwrap().into_token().unwrap()
+        }
+    }
+}