refactor method generation assist

Author: mahdifrmz

crates/ide_assists/src/handlers/generate_function.rs

@@ -6,7 +6,7 @@ use syntax::{
     ast::{
         self,
         edit::{AstNodeEdit, IndentLevel},
-        make, ArgList, ArgListOwner, AstNode, ModuleItemOwner,
+        make, ArgListOwner, AstNode, ModuleItemOwner,
     },
     SyntaxKind, SyntaxNode, TextSize,
 };
@@ -17,27 +17,6 @@ use crate::{
     AssistContext, AssistId, AssistKind, Assists,
 };
 
-enum FuncExpr {
-    Func(ast::CallExpr),
-    Method(ast::MethodCallExpr),
-}
-
-impl FuncExpr {
-    fn arg_list(&self) -> Option<ArgList> {
-        match self {
-            FuncExpr::Func(fn_call) => fn_call.arg_list(),
-            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
 //
 // Adds a stub function with a signature matching the function under the cursor.
@@ -64,6 +43,31 @@ impl FuncExpr {
 //
 // ```
 pub(crate) fn generate_function(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
+    gen_fn(acc, ctx).or_else(|| gen_method(acc, ctx))
+}
+
+enum FuncExpr {
+    Func(ast::CallExpr),
+    Method(ast::MethodCallExpr),
+}
+
+impl FuncExpr {
+    fn arg_list(&self) -> Option<ast::ArgList> {
+        match self {
+            FuncExpr::Func(fn_call) => fn_call.arg_list(),
+            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(),
+        }
+    }
+}
+
+fn gen_fn(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
     let path_expr: ast::PathExpr = ctx.find_node_at_offset()?;
     let call = path_expr.syntax().parent().and_then(ast::CallExpr::cast)?;
 
@@ -100,7 +104,7 @@ pub(crate) fn generate_function(acc: &mut Assists, ctx: &AssistContext) -> Optio
     )
 }
 
-pub(crate) fn generate_method(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
+fn gen_method(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
     let call: ast::MethodCallExpr = ctx.find_node_at_offset()?;
     let fn_name: ast::NameRef = ast::NameRef::cast(
         call.syntax().children().find(|child| child.kind() == SyntaxKind::NAME_REF)?,
@@ -1351,8 +1355,7 @@ fn bar(baz: ()) {}
 
     #[test]
     fn create_method_with_no_args() {
-        // FIXME: This is wrong, this should just work.
-        check_assist_not_applicable(
+        check_assist(
             generate_function,
             r#"
 struct Foo;
@@ -1361,7 +1364,19 @@ impl Foo {
         self.bar()$0;
     }
 }
-        "#,
+"#,
+            r#"
+struct Foo;
+impl Foo {
+    fn foo(&self) {
+        self.bar();
+    }
+
+    fn bar(&self) ${0:-> ()} {
+        todo!()
+    }
+}
+"#,
         )
     }
 
@@ -1389,52 +1404,38 @@ async fn bar(arg: i32) ${0:-> ()} {
     #[test]
     fn create_method() {
         check_assist(
-            generate_method,
+            generate_function,
             r"
 struct S;
-
-fn foo() {
-    S.bar$0();
-}
-
+fn foo() {S.bar$0();}
 ",
             r"
 struct S;
-
-fn foo() {
-    S.bar();
-}
+fn foo() {S.bar();}
 impl S {
 
 
 fn bar(&self) ${0:-> ()} {
     todo!()
 }
 }
-
 ",
         )
     }
 
     #[test]
     fn create_method_within_an_impl() {
         check_assist(
-            generate_method,
+            generate_function,
             r"
 struct S;
-
-fn foo() {
-    S.bar$0();
-}
+fn foo() {S.bar$0();}
 impl S {}
 
 ",
             r"
 struct S;
-
-fn foo() {
-    S.bar();
-}
+fn foo() {S.bar();}
 impl S {
     fn bar(&self) ${0:-> ()} {
         todo!()
@@ -1448,15 +1449,12 @@ impl S {
     #[test]
     fn create_method_from_different_module() {
         check_assist(
-            generate_method,
+            generate_function,
             r"
 mod s {
     pub struct S;
 }
-fn foo() {
-    s::S.bar$0();
-}
-
+fn foo() {s::S.bar$0();}
 ",
             r"
 mod s {
@@ -1469,18 +1467,15 @@ impl S {
     }
 }
 }
-fn foo() {
-    s::S.bar();
-}
-
+fn foo() {s::S.bar();}
 ",
         )
     }
 
     #[test]
     fn create_method_from_descendant_module() {
         check_assist(
-            generate_method,
+            generate_function,
             r"
 struct S;
 mod s {
@@ -1512,29 +1507,21 @@ fn bar(&self) ${0:-> ()} {
     #[test]
     fn create_method_with_cursor_anywhere_on_call_expresion() {
         check_assist(
-            generate_method,
+            generate_function,
             r"
 struct S;
-
-fn foo() {
-    $0S.bar();
-}
-
+fn foo() {$0S.bar();}
 ",
             r"
 struct S;
-
-fn foo() {
-    S.bar();
-}
+fn foo() {S.bar();}
 impl S {
 
 
 fn bar(&self) ${0:-> ()} {
     todo!()
 }
 }
-
 ",
         )
     }

crates/ide_assists/src/lib.rs

@@ -151,7 +151,6 @@ mod handlers {
             generate_enum_projection_method::generate_enum_try_into_method,
             generate_from_impl_for_enum::generate_from_impl_for_enum,
             generate_function::generate_function,
-            generate_function::generate_method,
             generate_getter::generate_getter,
             generate_getter::generate_getter_mut,
             generate_impl::generate_impl,