1 files changed, 11 insertions, 5 deletions

diff --git a/crates/ra_assists/src/assist_ctx.rs b/crates/ra_assists/src/assist_ctx.rs
index c3e653299..2fe7c3de3 100644
--- a/crates/ra_assists/src/assist_ctx.rs
+++ b/crates/ra_assists/src/assist_ctx.rs
@@ -5,12 +5,12 @@ use ra_fmt::{leading_indent, reindent};
 use ra_ide_db::RootDatabase;
 use ra_syntax::{
     algo::{self, find_covering_element, find_node_at_offset},
-    AstNode, SourceFile, SyntaxElement, SyntaxKind, SyntaxNode, SyntaxToken, TextRange, TextUnit,
+    AstNode, SourceFile, SyntaxElement, SyntaxKind, SyntaxNode, SyntaxToken, TextRange, TextSize,
     TokenAtOffset,
 };
 use ra_text_edit::TextEditBuilder;
 
-use crate::{AssistAction, AssistId, AssistLabel, GroupLabel, ResolvedAssist};
+use crate::{AssistAction, AssistFile, AssistId, AssistLabel, GroupLabel, ResolvedAssist};
 use algo::SyntaxRewriter;
 
 #[derive(Clone, Debug)]
@@ -178,8 +178,9 @@ impl<'a> AssistGroup<'a> {
 #[derive(Default)]
 pub(crate) struct ActionBuilder {
     edit: TextEditBuilder,
-    cursor_position: Option<TextUnit>,
+    cursor_position: Option<TextSize>,
     target: Option<TextRange>,
+    file: AssistFile,
 }
 
 impl ActionBuilder {
@@ -210,12 +211,12 @@ impl ActionBuilder {
     }
 
     /// Append specified `text` at the given `offset`
-    pub(crate) fn insert(&mut self, offset: TextUnit, text: impl Into<String>) {
+    pub(crate) fn insert(&mut self, offset: TextSize, text: impl Into<String>) {
         self.edit.insert(offset, text.into())
     }
 
     /// Specify desired position of the cursor after the assist is applied.
-    pub(crate) fn set_cursor(&mut self, offset: TextUnit) {
+    pub(crate) fn set_cursor(&mut self, offset: TextSize) {
         self.cursor_position = Some(offset)
     }
 
@@ -241,11 +242,16 @@ impl ActionBuilder {
         algo::diff(&node, &new).into_text_edit(&mut self.edit)
     }
 
+    pub(crate) fn set_file(&mut self, assist_file: AssistFile) {
+        self.file = assist_file
+    }
+
     fn build(self) -> AssistAction {
         AssistAction {
             edit: self.edit.finish(),
             cursor_position: self.cursor_position,
             target: self.target,
+            file: self.file,
         }
     }
 }