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