2 files changed, 36 insertions, 8 deletions

diff --git a/crates/ra_assists/src/assist_ctx.rs b/crates/ra_assists/src/assist_ctx.rs
index e80e35738..17a9041c5 100644
--- a/crates/ra_assists/src/assist_ctx.rs
+++ b/crates/ra_assists/src/assist_ctx.rs
@@ -144,6 +144,10 @@ impl AssistBuilder {
         self.replace(node.range(), replace_with)
     }
 
+    pub(crate) fn set_edit_builder(&mut self, edit: TextEditBuilder) {
+        self.edit = edit;
+    }
+
     #[allow(unused)]
     pub(crate) fn delete(&mut self, range: TextRange) {
         self.edit.delete(range)
diff --git a/crates/ra_assists/src/auto_import.rs b/crates/ra_assists/src/auto_import.rs
index 3fdf6b0d9..6bc5caf76 100644
--- a/crates/ra_assists/src/auto_import.rs
+++ b/crates/ra_assists/src/auto_import.rs
@@ -1,3 +1,4 @@
+use ra_text_edit::TextEditBuilder;
 use hir::db::HirDatabase;
 
 use ra_syntax::{
@@ -374,7 +375,7 @@ fn best_action_for_target<'b, 'a: 'b>(
     }
 }
 
-fn make_assist(action: &ImportAction, target: &[&ast::PathSegment], edit: &mut AssistBuilder) {
+fn make_assist(action: &ImportAction, target: &[&ast::PathSegment], edit: &mut TextEditBuilder) {
     match action {
         ImportAction::AddNewUse { anchor, add_after_anchor } => {
             make_assist_add_new_use(anchor, *add_after_anchor, target, edit)
@@ -408,7 +409,7 @@ fn make_assist_add_new_use(
     anchor: &Option<&SyntaxNode>,
     after: bool,
     target: &[&ast::PathSegment],
-    edit: &mut AssistBuilder,
+    edit: &mut TextEditBuilder,
 ) {
     if let Some(anchor) = anchor {
         let indent = ra_fmt::leading_indent(anchor);
@@ -437,7 +438,7 @@ fn make_assist_add_in_tree_list(
     tree_list: &ast::UseTreeList,
     target: &[&ast::PathSegment],
     add_self: bool,
-    edit: &mut AssistBuilder,
+    edit: &mut TextEditBuilder,
 ) {
     let last = tree_list.use_trees().last();
     if let Some(last) = last {
@@ -466,7 +467,7 @@ fn make_assist_add_nested_import(
     first_segment_to_split: &Option<&ast::PathSegment>,
     target: &[&ast::PathSegment],
     add_self: bool,
-    edit: &mut AssistBuilder,
+    edit: &mut TextEditBuilder,
 ) {
     let use_tree = path.syntax().ancestors().find_map(ast::UseTree::cast);
     if let Some(use_tree) = use_tree {
@@ -491,7 +492,7 @@ fn make_assist_add_nested_import(
             buf.push_str(", ");
         }
         edit.insert(start, buf);
-        edit.insert(end, "}");
+        edit.insert(end, "}".to_string());
     }
 }
 
@@ -499,7 +500,7 @@ fn apply_auto_import<'a>(
     container: &SyntaxNode,
     path: &ast::Path,
     target: &[&'a ast::PathSegment],
-    edit: &mut AssistBuilder,
+    edit: &mut TextEditBuilder,
 ) {
     let action = best_action_for_target(container, path, target);
     make_assist(&action, target, edit);
@@ -513,6 +514,25 @@ fn apply_auto_import<'a>(
     }
 }
 
+pub fn auto_import_text_edit<'a>(
+    position: &SyntaxNode,
+    path: &ast::Path,
+    target: &[&'a ast::PathSegment],
+    edit: &mut TextEditBuilder,
+) {
+    let container = position.ancestors().find_map(|n| {
+        if let Some(module) = ast::Module::cast(n) {
+            return module.item_list().map(ast::AstNode::syntax);
+        }
+        ast::SourceFile::cast(n).map(ast::AstNode::syntax)
+    });
+
+    if let Some(container) = container {
+        let action = best_action_for_target(container, path, target);
+        make_assist(&action, target, edit);
+    }
+}
+
 pub(crate) fn auto_import(mut ctx: AssistCtx<impl HirDatabase>) -> Option<Assist> {
     let path: &ast::Path = ctx.node_at_offset()?;
     // We don't want to mess with use statements
@@ -531,7 +551,9 @@ pub(crate) fn auto_import(mut ctx: AssistCtx<impl HirDatabase>) -> Option<Assist
                 AssistId("auto_import"),
                 format!("import {} in mod {}", fmt_segments(&segments), name.text()),
                 |edit| {
-                    apply_auto_import(item_list.syntax(), path, &segments, edit);
+                    let mut text_edit = TextEditBuilder::default();
+                    apply_auto_import(item_list.syntax(), path, &segments, &mut text_edit);
+                    edit.set_edit_builder(text_edit);
                 },
             );
         }
@@ -541,7 +563,9 @@ pub(crate) fn auto_import(mut ctx: AssistCtx<impl HirDatabase>) -> Option<Assist
             AssistId("auto_import"),
             format!("import {} in the current file", fmt_segments(&segments)),
             |edit| {
-                apply_auto_import(current_file.syntax(), path, &segments, edit);
+                let mut text_edit = TextEditBuilder::default();
+                apply_auto_import(current_file.syntax(), path, &segments, &mut text_edit);
+                edit.set_edit_builder(text_edit);
             },
         );
     }