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

   1 use hir::{HasSource, HirDisplay, InFile};
   2 use ide_db::assists::{AssistId, AssistKind};
   3 use syntax::{
   4     ast::{self, make, HasArgList},
   5     match_ast, AstNode, SyntaxNode,
   6 };
   7
   8 use crate::assist_context::{AssistContext, Assists};
   9
  10 // Assist: generate_enum_variant
  11 //
  12 // Adds a variant to an enum.
  13 //
  14 // ```
  15 // enum Countries {
  16 //     Ghana,
  17 // }
  18 //
  19 // fn main() {
  20 //     let country = Countries::Lesotho$0;
  21 // }
  22 // ```
  23 // ->
  24 // ```
  25 // enum Countries {
  26 //     Ghana,
  27 //     Lesotho,
  28 // }
  29 //
  30 // fn main() {
  31 //     let country = Countries::Lesotho;
  32 // }
  33 // ```
  34 pub(crate) fn generate_enum_variant(acc: &mut Assists, ctx: &AssistContext<'_>) -> Option<()> {
  35     let path: ast::Path = ctx.find_node_at_offset()?;
  36     let parent = path_parent(&path)?;
  37
  38     if ctx.sema.resolve_path(&path).is_some() {
  39         // No need to generate anything if the path resolves
  40         return None;
  41     }
  42
  43     let name_ref = path.segment()?.name_ref()?;
  44     if name_ref.text().starts_with(char::is_lowercase) {
  45         // Don't suggest generating variant if the name starts with a lowercase letter
  46         return None;
  47     }
  48
  49     if let Some(hir::PathResolution::Def(hir::ModuleDef::Adt(hir::Adt::Enum(e)))) =
  50         ctx.sema.resolve_path(&path.qualifier()?)
  51     {
  52         let target = path.syntax().text_range();
  53         return add_variant_to_accumulator(acc, ctx, target, e, &name_ref, parent);
  54     }
  55
  56     None
  57 }
  58
  59 #[derive(Debug)]
  60 enum PathParent {
  61     PathExpr(ast::PathExpr),
  62     RecordExpr(ast::RecordExpr),
  63     PathPat(ast::PathPat),
  64     UseTree(ast::UseTree),
  65 }
  66
  67 impl PathParent {
  68     fn syntax(&self) -> &SyntaxNode {
  69         match self {
  70             PathParent::PathExpr(it) => it.syntax(),
  71             PathParent::RecordExpr(it) => it.syntax(),
  72             PathParent::PathPat(it) => it.syntax(),
  73             PathParent::UseTree(it) => it.syntax(),
  74         }
  75     }
  76
  77     fn make_field_list(&self, ctx: &AssistContext<'_>) -> Option<ast::FieldList> {
  78         let scope = ctx.sema.scope(self.syntax())?;
  79
  80         match self {
  81             PathParent::PathExpr(it) => {
  82                 if let Some(call_expr) = it.syntax().parent().and_then(ast::CallExpr::cast) {
  83                     make_tuple_field_list(call_expr, ctx, &scope)
  84                 } else {
  85                     None
  86                 }
  87             }
  88             PathParent::RecordExpr(it) => make_record_field_list(it, ctx, &scope),
  89             PathParent::UseTree(_) | PathParent::PathPat(_) => None,
  90         }
  91     }
  92 }
  93
  94 fn path_parent(path: &ast::Path) -> Option<PathParent> {
  95     let parent = path.syntax().parent()?;
  96
  97     match_ast! {
  98         match parent {
  99             ast::PathExpr(it) => Some(PathParent::PathExpr(it)),
 100             ast::RecordExpr(it) => Some(PathParent::RecordExpr(it)),
 101             ast::PathPat(it) => Some(PathParent::PathPat(it)),
 102             ast::UseTree(it) => Some(PathParent::UseTree(it)),
 103             _ => None
 104         }
 105     }
 106 }
 107
 108 fn add_variant_to_accumulator(
 109     acc: &mut Assists,
 110     ctx: &AssistContext<'_>,
 111     target: syntax::TextRange,
 112     adt: hir::Enum,
 113     name_ref: &ast::NameRef,
 114     parent: PathParent,
 115 ) -> Option<()> {
 116     let db = ctx.db();
 117     let InFile { file_id, value: enum_node } = adt.source(db)?.original_ast_node(db)?;
 118
 119     acc.add(
 120         AssistId("generate_enum_variant", AssistKind::Generate),
 121         "Generate variant",
 122         target,
 123         |builder| {
 124             builder.edit_file(file_id.original_file(db));
 125             let node = builder.make_mut(enum_node);
 126             let variant = make_variant(ctx, name_ref, parent);
 127             node.variant_list().map(|it| it.add_variant(variant.clone_for_update()));
 128         },
 129     )
 130 }
 131
 132 fn make_variant(
 133     ctx: &AssistContext<'_>,
 134     name_ref: &ast::NameRef,
 135     parent: PathParent,
 136 ) -> ast::Variant {
 137     let field_list = parent.make_field_list(ctx);
 138     make::variant(make::name(&name_ref.text()), field_list)
 139 }
 140
 141 fn make_record_field_list(
 142     record: &ast::RecordExpr,
 143     ctx: &AssistContext<'_>,
 144     scope: &hir::SemanticsScope<'_>,
 145 ) -> Option<ast::FieldList> {
 146     let fields = record.record_expr_field_list()?.fields();
 147     let record_fields = fields.map(|field| {
 148         let name = name_from_field(&field);
 149
 150         let ty = field
 151             .expr()
 152             .and_then(|it| expr_ty(ctx, it, scope))
 153             .unwrap_or_else(make::ty_placeholder);
 154
 155         make::record_field(None, name, ty)
 156     });
 157     Some(make::record_field_list(record_fields).into())
 158 }
 159
 160 fn name_from_field(field: &ast::RecordExprField) -> ast::Name {
 161     let text = match field.name_ref() {
 162         Some(it) => it.to_string(),
 163         None => name_from_field_shorthand(field).unwrap_or("unknown".to_string()),
 164     };
 165     make::name(&text)
 166 }
 167
 168 fn name_from_field_shorthand(field: &ast::RecordExprField) -> Option<String> {
 169     let path = match field.expr()? {
 170         ast::Expr::PathExpr(path_expr) => path_expr.path(),
 171         _ => None,
 172     }?;
 173     Some(path.as_single_name_ref()?.to_string())
 174 }
 175
 176 fn make_tuple_field_list(
 177     call_expr: ast::CallExpr,
 178     ctx: &AssistContext<'_>,
 179     scope: &hir::SemanticsScope<'_>,
 180 ) -> Option<ast::FieldList> {
 181     let args = call_expr.arg_list()?.args();
 182     let tuple_fields = args.map(|arg| {
 183         let ty = expr_ty(ctx, arg, &scope).unwrap_or_else(make::ty_placeholder);
 184         make::tuple_field(None, ty)
 185     });
 186     Some(make::tuple_field_list(tuple_fields).into())
 187 }
 188
 189 fn expr_ty(
 190     ctx: &AssistContext<'_>,
 191     arg: ast::Expr,
 192     scope: &hir::SemanticsScope<'_>,
 193 ) -> Option<ast::Type> {
 194     let ty = ctx.sema.type_of_expr(&arg).map(|it| it.adjusted())?;
 195     let text = ty.display_source_code(ctx.db(), scope.module().into()).ok()?;
 196     Some(make::ty(&text))
 197 }
 198
 199 #[cfg(test)]
 200 mod tests {
 201     use crate::tests::{check_assist, check_assist_not_applicable};
 202
 203     use super::*;
 204
 205     #[test]
 206     fn generate_basic_enum_variant_in_empty_enum() {
 207         check_assist(
 208             generate_enum_variant,
 209             r"
 210 enum Foo {}
 211 fn main() {
 212     Foo::Bar$0
 213 }
 214 ",
 215             r"
 216 enum Foo {
 217     Bar,
 218 }
 219 fn main() {
 220     Foo::Bar
 221 }
 222 ",
 223         )
 224     }
 225
 226     #[test]
 227     fn generate_basic_enum_variant_in_non_empty_enum() {
 228         check_assist(
 229             generate_enum_variant,
 230             r"
 231 enum Foo {
 232     Bar,
 233 }
 234 fn main() {
 235     Foo::Baz$0
 236 }
 237 ",
 238             r"
 239 enum Foo {
 240     Bar,
 241     Baz,
 242 }
 243 fn main() {
 244     Foo::Baz
 245 }
 246 ",
 247         )
 248     }
 249
 250     #[test]
 251     fn generate_basic_enum_variant_in_different_file() {
 252         check_assist(
 253             generate_enum_variant,
 254             r"
 255 //- /main.rs
 256 mod foo;
 257 use foo::Foo;
 258
 259 fn main() {
 260     Foo::Baz$0
 261 }
 262
 263 //- /foo.rs
 264 enum Foo {
 265     Bar,
 266 }
 267 ",
 268             r"
 269 enum Foo {
 270     Bar,
 271     Baz,
 272 }
 273 ",
 274         )
 275     }
 276
 277     #[test]
 278     fn not_applicable_for_existing_variant() {
 279         check_assist_not_applicable(
 280             generate_enum_variant,
 281             r"
 282 enum Foo {
 283     Bar,
 284 }
 285 fn main() {
 286     Foo::Bar$0
 287 }
 288 ",
 289         )
 290     }
 291
 292     #[test]
 293     fn not_applicable_for_lowercase() {
 294         check_assist_not_applicable(
 295             generate_enum_variant,
 296             r"
 297 enum Foo {
 298     Bar,
 299 }
 300 fn main() {
 301     Foo::new$0
 302 }
 303 ",
 304         )
 305     }
 306
 307     #[test]
 308     fn indentation_level_is_correct() {
 309         check_assist(
 310             generate_enum_variant,
 311             r"
 312 mod m {
 313     enum Foo {
 314         Bar,
 315     }
 316 }
 317 fn main() {
 318     m::Foo::Baz$0
 319 }
 320 ",
 321             r"
 322 mod m {
 323     enum Foo {
 324         Bar,
 325         Baz,
 326     }
 327 }
 328 fn main() {
 329     m::Foo::Baz
 330 }
 331 ",
 332         )
 333     }
 334
 335     #[test]
 336     fn associated_single_element_tuple() {
 337         check_assist(
 338             generate_enum_variant,
 339             r"
 340 enum Foo {}
 341 fn main() {
 342     Foo::Bar$0(true)
 343 }
 344 ",
 345             r"
 346 enum Foo {
 347     Bar(bool),
 348 }
 349 fn main() {
 350     Foo::Bar(true)
 351 }
 352 ",
 353         )
 354     }
 355
 356     #[test]
 357     fn associated_single_element_tuple_unknown_type() {
 358         check_assist(
 359             generate_enum_variant,
 360             r"
 361 enum Foo {}
 362 fn main() {
 363     Foo::Bar$0(x)
 364 }
 365 ",
 366             r"
 367 enum Foo {
 368     Bar(_),
 369 }
 370 fn main() {
 371     Foo::Bar(x)
 372 }
 373 ",
 374         )
 375     }
 376
 377     #[test]
 378     fn associated_multi_element_tuple() {
 379         check_assist(
 380             generate_enum_variant,
 381             r"
 382 struct Struct {}
 383 enum Foo {}
 384 fn main() {
 385     Foo::Bar$0(true, x, Struct {})
 386 }
 387 ",
 388             r"
 389 struct Struct {}
 390 enum Foo {
 391     Bar(bool, _, Struct),
 392 }
 393 fn main() {
 394     Foo::Bar(true, x, Struct {})
 395 }
 396 ",
 397         )
 398     }
 399
 400     #[test]
 401     fn associated_record() {
 402         check_assist(
 403             generate_enum_variant,
 404             r"
 405 enum Foo {}
 406 fn main() {
 407     Foo::$0Bar { x: true }
 408 }
 409 ",
 410             r"
 411 enum Foo {
 412     Bar { x: bool },
 413 }
 414 fn main() {
 415     Foo::Bar { x: true }
 416 }
 417 ",
 418         )
 419     }
 420
 421     #[test]
 422     fn associated_record_unknown_type() {
 423         check_assist(
 424             generate_enum_variant,
 425             r"
 426 enum Foo {}
 427 fn main() {
 428     Foo::$0Bar { x: y }
 429 }
 430 ",
 431             r"
 432 enum Foo {
 433     Bar { x: _ },
 434 }
 435 fn main() {
 436     Foo::Bar { x: y }
 437 }
 438 ",
 439         )
 440     }
 441
 442     #[test]
 443     fn associated_record_field_shorthand() {
 444         check_assist(
 445             generate_enum_variant,
 446             r"
 447 enum Foo {}
 448 fn main() {
 449     let x = true;
 450     Foo::$0Bar { x }
 451 }
 452 ",
 453             r"
 454 enum Foo {
 455     Bar { x: bool },
 456 }
 457 fn main() {
 458     let x = true;
 459     Foo::Bar { x }
 460 }
 461 ",
 462         )
 463     }
 464
 465     #[test]
 466     fn associated_record_field_shorthand_unknown_type() {
 467         check_assist(
 468             generate_enum_variant,
 469             r"
 470 enum Foo {}
 471 fn main() {
 472     Foo::$0Bar { x }
 473 }
 474 ",
 475             r"
 476 enum Foo {
 477     Bar { x: _ },
 478 }
 479 fn main() {
 480     Foo::Bar { x }
 481 }
 482 ",
 483         )
 484     }
 485
 486     #[test]
 487     fn associated_record_field_multiple_fields() {
 488         check_assist(
 489             generate_enum_variant,
 490             r"
 491 struct Struct {}
 492 enum Foo {}
 493 fn main() {
 494     Foo::$0Bar { x, y: x, s: Struct {} }
 495 }
 496 ",
 497             r"
 498 struct Struct {}
 499 enum Foo {
 500     Bar { x: _, y: _, s: Struct },
 501 }
 502 fn main() {
 503     Foo::Bar { x, y: x, s: Struct {} }
 504 }
 505 ",
 506         )
 507     }
 508
 509     #[test]
 510     fn use_tree() {
 511         check_assist(
 512             generate_enum_variant,
 513             r"
 514 //- /main.rs
 515 mod foo;
 516 use foo::Foo::Bar$0;
 517
 518 //- /foo.rs
 519 enum Foo {}
 520 ",
 521             r"
 522 enum Foo {
 523     Bar,
 524 }
 525 ",
 526         )
 527     }
 528
 529     #[test]
 530     fn not_applicable_for_path_type() {
 531         check_assist_not_applicable(
 532             generate_enum_variant,
 533             r"
 534 enum Foo {}
 535 impl Foo::Bar$0 {}
 536 ",
 537         )
 538     }
 539
 540     #[test]
 541     fn path_pat() {
 542         check_assist(
 543             generate_enum_variant,
 544             r"
 545 enum Foo {}
 546 fn foo(x: Foo) {
 547     match x {
 548         Foo::Bar$0 =>
 549     }
 550 }
 551 ",
 552             r"
 553 enum Foo {
 554     Bar,
 555 }
 556 fn foo(x: Foo) {
 557     match x {
 558         Foo::Bar =>
 559     }
 560 }
 561 ",
 562         )
 563     }
 564 }