From 99570f32d863504946bde3151eb02b7ca97f9399 Mon Sep 17 00:00:00 2001
From: mahdi-frms <mahdif1380@outlook.com>
Date: Thu, 5 Aug 2021 16:48:51 +0430
Subject: [PATCH] refactor: use single next space

---
 .../src/handlers/generate_function.rs         | 35 ++++++-------------
 1 file changed, 10 insertions(+), 25 deletions(-)

diff --git a/crates/ide_assists/src/handlers/generate_function.rs b/crates/ide_assists/src/handlers/generate_function.rs
index 6791392c50d8..bd1be5f98d65 100644
--- a/crates/ide_assists/src/handlers/generate_function.rs
+++ b/crates/ide_assists/src/handlers/generate_function.rs
@@ -29,6 +29,13 @@ impl<'a> FuncExpr<'a> {
             FuncExpr::Method(m_call) => m_call.arg_list(),
         }
     }
+
+    fn syntax(&self) -> &SyntaxNode {
+        match *self {
+            FuncExpr::Func(fn_call) => fn_call.syntax(),
+            FuncExpr::Method(m_call) => m_call.syntax(),
+        }
+    }
 }
 
 // Assist: generate_function
@@ -173,7 +180,7 @@ impl FunctionBuilder {
                 file = in_file;
                 target
             }
-            None => next_space_for_fn_after_call_site(call)?,
+            None => next_space_for_fn_after_call_site(FuncExpr::Func(call))?,
         };
         let needs_pub = target_module.is_some();
         let target_module = target_module.or_else(|| ctx.sema.scope(target.syntax()).module())?;
@@ -238,7 +245,7 @@ impl FunctionBuilder {
                 file = in_file;
                 target
             }
-            None => next_space_for_fn_after_method_call_site(call)?,
+            None => next_space_for_fn_after_call_site(FuncExpr::Method(call))?,
         };
         let needs_pub = false;
         let target_module = target_module.or_else(|| ctx.sema.scope(target.syntax()).module())?;
@@ -465,29 +472,7 @@ fn fn_arg_type(
 /// directly after the current block
 /// We want to write the generated function directly after
 /// fns, impls or macro calls, but inside mods
-fn next_space_for_fn_after_call_site(expr: &ast::CallExpr) -> Option<GeneratedFunctionTarget> {
-    let mut ancestors = expr.syntax().ancestors().peekable();
-    let mut last_ancestor: Option<SyntaxNode> = None;
-    while let Some(next_ancestor) = ancestors.next() {
-        match next_ancestor.kind() {
-            SyntaxKind::SOURCE_FILE => {
-                break;
-            }
-            SyntaxKind::ITEM_LIST => {
-                if ancestors.peek().map(|a| a.kind()) == Some(SyntaxKind::MODULE) {
-                    break;
-                }
-            }
-            _ => {}
-        }
-        last_ancestor = Some(next_ancestor);
-    }
-    last_ancestor.map(GeneratedFunctionTarget::BehindItem)
-}
-
-fn next_space_for_fn_after_method_call_site(
-    expr: &ast::MethodCallExpr,
-) -> Option<GeneratedFunctionTarget> {
+fn next_space_for_fn_after_call_site(expr: FuncExpr) -> Option<GeneratedFunctionTarget> {
     let mut ancestors = expr.syntax().ancestors().peekable();
     let mut last_ancestor: Option<SyntaxNode> = None;
     while let Some(next_ancestor) = ancestors.next() {