crates/ide_assists/src/handlers/add_turbo_fish.rs

   1 use ide_db::defs::{Definition, NameRefClass};
   2 use syntax::{ast, AstNode, SyntaxKind, T};
   3
   4 use crate::{
   5     assist_context::{AssistContext, Assists},
   6     AssistId, AssistKind,
   7 };
   8
   9 // Assist: add_turbo_fish
  10 //
  11 // Adds `::<_>` to a call of a generic method or function.
  12 //
  13 // ```
  14 // fn make<T>() -> T { todo!() }
  15 // fn main() {
  16 //     let x = make$0();
  17 // }
  18 // ```
  19 // ->
  20 // ```
  21 // fn make<T>() -> T { todo!() }
  22 // fn main() {
  23 //     let x = make::<${0:_}>();
  24 // }
  25 // ```
  26 pub(crate) fn add_turbo_fish(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
  27     let ident = ctx.find_token_syntax_at_offset(SyntaxKind::IDENT).or_else(|| {
  28         let arg_list = ctx.find_node_at_offset::<ast::ArgList>()?;
  29         if arg_list.args().next().is_some() {
  30             return None;
  31         }
  32         cov_mark::hit!(add_turbo_fish_after_call);
  33         cov_mark::hit!(add_type_ascription_after_call);
  34         arg_list.l_paren_token()?.prev_token().filter(|it| it.kind() == SyntaxKind::IDENT)
  35     })?;
  36     let next_token = ident.next_token()?;
  37     if next_token.kind() == T![::] {
  38         cov_mark::hit!(add_turbo_fish_one_fish_is_enough);
  39         return None;
  40     }
  41     let name_ref = ast::NameRef::cast(ident.parent()?)?;
  42     let def = match NameRefClass::classify(&ctx.sema, &name_ref)? {
  43         NameRefClass::Definition(def) => def,
  44         NameRefClass::FieldShorthand { .. } => return None,
  45     };
  46     let fun = match def {
  47         Definition::Function(it) => it,
  48         _ => return None,
  49     };
  50     let generics = hir::GenericDef::Function(fun).params(ctx.sema.db);
  51     if generics.is_empty() {
  52         cov_mark::hit!(add_turbo_fish_non_generic);
  53         return None;
  54     }
  55
  56     if let Some(let_stmt) = ctx.find_node_at_offset::<ast::LetStmt>() {
  57         if let_stmt.colon_token().is_none() {
  58             let type_pos = let_stmt.pat()?.syntax().last_token()?.text_range().end();
  59             let semi_pos = let_stmt.syntax().last_token()?.text_range().end();
  60
  61             acc.add(
  62                 AssistId("add_type_ascription", AssistKind::RefactorRewrite),
  63                 "Add `: _` before assignment operator",
  64                 ident.text_range(),
  65                 |builder| {
  66                     if let_stmt.semicolon_token().is_none() {
  67                         builder.insert(semi_pos, ";");
  68                     }
  69                     match ctx.config.snippet_cap {
  70                         Some(cap) => builder.insert_snippet(cap, type_pos, ": ${0:_}"),
  71                         None => builder.insert(type_pos, ": _"),
  72                     }
  73                 },
  74             )?
  75         } else {
  76             cov_mark::hit!(add_type_ascription_already_typed);
  77         }
  78     }
  79
  80     let number_of_arguments = generics
  81         .iter()
  82         .filter(|param| {
  83             matches!(param, hir::GenericParam::TypeParam(_) | hir::GenericParam::ConstParam(_))
  84         })
  85         .count();
  86
  87     acc.add(
  88         AssistId("add_turbo_fish", AssistKind::RefactorRewrite),
  89         "Add `::<>`",
  90         ident.text_range(),
  91         |builder| match ctx.config.snippet_cap {
  92             Some(cap) => {
  93                 let snip = format!("::<{}>", get_snippet_fish_head(number_of_arguments));
  94                 builder.insert_snippet(cap, ident.text_range().end(), snip)
  95             }
  96             None => {
  97                 let fish_head =
  98                     std::iter::repeat("_").take(number_of_arguments).collect::<Vec<_>>().join(", ");
  99                 let snip = format!("::<{}>", fish_head);
 100                 builder.insert(ident.text_range().end(), snip);
 101             }
 102         },
 103     )
 104 }
 105
 106 /// This will create a snippet string with tabstops marked
 107 fn get_snippet_fish_head(number_of_arguments: usize) -> String {
 108     let mut fish_head = String::new();
 109     let mut i = 1;
 110     while i < number_of_arguments {
 111         fish_head.push_str(&format!("${{{}:_}},", i));
 112         i = i + 1;
 113     }
 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 }