src/tools/rust-analyzer/crates/ide-assists/src/handlers/add_turbo_fish.rs

   1 use ide_db::defs::{Definition, NameRefClass};
   2 use itertools::Itertools;
   3 use syntax::{ast, AstNode, SyntaxKind, T};
   4
   5 use crate::{
   6     assist_context::{AssistContext, Assists},
   7     AssistId, AssistKind,
   8 };
   9
  10 // Assist: add_turbo_fish
  11 //
  12 // Adds `::<_>` to a call of a generic method or function.
  13 //
  14 // ```
  15 // fn make<T>() -> T { todo!() }
  16 // fn main() {
  17 //     let x = make$0();
  18 // }
  19 // ```
  20 // ->
  21 // ```
  22 // fn make<T>() -> T { todo!() }
  23 // fn main() {
  24 //     let x = make::<${0:_}>();
  25 // }
  26 // ```
  27 pub(crate) fn add_turbo_fish(acc: &mut Assists, ctx: &AssistContext<'_>) -> Option<()> {
  28     let ident = ctx.find_token_syntax_at_offset(SyntaxKind::IDENT).or_else(|| {
  29         let arg_list = ctx.find_node_at_offset::<ast::ArgList>()?;
  30         if arg_list.args().next().is_some() {
  31             return None;
  32         }
  33         cov_mark::hit!(add_turbo_fish_after_call);
  34         cov_mark::hit!(add_type_ascription_after_call);
  35         arg_list.l_paren_token()?.prev_token().filter(|it| it.kind() == SyntaxKind::IDENT)
  36     })?;
  37     let next_token = ident.next_token()?;
  38     if next_token.kind() == T![::] {
  39         cov_mark::hit!(add_turbo_fish_one_fish_is_enough);
  40         return None;
  41     }
  42     let name_ref = ast::NameRef::cast(ident.parent()?)?;
  43     let def = match NameRefClass::classify(&ctx.sema, &name_ref)? {
  44         NameRefClass::Definition(def) => def,
  45         NameRefClass::FieldShorthand { .. } => return None,
  46     };
  47     let fun = match def {
  48         Definition::Function(it) => it,
  49         _ => return None,
  50     };
  51     let generics = hir::GenericDef::Function(fun).params(ctx.sema.db);
  52     if generics.is_empty() {
  53         cov_mark::hit!(add_turbo_fish_non_generic);
  54         return None;
  55     }
  56
  57     if let Some(let_stmt) = ctx.find_node_at_offset::<ast::LetStmt>() {
  58         if let_stmt.colon_token().is_none() {
  59             let type_pos = let_stmt.pat()?.syntax().last_token()?.text_range().end();
  60             let semi_pos = let_stmt.syntax().last_token()?.text_range().end();
  61
  62             acc.add(
  63                 AssistId("add_type_ascription", AssistKind::RefactorRewrite),
  64                 "Add `: _` before assignment operator",
  65                 ident.text_range(),
  66                 |builder| {
  67                     if let_stmt.semicolon_token().is_none() {
  68                         builder.insert(semi_pos, ";");
  69                     }
  70                     match ctx.config.snippet_cap {
  71                         Some(cap) => builder.insert_snippet(cap, type_pos, ": ${0:_}"),
  72                         None => builder.insert(type_pos, ": _"),
  73                     }
  74                 },
  75             )?
  76         } else {
  77             cov_mark::hit!(add_type_ascription_already_typed);
  78         }
  79     }
  80
  81     let number_of_arguments = generics
  82         .iter()
  83         .filter(|param| {
  84             matches!(param, hir::GenericParam::TypeParam(_) | hir::GenericParam::ConstParam(_))
  85         })
  86         .count();
  87
  88     acc.add(
  89         AssistId("add_turbo_fish", AssistKind::RefactorRewrite),
  90         "Add `::<>`",
  91         ident.text_range(),
  92         |builder| {
  93             builder.trigger_signature_help();
  94             match ctx.config.snippet_cap {
  95                 Some(cap) => {
  96                     let fish_head = get_snippet_fish_head(number_of_arguments);
  97                     let snip = format!("::<{fish_head}>");
  98                     builder.insert_snippet(cap, ident.text_range().end(), snip)
  99                 }
 100                 None => {
 101                     let fish_head = std::iter::repeat("_").take(number_of_arguments).format(", ");
 102                     let snip = format!("::<{fish_head}>");
 103                     builder.insert(ident.text_range().end(), snip);
 104                 }
 105             }
 106         },
 107     )
 108 }
 109
 110 /// This will create a snippet string with tabstops marked
 111 fn get_snippet_fish_head(number_of_arguments: usize) -> String {
 112     let mut fish_head = (1..number_of_arguments)
 113         .format_with("", |i, f| f(&format_args!("${{{i}:_}}, ")))
 114         .to_string();
 115
 116     // tabstop 0 is a special case and always the last one
 117     fish_head.push_str("${0:_}");
 118     fish_head
 119 }
 120
 121 #[cfg(test)]
 122 mod tests {
 123     use crate::tests::{check_assist, check_assist_by_label, check_assist_not_applicable};
 124
 125     use super::*;
 126
 127     #[test]
 128     fn add_turbo_fish_function() {
 129         check_assist(
 130             add_turbo_fish,
 131             r#"
 132 fn make<T>() -> T {}
 133 fn main() {
 134     make$0();
 135 }
 136 "#,
 137             r#"
 138 fn make<T>() -> T {}
 139 fn main() {
 140     make::<${0:_}>();
 141 }
 142 "#,
 143         );
 144     }
 145
 146     #[test]
 147     fn add_turbo_fish_function_multiple_generic_types() {
 148         check_assist(
 149             add_turbo_fish,
 150             r#"
 151 fn make<T, A>() -> T {}
 152 fn main() {
 153     make$0();
 154 }
 155 "#,
 156             r#"
 157 fn make<T, A>() -> T {}
 158 fn main() {
 159     make::<${1:_}, ${0:_}>();
 160 }
 161 "#,
 162         );
 163     }
 164
 165     #[test]
 166     fn add_turbo_fish_function_many_generic_types() {
 167         check_assist(
 168             add_turbo_fish,
 169             r#"
 170 fn make<T, A, B, C, D, E, F>() -> T {}
 171 fn main() {
 172     make$0();
 173 }
 174 "#,
 175             r#"
 176 fn make<T, A, B, C, D, E, F>() -> T {}
 177 fn main() {
 178     make::<${1:_}, ${2:_}, ${3:_}, ${4:_}, ${5:_}, ${6:_}, ${0:_}>();
 179 }
 180 "#,
 181         );
 182     }
 183
 184     #[test]
 185     fn add_turbo_fish_after_call() {
 186         cov_mark::check!(add_turbo_fish_after_call);
 187         check_assist(
 188             add_turbo_fish,
 189             r#"
 190 fn make<T>() -> T {}
 191 fn main() {
 192     make()$0;
 193 }
 194 "#,
 195             r#"
 196 fn make<T>() -> T {}
 197 fn main() {
 198     make::<${0:_}>();
 199 }
 200 "#,
 201         );
 202     }
 203
 204     #[test]
 205     fn add_turbo_fish_method() {
 206         check_assist(
 207             add_turbo_fish,
 208             r#"
 209 struct S;
 210 impl S {
 211     fn make<T>(&self) -> T {}
 212 }
 213 fn main() {
 214     S.make$0();
 215 }
 216 "#,
 217             r#"
 218 struct S;
 219 impl S {
 220     fn make<T>(&self) -> T {}
 221 }
 222 fn main() {
 223     S.make::<${0:_}>();
 224 }
 225 "#,
 226         );
 227     }
 228
 229     #[test]
 230     fn add_turbo_fish_one_fish_is_enough() {
 231         cov_mark::check!(add_turbo_fish_one_fish_is_enough);
 232         check_assist_not_applicable(
 233             add_turbo_fish,
 234             r#"
 235 fn make<T>() -> T {}
 236 fn main() {
 237     make$0::<()>();
 238 }
 239 "#,
 240         );
 241     }
 242
 243     #[test]
 244     fn add_turbo_fish_non_generic() {
 245         cov_mark::check!(add_turbo_fish_non_generic);
 246         check_assist_not_applicable(
 247             add_turbo_fish,
 248             r#"
 249 fn make() -> () {}
 250 fn main() {
 251     make$0();
 252 }
 253 "#,
 254         );
 255     }
 256
 257     #[test]
 258     fn add_type_ascription_function() {
 259         check_assist_by_label(
 260             add_turbo_fish,
 261             r#"
 262 fn make<T>() -> T {}
 263 fn main() {
 264     let x = make$0();
 265 }
 266 "#,
 267             r#"
 268 fn make<T>() -> T {}
 269 fn main() {
 270     let x: ${0:_} = make();
 271 }
 272 "#,
 273             "Add `: _` before assignment operator",
 274         );
 275     }
 276
 277     #[test]
 278     fn add_type_ascription_after_call() {
 279         cov_mark::check!(add_type_ascription_after_call);
 280         check_assist_by_label(
 281             add_turbo_fish,
 282             r#"
 283 fn make<T>() -> T {}
 284 fn main() {
 285     let x = make()$0;
 286 }
 287 "#,
 288             r#"
 289 fn make<T>() -> T {}
 290 fn main() {
 291     let x: ${0:_} = make();
 292 }
 293 "#,
 294             "Add `: _` before assignment operator",
 295         );
 296     }
 297
 298     #[test]
 299     fn add_type_ascription_method() {
 300         check_assist_by_label(
 301             add_turbo_fish,
 302             r#"
 303 struct S;
 304 impl S {
 305     fn make<T>(&self) -> T {}
 306 }
 307 fn main() {
 308     let x = S.make$0();
 309 }
 310 "#,
 311             r#"
 312 struct S;
 313 impl S {
 314     fn make<T>(&self) -> T {}
 315 }
 316 fn main() {
 317     let x: ${0:_} = S.make();
 318 }
 319 "#,
 320             "Add `: _` before assignment operator",
 321         );
 322     }
 323
 324     #[test]
 325     fn add_type_ascription_already_typed() {
 326         cov_mark::check!(add_type_ascription_already_typed);
 327         check_assist(
 328             add_turbo_fish,
 329             r#"
 330 fn make<T>() -> T {}
 331 fn main() {
 332     let x: () = make$0();
 333 }
 334 "#,
 335             r#"
 336 fn make<T>() -> T {}
 337 fn main() {
 338     let x: () = make::<${0:_}>();
 339 }
 340 "#,
 341         );
 342     }
 343
 344     #[test]
 345     fn add_type_ascription_append_semicolon() {
 346         check_assist_by_label(
 347             add_turbo_fish,
 348             r#"
 349 fn make<T>() -> T {}
 350 fn main() {
 351     let x = make$0()
 352 }
 353 "#,
 354             r#"
 355 fn make<T>() -> T {}
 356 fn main() {
 357     let x: ${0:_} = make();
 358 }
 359 "#,
 360             "Add `: _` before assignment operator",
 361         );
 362     }
 363
 364     #[test]
 365     fn add_turbo_fish_function_lifetime_parameter() {
 366         check_assist(
 367             add_turbo_fish,
 368             r#"
 369 fn make<'a, T, A>(t: T, a: A) {}
 370 fn main() {
 371     make$0(5, 2);
 372 }
 373 "#,
 374             r#"
 375 fn make<'a, T, A>(t: T, a: A) {}
 376 fn main() {
 377     make::<${1:_}, ${0:_}>(5, 2);
 378 }
 379 "#,
 380         );
 381     }
 382
 383     #[test]
 384     fn add_turbo_fish_function_const_parameter() {
 385         check_assist(
 386             add_turbo_fish,
 387             r#"
 388 fn make<T, const N: usize>(t: T) {}
 389 fn main() {
 390     make$0(3);
 391 }
 392 "#,
 393             r#"
 394 fn make<T, const N: usize>(t: T) {}
 395 fn main() {
 396     make::<${1:_}, ${0:_}>(3);
 397 }
 398 "#,
 399         );
 400     }
 401 }