Generate function assist creates bad param names for const/static item args

Author: xffxff

crates/ide_assists/src/handlers/generate_function.rs

@@ -1,5 +1,9 @@
-use hir::{HasSource, HirDisplay, Module, TypeInfo};
-use ide_db::{base_db::FileId, helpers::SnippetCap};
+use hir::{HasSource, HirDisplay, Module, ModuleDef, TypeInfo};
+use ide_db::{
+    base_db::FileId,
+    defs::{Definition, NameRefClass},
+    helpers::SnippetCap,
+};
 use rustc_hash::{FxHashMap, FxHashSet};
 use stdx::to_lower_snake_case;
 use syntax::{
@@ -438,7 +442,7 @@ fn fn_args(
     let mut arg_names = Vec::new();
     let mut arg_types = Vec::new();
     for arg in call.arg_list()?.args() {
-        arg_names.push(fn_arg_name(&arg));
+        arg_names.push(fn_arg_name(ctx, &arg));
         arg_types.push(match fn_arg_type(ctx, target_module, &arg) {
             Some(ty) => {
                 if !ty.is_empty() && ty.starts_with('&') {
@@ -503,12 +507,25 @@ fn deduplicate_arg_names(arg_names: &mut Vec<String>) {
     }
 }
 
-fn fn_arg_name(arg_expr: &ast::Expr) -> String {
+fn fn_arg_name(ctx: &AssistContext, arg_expr: &ast::Expr) -> String {
     let name = (|| match arg_expr {
-        ast::Expr::CastExpr(cast_expr) => Some(fn_arg_name(&cast_expr.expr()?)),
+        ast::Expr::CastExpr(cast_expr) => Some(fn_arg_name(ctx, &cast_expr.expr()?)),
         expr => {
-            let s = expr.syntax().descendants().filter_map(ast::NameRef::cast).last()?.to_string();
-            Some(to_lower_snake_case(&s))
+            let name_ref = expr.syntax().descendants().filter_map(ast::NameRef::cast).last()?;
+            if let Some(NameRefClass::Definition(def)) =
+                NameRefClass::classify(&ctx.sema, &name_ref)
+            {
+                match def {
+                    Definition::ModuleDef(ModuleDef::Const(_)) => {
+                        return Some(name_ref.to_string().to_lowercase());
+                    }
+                    Definition::ModuleDef(ModuleDef::Static(_)) => {
+                        return Some(name_ref.to_string().to_lowercase());
+                    }
+                    _ => {}
+                }
+            };
+            Some(to_lower_snake_case(&name_ref.to_string()))
         }
     })();
     match name {
@@ -1683,6 +1700,75 @@ fn main() {
 fn foo(arg0: ()) ${0:-> _} {
     todo!()
 }
+",
+        )
+    }
+
+    #[test]
+    fn add_function_with_const_arg() {
+        check_assist(
+            generate_function,
+            r"
+const VALUE: usize = 0;
+fn main() {
+    foo$0(VALUE);
+}
+",
+            r"
+const VALUE: usize = 0;
+fn main() {
+    foo(VALUE);
+}
+
+fn foo(value: usize) ${0:-> _} {
+    todo!()
+}
+",
+        )
+    }
+
+    #[test]
+    fn add_function_with_static_arg() {
+        check_assist(
+            generate_function,
+            r"
+static VALUE: usize = 0;
+fn main() {
+    foo$0(VALUE);
+}
+",
+            r"
+static VALUE: usize = 0;
+fn main() {
+    foo(VALUE);
+}
+
+fn foo(value: usize) ${0:-> _} {
+    todo!()
+}
+",
+        )
+    }
+
+    #[test]
+    fn add_function_with_static_mut_arg() {
+        check_assist(
+            generate_function,
+            r"
+static mut VALUE: usize = 0;
+fn main() {
+    foo$0(VALUE);
+}
+",
+            r"
+static mut VALUE: usize = 0;
+fn main() {
+    foo(VALUE);
+}
+
+fn foo(value: usize) ${0:-> _} {
+    todo!()
+}
 ",
         )
     }