1 files changed, 356 insertions, 19 deletions
diff --git a/src/tools/rust-analyzer/crates/ide-assists/src/handlers/generate_enum_variant.rs b/src/tools/rust-analyzer/crates/ide-assists/src/handlers/generate_enum_variant.rs
index 4461fbd5a..35cd42908 100644
--- a/src/tools/rust-analyzer/crates/ide-assists/src/handlers/generate_enum_variant.rs
+++ b/src/tools/rust-analyzer/crates/ide-assists/src/handlers/generate_enum_variant.rs
@@ -1,8 +1,8 @@
-use hir::{HasSource, InFile};
+use hir::{HasSource, HirDisplay, InFile};
 use ide_db::assists::{AssistId, AssistKind};
 use syntax::{
-    ast::{self, edit::IndentLevel},
-    AstNode, TextSize,
+    ast::{self, make, HasArgList},
+    match_ast, AstNode, SyntaxNode,
 };
 
 use crate::assist_context::{AssistContext, Assists};
@@ -32,8 +32,8 @@ use crate::assist_context::{AssistContext, Assists};
 // }
 // ```
 pub(crate) fn generate_enum_variant(acc: &mut Assists, ctx: &AssistContext<'_>) -> Option<()> {
-    let path_expr: ast::PathExpr = ctx.find_node_at_offset()?;
-    let path = path_expr.path()?;
+    let path: ast::Path = ctx.find_node_at_offset()?;
+    let parent = path_parent(&path)?;
 
     if ctx.sema.resolve_path(&path).is_some() {
         // No need to generate anything if the path resolves
@@ -50,26 +50,71 @@ pub(crate) fn generate_enum_variant(acc: &mut Assists, ctx: &AssistContext<'_>)
         ctx.sema.resolve_path(&path.qualifier()?)
     {
         let target = path.syntax().text_range();
-        return add_variant_to_accumulator(acc, ctx, target, e, &name_ref);
+        return add_variant_to_accumulator(acc, ctx, target, e, &name_ref, parent);
     }
 
     None
 }
 
+#[derive(Debug)]
+enum PathParent {
+    PathExpr(ast::PathExpr),
+    RecordExpr(ast::RecordExpr),
+    PathPat(ast::PathPat),
+    UseTree(ast::UseTree),
+}
+
+impl PathParent {
+    fn syntax(&self) -> &SyntaxNode {
+        match self {
+            PathParent::PathExpr(it) => it.syntax(),
+            PathParent::RecordExpr(it) => it.syntax(),
+            PathParent::PathPat(it) => it.syntax(),
+            PathParent::UseTree(it) => it.syntax(),
+        }
+    }
+
+    fn make_field_list(&self, ctx: &AssistContext<'_>) -> Option<ast::FieldList> {
+        let scope = ctx.sema.scope(self.syntax())?;
+
+        match self {
+            PathParent::PathExpr(it) => {
+                if let Some(call_expr) = it.syntax().parent().and_then(ast::CallExpr::cast) {
+                    make_tuple_field_list(call_expr, ctx, &scope)
+                } else {
+                    None
+                }
+            }
+            PathParent::RecordExpr(it) => make_record_field_list(it, ctx, &scope),
+            PathParent::UseTree(_) | PathParent::PathPat(_) => None,
+        }
+    }
+}
+
+fn path_parent(path: &ast::Path) -> Option<PathParent> {
+    let parent = path.syntax().parent()?;
+
+    match_ast! {
+        match parent {
+            ast::PathExpr(it) => Some(PathParent::PathExpr(it)),
+            ast::RecordExpr(it) => Some(PathParent::RecordExpr(it)),
+            ast::PathPat(it) => Some(PathParent::PathPat(it)),
+            ast::UseTree(it) => Some(PathParent::UseTree(it)),
+            _ => None
+        }
+    }
+}
+
 fn add_variant_to_accumulator(
     acc: &mut Assists,
     ctx: &AssistContext<'_>,
     target: syntax::TextRange,
     adt: hir::Enum,
     name_ref: &ast::NameRef,
+    parent: PathParent,
 ) -> Option<()> {
     let db = ctx.db();
     let InFile { file_id, value: enum_node } = adt.source(db)?.original_ast_node(db)?;
-    let enum_indent = IndentLevel::from_node(&enum_node.syntax());
-
-    let variant_list = enum_node.variant_list()?;
-    let offset = variant_list.syntax().text_range().end() - TextSize::of('}');
-    let empty_enum = variant_list.variants().next().is_none();
 
     acc.add(
         AssistId("generate_enum_variant", AssistKind::Generate),
@@ -77,18 +122,80 @@ fn add_variant_to_accumulator(
         target,
         |builder| {
             builder.edit_file(file_id.original_file(db));
-            let text = format!(
-                "{maybe_newline}{indent_1}{name},\n{enum_indent}",
-                maybe_newline = if empty_enum { "\n" } else { "" },
-                indent_1 = IndentLevel(1),
-                name = name_ref,
-                enum_indent = enum_indent
-            );
-            builder.insert(offset, text)
+            let node = builder.make_mut(enum_node);
+            let variant = make_variant(ctx, name_ref, parent);
+            node.variant_list().map(|it| it.add_variant(variant.clone_for_update()));
         },
     )
 }
 
+fn make_variant(
+    ctx: &AssistContext<'_>,
+    name_ref: &ast::NameRef,
+    parent: PathParent,
+) -> ast::Variant {
+    let field_list = parent.make_field_list(ctx);
+    make::variant(make::name(&name_ref.text()), field_list)
+}
+
+fn make_record_field_list(
+    record: &ast::RecordExpr,
+    ctx: &AssistContext<'_>,
+    scope: &hir::SemanticsScope<'_>,
+) -> Option<ast::FieldList> {
+    let fields = record.record_expr_field_list()?.fields();
+    let record_fields = fields.map(|field| {
+        let name = name_from_field(&field);
+
+        let ty = field
+            .expr()
+            .and_then(|it| expr_ty(ctx, it, scope))
+            .unwrap_or_else(make::ty_placeholder);
+
+        make::record_field(None, name, ty)
+    });
+    Some(make::record_field_list(record_fields).into())
+}
+
+fn name_from_field(field: &ast::RecordExprField) -> ast::Name {
+    let text = match field.name_ref() {
+        Some(it) => it.to_string(),
+        None => name_from_field_shorthand(field).unwrap_or("unknown".to_string()),
+    };
+    make::name(&text)
+}
+
+fn name_from_field_shorthand(field: &ast::RecordExprField) -> Option<String> {
+    let path = match field.expr()? {
+        ast::Expr::PathExpr(path_expr) => path_expr.path(),
+        _ => None,
+    }?;
+    Some(path.as_single_name_ref()?.to_string())
+}
+
+fn make_tuple_field_list(
+    call_expr: ast::CallExpr,
+    ctx: &AssistContext<'_>,
+    scope: &hir::SemanticsScope<'_>,
+) -> Option<ast::FieldList> {
+    let args = call_expr.arg_list()?.args();
+    let tuple_fields = args.map(|arg| {
+        let ty = expr_ty(ctx, arg, &scope).unwrap_or_else(make::ty_placeholder);
+        make::tuple_field(None, ty)
+    });
+    Some(make::tuple_field_list(tuple_fields).into())
+}
+
+fn expr_ty(
+    ctx: &AssistContext<'_>,
+    arg: ast::Expr,
+    scope: &hir::SemanticsScope<'_>,
+) -> Option<ast::Type> {
+    let ty = ctx.sema.type_of_expr(&arg).map(|it| it.adjusted())?;
+    let text = ty.display_source_code(ctx.db(), scope.module().into()).ok()?;
+    Some(make::ty(&text))
+}
+
 #[cfg(test)]
 mod tests {
     use crate::tests::{check_assist, check_assist_not_applicable};
@@ -224,4 +331,234 @@ fn main() {
 ",
         )
     }
+
+    #[test]
+    fn associated_single_element_tuple() {
+        check_assist(
+            generate_enum_variant,
+            r"
+enum Foo {}
+fn main() {
+    Foo::Bar$0(true)
+}
+",
+            r"
+enum Foo {
+    Bar(bool),
+}
+fn main() {
+    Foo::Bar(true)
+}
+",
+        )
+    }
+
+    #[test]
+    fn associated_single_element_tuple_unknown_type() {
+        check_assist(
+            generate_enum_variant,
+            r"
+enum Foo {}
+fn main() {
+    Foo::Bar$0(x)
+}
+",
+            r"
+enum Foo {
+    Bar(_),
+}
+fn main() {
+    Foo::Bar(x)
+}
+",
+        )
+    }
+
+    #[test]
+    fn associated_multi_element_tuple() {
+        check_assist(
+            generate_enum_variant,
+            r"
+struct Struct {}
+enum Foo {}
+fn main() {
+    Foo::Bar$0(true, x, Struct {})
+}
+",
+            r"
+struct Struct {}
+enum Foo {
+    Bar(bool, _, Struct),
+}
+fn main() {
+    Foo::Bar(true, x, Struct {})
+}
+",
+        )
+    }
+
+    #[test]
+    fn associated_record() {
+        check_assist(
+            generate_enum_variant,
+            r"
+enum Foo {}
+fn main() {
+    Foo::$0Bar { x: true }
+}
+",
+            r"
+enum Foo {
+    Bar { x: bool },
+}
+fn main() {
+    Foo::Bar { x: true }
+}
+",
+        )
+    }
+
+    #[test]
+    fn associated_record_unknown_type() {
+        check_assist(
+            generate_enum_variant,
+            r"
+enum Foo {}
+fn main() {
+    Foo::$0Bar { x: y }
+}
+",
+            r"
+enum Foo {
+    Bar { x: _ },
+}
+fn main() {
+    Foo::Bar { x: y }
+}
+",
+        )
+    }
+
+    #[test]
+    fn associated_record_field_shorthand() {
+        check_assist(
+            generate_enum_variant,
+            r"
+enum Foo {}
+fn main() {
+    let x = true;
+    Foo::$0Bar { x }
+}
+",
+            r"
+enum Foo {
+    Bar { x: bool },
+}
+fn main() {
+    let x = true;
+    Foo::Bar { x }
+}
+",
+        )
+    }
+
+    #[test]
+    fn associated_record_field_shorthand_unknown_type() {
+        check_assist(
+            generate_enum_variant,
+            r"
+enum Foo {}
+fn main() {
+    Foo::$0Bar { x }
+}
+",
+            r"
+enum Foo {
+    Bar { x: _ },
+}
+fn main() {
+    Foo::Bar { x }
+}
+",
+        )
+    }
+
+    #[test]
+    fn associated_record_field_multiple_fields() {
+        check_assist(
+            generate_enum_variant,
+            r"
+struct Struct {}
+enum Foo {}
+fn main() {
+    Foo::$0Bar { x, y: x, s: Struct {} }
+}
+",
+            r"
+struct Struct {}
+enum Foo {
+    Bar { x: _, y: _, s: Struct },
+}
+fn main() {
+    Foo::Bar { x, y: x, s: Struct {} }
+}
+",
+        )
+    }
+
+    #[test]
+    fn use_tree() {
+        check_assist(
+            generate_enum_variant,
+            r"
+//- /main.rs
+mod foo;
+use foo::Foo::Bar$0;
+
+//- /foo.rs
+enum Foo {}
+",
+            r"
+enum Foo {
+    Bar,
+}
+",
+        )
+    }
+
+    #[test]
+    fn not_applicable_for_path_type() {
+        check_assist_not_applicable(
+            generate_enum_variant,
+            r"
+enum Foo {}
+impl Foo::Bar$0 {}
+",
+        )
+    }
+
+    #[test]
+    fn path_pat() {
+        check_assist(
+            generate_enum_variant,
+            r"
+enum Foo {}
+fn foo(x: Foo) {
+    match x {
+        Foo::Bar$0 =>
+    }
+}
+",
+            r"
+enum Foo {
+    Bar,
+}
+fn foo(x: Foo) {
+    match x {
+        Foo::Bar =>
+    }
+}
+",
+        )
+    }
 }