3 files changed, 19 insertions, 18 deletions

diff --git a/crates/syntax/src/algo.rs b/crates/syntax/src/algo.rs
index a153a9e1c..c9229c4e0 100644
--- a/crates/syntax/src/algo.rs
+++ b/crates/syntax/src/algo.rs
@@ -342,10 +342,10 @@ enum InsertPos {
 
 #[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>>,
+    _pd: std::marker::PhantomData<&'a ()>,
 }
 
 impl fmt::Debug for SyntaxRewriter<'_> {
@@ -357,14 +357,7 @@ 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(),
-            insertions: IndexMap::default(),
-        }
-    }
+impl SyntaxRewriter<'_> {
     pub fn delete<T: Clone + Into<SyntaxElement>>(&mut self, what: &T) {
         let what = what.clone().into();
         let replacement = Replacement::Delete;
@@ -470,7 +463,7 @@ impl<'a> SyntaxRewriter<'a> {
     pub fn rewrite(&self, node: &SyntaxNode) -> SyntaxNode {
         let _p = profile::span("rewrite");
 
-        if self.f.is_none() && self.replacements.is_empty() && self.insertions.is_empty() {
+        if self.replacements.is_empty() && self.insertions.is_empty() {
             return node.clone();
         }
         let green = self.rewrite_children(node);
@@ -495,7 +488,6 @@ impl<'a> SyntaxRewriter<'a> {
             }
         }
 
-        assert!(self.f.is_none());
         self.replacements
             .keys()
             .filter_map(element_to_node_or_parent)
@@ -510,10 +502,6 @@ impl<'a> SyntaxRewriter<'a> {
     }
 
     fn replacement(&self, element: &SyntaxElement) -> Option<Replacement> {
-        if let Some(f) = &self.f {
-            assert!(self.replacements.is_empty());
-            return f(element).map(Replacement::Single);
-        }
         self.replacements.get(element).cloned()
     }
 
@@ -574,7 +562,6 @@ fn element_to_green(element: SyntaxElement) -> NodeOrToken<rowan::GreenNode, row
 
 impl ops::AddAssign for SyntaxRewriter<'_> {
     fn add_assign(&mut self, rhs: SyntaxRewriter) {
-        assert!(rhs.f.is_none());
         self.replacements.extend(rhs.replacements);
         for (pos, insertions) in rhs.insertions.into_iter() {
             match self.insertions.entry(pos) {
diff --git a/crates/syntax/src/ast/make.rs b/crates/syntax/src/ast/make.rs
index 222b7e212..42da09606 100644
--- a/crates/syntax/src/ast/make.rs
+++ b/crates/syntax/src/ast/make.rs
@@ -632,6 +632,7 @@ pub mod tokens {
         SOURCE_FILE
             .tree()
             .syntax()
+            .clone_for_update()
             .descendants_with_tokens()
             .filter_map(|it| it.into_token())
             .find(|it| it.kind() == WHITESPACE && it.text() == "\n\n")
diff --git a/crates/syntax/src/ted.rs b/crates/syntax/src/ted.rs
index 450f2e447..91a06101f 100644
--- a/crates/syntax/src/ted.rs
+++ b/crates/syntax/src/ted.rs
@@ -7,7 +7,7 @@ use std::{mem, ops::RangeInclusive};
 use parser::T;
 
 use crate::{
-    ast::{edit::IndentLevel, make},
+    ast::{self, edit::IndentLevel, make, AstNode},
     SyntaxElement, SyntaxKind, SyntaxNode, SyntaxToken,
 };
 
@@ -147,6 +147,16 @@ pub fn append_child_raw(node: &(impl Into<SyntaxNode> + Clone), child: impl Elem
 fn ws_before(position: &Position, new: &SyntaxElement) -> Option<SyntaxToken> {
     let prev = match &position.repr {
         PositionRepr::FirstChild(_) => return None,
+        PositionRepr::After(it) if it.kind() == SyntaxKind::L_CURLY => {
+            if new.kind() == SyntaxKind::USE {
+                if let Some(item_list) = it.parent().and_then(ast::ItemList::cast) {
+                    let mut indent = IndentLevel::from_element(&item_list.syntax().clone().into());
+                    indent.0 += 1;
+                    return Some(make::tokens::whitespace(&format!("\n{}", indent)));
+                }
+            }
+            it
+        }
         PositionRepr::After(it) => it,
     };
     ws_between(prev, new)
@@ -173,7 +183,10 @@ fn ws_between(left: &SyntaxElement, right: &SyntaxElement) -> Option<SyntaxToken
     }
 
     if right.kind() == SyntaxKind::USE {
-        let indent = IndentLevel::from_element(left);
+        let mut indent = IndentLevel::from_element(left);
+        if left.kind() == SyntaxKind::USE {
+            indent.0 = IndentLevel::from_element(right).0.max(indent.0);
+        }
         return Some(make::tokens::whitespace(&format!("\n{}", indent)));
     }
     Some(make::tokens::single_space())