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 snip = format!("::<{}>", get_snippet_fish_head(number_of_arguments));
  97                     builder.insert_snippet(cap, ident.text_range().end(), snip)
  98                 }
  99                 None => {
 100                     let fish_head = std::iter::repeat("_").take(number_of_arguments).format(", ");
 101                     let snip = format!("::<{}>", fish_head);
 102                     builder.insert(ident.text_range().end(), snip);
 103                 }
 104             }
 105         },
 106     )
 107 }
 108
 109 /// This will create a snippet string with tabstops marked
 110 fn get_snippet_fish_head(number_of_arguments: usize) -> String {
 111     let mut fish_head = (1..number_of_arguments)
 112         .format_with("", |i, f| f(&format_args!("${{{}:_}}, ", i)))
 113         .to_string();
 114
 115     // tabstop 0 is a special case and always the last one
 116     fish_head.push_str("${0:_}");
 117     fish_head
 118 }
 119
 120 #[cfg(test)]
 121 mod tests {
 122     use crate::tests::{check_assist, check_assist_by_label, check_assist_not_applicable};
 123
 124     use super::*;
 125
 126     #[test]
 127     fn add_turbo_fish_function() {
 128         check_assist(
 129             add_turbo_fish,
 130             r#"
 131 fn make<T>() -> T {}
 132 fn main() {
 133     make$0();
 134 }
 135 "#,
 136             r#"
 137 fn make<T>() -> T {}
 138 fn main() {
 139     make::<${0:_}>();
 140 }
 141 "#,
 142         );
 143     }
 144
 145     #[test]
 146     fn add_turbo_fish_function_multiple_generic_types() {
 147         check_assist(
 148             add_turbo_fish,
 149             r#"
 150 fn make<T, A>() -> T {}
 151 fn main() {
 152     make$0();
 153 }
 154 "#,
 155             r#"
 156 fn make<T, A>() -> T {}
 157 fn main() {
 158     make::<${1:_}, ${0:_}>();
 159 }
 160 "#,
 161         );
 162     }
 163
 164     #[test]
 165     fn add_turbo_fish_function_many_generic_types() {
 166         check_assist(
 167             add_turbo_fish,
 168             r#"
 169 fn make<T, A, B, C, D, E, F>() -> T {}
 170 fn main() {
 171     make$0();
 172 }
 173 "#,
 174             r#"
 175 fn make<T, A, B, C, D, E, F>() -> T {}
 176 fn main() {
 177     make::<${1:_}, ${2:_}, ${3:_}, ${4:_}, ${5:_}, ${6:_}, ${0:_}>();
 178 }
 179 "#,
 180         );
 181     }
 182
 183     #[test]
 184     fn add_turbo_fish_after_call() {
 185         cov_mark::check!(add_turbo_fish_after_call);
 186         check_assist(
 187             add_turbo_fish,
 188             r#"
 189 fn make<T>() -> T {}
 190 fn main() {
 191     make()$0;
 192 }
 193 "#,
 194             r#"
 195 fn make<T>() -> T {}
 196 fn main() {
 197     make::<${0:_}>();
 198 }
 199 "#,
 200         );
 201     }
 202
 203     #[test]
 204     fn add_turbo_fish_method() {
 205         check_assist(
 206             add_turbo_fish,
 207             r#"
 208 struct S;
 209 impl S {
 210     fn make<T>(&self) -> T {}
 211 }
 212 fn main() {
 213     S.make$0();
 214 }
 215 "#,
 216             r#"
 217 struct S;
 218 impl S {
 219     fn make<T>(&self) -> T {}
 220 }
 221 fn main() {
 222     S.make::<${0:_}>();
 223 }
 224 "#,
 225         );
 226     }
 227
 228     #[test]
 229     fn add_turbo_fish_one_fish_is_enough() {
 230         cov_mark::check!(add_turbo_fish_one_fish_is_enough);
 231         check_assist_not_applicable(
 232             add_turbo_fish,
 233             r#"
 234 fn make<T>() -> T {}
 235 fn main() {
 236     make$0::<()>();
 237 }
 238 "#,
 239         );
 240     }
 241
 242     #[test]
 243     fn add_turbo_fish_non_generic() {
 244         cov_mark::check!(add_turbo_fish_non_generic);
 245         check_assist_not_applicable(
 246             add_turbo_fish,
 247             r#"
 248 fn make() -> () {}
 249 fn main() {
 250     make$0();
 251 }
 252 "#,
 253         );
 254     }
 255
 256     #[test]
 257     fn add_type_ascription_function() {
 258         check_assist_by_label(
 259             add_turbo_fish,
 260             r#"
 261 fn make<T>() -> T {}
 262 fn main() {
 263     let x = make$0();
 264 }
 265 "#,
 266             r#"
 267 fn make<T>() -> T {}
 268 fn main() {
 269     let x: ${0:_} = make();
 270 }
 271 "#,
 272             "Add `: _` before assignment operator",
 273         );
 274     }
 275
 276     #[test]
 277     fn add_type_ascription_after_call() {
 278         cov_mark::check!(add_type_ascription_after_call);
 279         check_assist_by_label(
 280             add_turbo_fish,
 281             r#"
 282 fn make<T>() -> T {}
 283 fn main() {
 284     let x = make()$0;
 285 }
 286 "#,
 287             r#"
 288 fn make<T>() -> T {}
 289 fn main() {
 290     let x: ${0:_} = make();
 291 }
 292 "#,
 293             "Add `: _` before assignment operator",
 294         );
 295     }
 296
 297     #[test]
 298     fn add_type_ascription_method() {
 299         check_assist_by_label(
 300             add_turbo_fish,
 301             r#"
 302 struct S;
 303 impl S {
 304     fn make<T>(&self) -> T {}
 305 }
 306 fn main() {
 307     let x = S.make$0();
 308 }
 309 "#,
 310             r#"
 311 struct S;
 312 impl S {
 313     fn make<T>(&self) -> T {}
 314 }
 315 fn main() {
 316     let x: ${0:_} = S.make();
 317 }
 318 "#,
 319             "Add `: _` before assignment operator",
 320         );
 321     }
 322
 323     #[test]
 324     fn add_type_ascription_already_typed() {
 325         cov_mark::check!(add_type_ascription_already_typed);
 326         check_assist(
 327             add_turbo_fish,
 328             r#"
 329 fn make<T>() -> T {}
 330 fn main() {
 331     let x: () = make$0();
 332 }
 333 "#,
 334             r#"
 335 fn make<T>() -> T {}
 336 fn main() {
 337     let x: () = make::<${0:_}>();
 338 }
 339 "#,
 340         );
 341     }
 342
 343     #[test]
 344     fn add_type_ascription_append_semicolon() {
 345         check_assist_by_label(
 346             add_turbo_fish,
 347             r#"
 348 fn make<T>() -> T {}
 349 fn main() {
 350     let x = make$0()
 351 }
 352 "#,
 353             r#"
 354 fn make<T>() -> T {}
 355 fn main() {
 356     let x: ${0:_} = make();
 357 }
 358 "#,
 359             "Add `: _` before assignment operator",
 360         );
 361     }
 362
 363     #[test]
 364     fn add_turbo_fish_function_lifetime_parameter() {
 365         check_assist(
 366             add_turbo_fish,
 367             r#"
 368 fn make<'a, T, A>(t: T, a: A) {}
 369 fn main() {
 370     make$0(5, 2);
 371 }
 372 "#,
 373             r#"
 374 fn make<'a, T, A>(t: T, a: A) {}
 375 fn main() {
 376     make::<${1:_}, ${0:_}>(5, 2);
 377 }
 378 "#,
 379         );
 380     }
 381
 382     #[test]
 383     fn add_turbo_fish_function_const_parameter() {
 384         check_assist(
 385             add_turbo_fish,
 386             r#"
 387 fn make<T, const N: usize>(t: T) {}
 388 fn main() {
 389     make$0(3);
 390 }
 391 "#,
 392             r#"
 393 fn make<T, const N: usize>(t: T) {}
 394 fn main() {
 395     make::<${1:_}, ${0:_}>(3);
 396 }
 397 "#,
 398         );
 399     }
 400 }