From a5edf6de79126c350e7f1544041079f5f25ceaa8 Mon Sep 17 00:00:00 2001
From: mahdi-frms <mahdif1380@outlook.com>
Date: Sat, 31 Jul 2021 15:43:22 +0430
Subject: [PATCH] generate function assist favors deref cmpt types

---
 .../src/handlers/generate_function.rs         | 13 +++++++++-
 .../src/handlers/generate_getter.rs           | 26 ++-----------------
 crates/ide_assists/src/utils.rs               | 23 ++++++++++++++++
 3 files changed, 37 insertions(+), 25 deletions(-)

diff --git a/crates/ide_assists/src/handlers/generate_function.rs b/crates/ide_assists/src/handlers/generate_function.rs
index c620cfac486b..d06e3f709424 100644
--- a/crates/ide_assists/src/handlers/generate_function.rs
+++ b/crates/ide_assists/src/handlers/generate_function.rs
@@ -12,6 +12,7 @@ use syntax::{
 };
 
 use crate::{
+    utils::useless_type_special_case,
     utils::{render_snippet, Cursor},
     AssistContext, AssistId, AssistKind, Assists,
 };
@@ -257,7 +258,17 @@ fn fn_args(
             None => String::from("arg"),
         });
         arg_types.push(match fn_arg_type(ctx, target_module, &arg) {
-            Some(ty) => ty,
+            Some(ty) => {
+                if ty.len() > 0 && ty.starts_with('&') {
+                    if let Some((new_ty, _)) = useless_type_special_case("", &ty[1..].to_owned()) {
+                        new_ty
+                    } else {
+                        ty
+                    }
+                } else {
+                    ty
+                }
+            }
             None => String::from("()"),
         });
     }
diff --git a/crates/ide_assists/src/handlers/generate_getter.rs b/crates/ide_assists/src/handlers/generate_getter.rs
index cc020c92c522..739f3e592fff 100644
--- a/crates/ide_assists/src/handlers/generate_getter.rs
+++ b/crates/ide_assists/src/handlers/generate_getter.rs
@@ -2,6 +2,7 @@ use stdx::{format_to, to_lower_snake_case};
 use syntax::ast::{self, AstNode, NameOwner, VisibilityOwner};
 
 use crate::{
+    utils::useless_type_special_case,
     utils::{find_impl_block_end, find_struct_impl, generate_impl_text},
     AssistContext, AssistId, AssistKind, Assists, GroupLabel,
 };
@@ -99,7 +100,7 @@ pub(crate) fn generate_getter_impl(
             let (ty, body) = if mutable {
                 (format!("&mut {}", field_ty), format!("&mut self.{}", field_name))
             } else {
-                useless_type_special_case(&field_name.to_string(), &field_ty)
+                useless_type_special_case(&field_name.to_string(), &field_ty.to_string())
                     .unwrap_or_else(|| (format!("&{}", field_ty), format!("&self.{}", field_name)))
             };
 
@@ -136,29 +137,6 @@ pub(crate) fn generate_getter_impl(
     )
 }
 
-fn useless_type_special_case(field_name: &str, field_ty: &ast::Type) -> Option<(String, String)> {
-    if field_ty.to_string() == "String" {
-        cov_mark::hit!(useless_type_special_case);
-        return Some(("&str".to_string(), format!("self.{}.as_str()", field_name)));
-    }
-    if let Some(arg) = ty_ctor(field_ty, "Vec") {
-        return Some((format!("&[{}]", arg), format!("self.{}.as_slice()", field_name)));
-    }
-    if let Some(arg) = ty_ctor(field_ty, "Box") {
-        return Some((format!("&{}", arg), format!("self.{}.as_ref()", field_name)));
-    }
-    if let Some(arg) = ty_ctor(field_ty, "Option") {
-        return Some((format!("Option<&{}>", arg), format!("self.{}.as_ref()", field_name)));
-    }
-    None
-}
-
-// FIXME: This should rely on semantic info.
-fn ty_ctor(ty: &ast::Type, ctor: &str) -> Option<String> {
-    let res = ty.to_string().strip_prefix(ctor)?.strip_prefix('<')?.strip_suffix('>')?.to_string();
-    Some(res)
-}
-
 #[cfg(test)]
 mod tests {
     use crate::tests::{check_assist, check_assist_not_applicable};
diff --git a/crates/ide_assists/src/utils.rs b/crates/ide_assists/src/utils.rs
index 99c91463b98c..81463745a461 100644
--- a/crates/ide_assists/src/utils.rs
+++ b/crates/ide_assists/src/utils.rs
@@ -493,3 +493,26 @@ pub(crate) fn add_method_to_adt(
 
     builder.insert(start_offset, buf);
 }
+
+pub fn useless_type_special_case(field_name: &str, field_ty: &String) -> Option<(String, String)> {
+    if field_ty.to_string() == "String" {
+        cov_mark::hit!(useless_type_special_case);
+        return Some(("&str".to_string(), format!("self.{}.as_str()", field_name)));
+    }
+    if let Some(arg) = ty_ctor(field_ty, "Vec") {
+        return Some((format!("&[{}]", arg), format!("self.{}.as_slice()", field_name)));
+    }
+    if let Some(arg) = ty_ctor(field_ty, "Box") {
+        return Some((format!("&{}", arg), format!("self.{}.as_ref()", field_name)));
+    }
+    if let Some(arg) = ty_ctor(field_ty, "Option") {
+        return Some((format!("Option<&{}>", arg), format!("self.{}.as_ref()", field_name)));
+    }
+    None
+}
+
+// FIXME: This should rely on semantic info.
+fn ty_ctor(ty: &String, ctor: &str) -> Option<String> {
+    let res = ty.to_string().strip_prefix(ctor)?.strip_prefix('<')?.strip_suffix('>')?.to_string();
+    Some(res)
+}