1 files changed, 19 insertions, 25 deletions

diff --git a/crates/ra_assists/src/handlers/add_function.rs b/crates/ra_assists/src/handlers/add_function.rs
index 6c7456579..278079665 100644
--- a/crates/ra_assists/src/handlers/add_function.rs
+++ b/crates/ra_assists/src/handlers/add_function.rs
@@ -3,9 +3,10 @@ use ra_syntax::{
     SyntaxKind, SyntaxNode, TextSize,
 };
 
-use crate::{Assist, AssistCtx, AssistFile, AssistId};
+use crate::{Assist, AssistCtx, AssistId};
 use ast::{edit::IndentLevel, ArgListOwner, ModuleItemOwner};
 use hir::HirDisplay;
+use ra_db::FileId;
 use rustc_hash::{FxHashMap, FxHashSet};
 
 // Assist: add_function
@@ -57,14 +58,12 @@ pub(crate) fn add_function(ctx: AssistCtx) -> Option<Assist> {
 
     let function_builder = FunctionBuilder::from_call(&ctx, &call, &path, target_module)?;
 
-    ctx.add_assist(AssistId("add_function"), "Add function", |edit| {
-        edit.target(call.syntax().text_range());
-
-        if let Some(function_template) = function_builder.render() {
-            edit.set_file(function_template.file);
-            edit.set_cursor(function_template.cursor_offset);
-            edit.insert(function_template.insert_offset, function_template.fn_def.to_string());
-        }
+    let target = call.syntax().text_range();
+    ctx.add_assist(AssistId("add_function"), "Add function", target, |edit| {
+        let function_template = function_builder.render();
+        edit.set_file(function_template.file);
+        edit.set_cursor(function_template.cursor_offset);
+        edit.insert(function_template.insert_offset, function_template.fn_def.to_string());
     })
 }
 
@@ -72,7 +71,7 @@ struct FunctionTemplate {
     insert_offset: TextSize,
     cursor_offset: TextSize,
     fn_def: ast::SourceFile,
-    file: AssistFile,
+    file: FileId,
 }
 
 struct FunctionBuilder {
@@ -80,7 +79,7 @@ struct FunctionBuilder {
     fn_name: ast::Name,
     type_params: Option<ast::TypeParamList>,
     params: ast::ParamList,
-    file: AssistFile,
+    file: FileId,
     needs_pub: bool,
 }
 
@@ -94,7 +93,7 @@ impl FunctionBuilder {
         target_module: Option<hir::InFile<hir::ModuleSource>>,
     ) -> Option<Self> {
         let needs_pub = target_module.is_some();
-        let mut file = AssistFile::default();
+        let mut file = ctx.frange.file_id;
         let target = if let Some(target_module) = target_module {
             let (in_file, target) = next_space_for_fn_in_module(ctx.sema.db, target_module)?;
             file = in_file;
@@ -107,7 +106,7 @@ impl FunctionBuilder {
         Some(Self { target, fn_name, type_params, params, file, needs_pub })
     }
 
-    fn render(self) -> Option<FunctionTemplate> {
+    fn render(self) -> FunctionTemplate {
         let placeholder_expr = ast::make::expr_todo();
         let fn_body = ast::make::block_expr(vec![], Some(placeholder_expr));
         let mut fn_def = ast::make::fn_def(self.fn_name, self.type_params, self.params, fn_body);
@@ -133,15 +132,11 @@ impl FunctionBuilder {
             }
         };
 
-        let cursor_offset_from_fn_start = fn_def
-            .syntax()
-            .descendants()
-            .find_map(ast::MacroCall::cast)?
-            .syntax()
-            .text_range()
-            .start();
+        let placeholder_expr =
+            fn_def.syntax().descendants().find_map(ast::MacroCall::cast).unwrap();
+        let cursor_offset_from_fn_start = placeholder_expr.syntax().text_range().start();
         let cursor_offset = insert_offset + cursor_offset_from_fn_start;
-        Some(FunctionTemplate { insert_offset, cursor_offset, fn_def, file: self.file })
+        FunctionTemplate { insert_offset, cursor_offset, fn_def, file: self.file }
     }
 }
 
@@ -259,9 +254,8 @@ fn next_space_for_fn_after_call_site(expr: &ast::CallExpr) -> Option<GeneratedFu
 fn next_space_for_fn_in_module(
     db: &dyn hir::db::AstDatabase,
     module: hir::InFile<hir::ModuleSource>,
-) -> Option<(AssistFile, GeneratedFunctionTarget)> {
+) -> Option<(FileId, GeneratedFunctionTarget)> {
     let file = module.file_id.original_file(db);
-    let assist_file = AssistFile::TargetFile(file);
     let assist_item = match module.value {
         hir::ModuleSource::SourceFile(it) => {
             if let Some(last_item) = it.items().last() {
@@ -278,12 +272,12 @@ fn next_space_for_fn_in_module(
             }
         }
     };
-    Some((assist_file, assist_item))
+    Some((file, assist_item))
 }
 
 #[cfg(test)]
 mod tests {
-    use crate::helpers::{check_assist, check_assist_not_applicable};
+    use crate::tests::{check_assist, check_assist_not_applicable};
 
     use super::*;