2 files changed, 224 insertions, 76 deletions

diff --git a/crates/assists/src/handlers/add_custom_impl.rs b/crates/assists/src/handlers/add_custom_impl.rs
index 8757fa33f..669dd9b21 100644
--- a/crates/assists/src/handlers/add_custom_impl.rs
+++ b/crates/assists/src/handlers/add_custom_impl.rs
@@ -1,13 +1,16 @@
+use ide_db::imports_locator;
 use itertools::Itertools;
 use syntax::{
-    ast::{self, AstNode},
+    ast::{self, make, AstNode},
     Direction, SmolStr,
     SyntaxKind::{IDENT, WHITESPACE},
     TextRange, TextSize,
 };
 
 use crate::{
-    assist_context::{AssistContext, Assists},
+    assist_config::SnippetCap,
+    assist_context::{AssistBuilder, AssistContext, Assists},
+    utils::mod_path_to_ast,
     AssistId, AssistKind,
 };
 
@@ -30,72 +33,116 @@ use crate::{
 // ```
 pub(crate) fn add_custom_impl(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
     let attr = ctx.find_node_at_offset::<ast::Attr>()?;
-    let input = attr.token_tree()?;
 
     let attr_name = attr
         .syntax()
         .descendants_with_tokens()
         .filter(|t| t.kind() == IDENT)
-        .find_map(|i| i.into_token())
-        .filter(|t| *t.text() == "derive")?
+        .find_map(syntax::NodeOrToken::into_token)
+        .filter(|t| t.text() == "derive")?
         .text()
         .clone();
 
     let trait_token =
         ctx.token_at_offset().find(|t| t.kind() == IDENT && *t.text() != attr_name)?;
+    let trait_path = make::path_unqualified(make::path_segment(make::name_ref(trait_token.text())));
 
     let annotated = attr.syntax().siblings(Direction::Next).find_map(ast::Name::cast)?;
     let annotated_name = annotated.syntax().text().to_string();
-    let start_offset = annotated.syntax().parent()?.text_range().end();
+    let insert_pos = annotated.syntax().parent()?.text_range().end();
+
+    let current_module = ctx.sema.scope(annotated.syntax()).module()?;
+    let current_crate = current_module.krate();
+
+    let found_traits = imports_locator::find_imports(&ctx.sema, current_crate, trait_token.text())
+        .into_iter()
+        .filter_map(|candidate: either::Either<hir::ModuleDef, hir::MacroDef>| match candidate {
+            either::Either::Left(hir::ModuleDef::Trait(trait_)) => Some(trait_),
+            _ => None,
+        })
+        .flat_map(|trait_| {
+            current_module
+                .find_use_path(ctx.sema.db, hir::ModuleDef::Trait(trait_))
+                .as_ref()
+                .map(mod_path_to_ast)
+                .zip(Some(trait_))
+        });
 
-    let label =
-        format!("Add custom impl `{}` for `{}`", trait_token.text().as_str(), annotated_name);
+    let mut no_traits_found = true;
+    for (trait_path, _trait) in found_traits.inspect(|_| no_traits_found = false) {
+        add_assist(acc, ctx.config.snippet_cap, &attr, &trait_path, &annotated_name, insert_pos)?;
+    }
+    if no_traits_found {
+        add_assist(acc, ctx.config.snippet_cap, &attr, &trait_path, &annotated_name, insert_pos)?;
+    }
+    Some(())
+}
 
+fn add_assist(
+    acc: &mut Assists,
+    snippet_cap: Option<SnippetCap>,
+    attr: &ast::Attr,
+    trait_path: &ast::Path,
+    annotated_name: &str,
+    insert_pos: TextSize,
+) -> Option<()> {
     let target = attr.syntax().text_range();
-    acc.add(AssistId("add_custom_impl", AssistKind::Refactor), label, target, |builder| {
-        let new_attr_input = input
-            .syntax()
-            .descendants_with_tokens()
-            .filter(|t| t.kind() == IDENT)
-            .filter_map(|t| t.into_token().map(|t| t.text().clone()))
-            .filter(|t| t != trait_token.text())
-            .collect::<Vec<SmolStr>>();
-        let has_more_derives = !new_attr_input.is_empty();
-
-        if has_more_derives {
-            let new_attr_input = format!("({})", new_attr_input.iter().format(", "));
-            builder.replace(input.syntax().text_range(), new_attr_input);
-        } else {
-            let attr_range = attr.syntax().text_range();
-            builder.delete(attr_range);
-
-            let line_break_range = attr
-                .syntax()
-                .next_sibling_or_token()
-                .filter(|t| t.kind() == WHITESPACE)
-                .map(|t| t.text_range())
-                .unwrap_or_else(|| TextRange::new(TextSize::from(0), TextSize::from(0)));
-            builder.delete(line_break_range);
-        }
+    let input = attr.token_tree()?;
+    let label = format!("Add custom impl `{}` for `{}`", trait_path, annotated_name);
+    let trait_name = trait_path.segment().and_then(|seg| seg.name_ref())?;
 
-        match ctx.config.snippet_cap {
+    acc.add(AssistId("add_custom_impl", AssistKind::Refactor), label, target, |builder| {
+        update_attribute(builder, &input, &trait_name, &attr);
+        match snippet_cap {
             Some(cap) => {
                 builder.insert_snippet(
                     cap,
-                    start_offset,
-                    format!("\n\nimpl {} for {} {{\n    $0\n}}", trait_token, annotated_name),
+                    insert_pos,
+                    format!("\n\nimpl {} for {} {{\n    $0\n}}", trait_path, annotated_name),
                 );
             }
             None => {
                 builder.insert(
-                    start_offset,
-                    format!("\n\nimpl {} for {} {{\n\n}}", trait_token, annotated_name),
+                    insert_pos,
+                    format!("\n\nimpl {} for {} {{\n\n}}", trait_path, annotated_name),
                 );
             }
         }
     })
 }
 
+fn update_attribute(
+    builder: &mut AssistBuilder,
+    input: &ast::TokenTree,
+    trait_name: &ast::NameRef,
+    attr: &ast::Attr,
+) {
+    let new_attr_input = input
+        .syntax()
+        .descendants_with_tokens()
+        .filter(|t| t.kind() == IDENT)
+        .filter_map(|t| t.into_token().map(|t| t.text().clone()))
+        .filter(|t| t != trait_name.text())
+        .collect::<Vec<SmolStr>>();
+    let has_more_derives = !new_attr_input.is_empty();
+
+    if has_more_derives {
+        let new_attr_input = format!("({})", new_attr_input.iter().format(", "));
+        builder.replace(input.syntax().text_range(), new_attr_input);
+    } else {
+        let attr_range = attr.syntax().text_range();
+        builder.delete(attr_range);
+
+        let line_break_range = attr
+            .syntax()
+            .next_sibling_or_token()
+            .filter(|t| t.kind() == WHITESPACE)
+            .map(|t| t.text_range())
+            .unwrap_or_else(|| TextRange::new(TextSize::from(0), TextSize::from(0)));
+        builder.delete(line_break_range);
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use crate::tests::{check_assist, check_assist_not_applicable};
@@ -103,6 +150,35 @@ mod tests {
     use super::*;
 
     #[test]
+    fn add_custom_impl_qualified() {
+        check_assist(
+            add_custom_impl,
+            "
+mod fmt {
+    pub trait Debug {}
+}
+
+#[derive(Debu<|>g)]
+struct Foo {
+    bar: String,
+}
+",
+            "
+mod fmt {
+    pub trait Debug {}
+}
+
+struct Foo {
+    bar: String,
+}
+
+impl fmt::Debug for Foo {
+    $0
+}
+",
+        )
+    }
+    #[test]
     fn add_custom_impl_for_unique_input() {
         check_assist(
             add_custom_impl,
diff --git a/crates/assists/src/handlers/extract_struct_from_enum_variant.rs b/crates/assists/src/handlers/extract_struct_from_enum_variant.rs
index dddab255e..14209b771 100644
--- a/crates/assists/src/handlers/extract_struct_from_enum_variant.rs
+++ b/crates/assists/src/handlers/extract_struct_from_enum_variant.rs
@@ -1,3 +1,6 @@
+use std::iter;
+
+use either::Either;
 use hir::{AsName, EnumVariant, Module, ModuleDef, Name};
 use ide_db::{defs::Definition, search::Reference, RootDatabase};
 use rustc_hash::{FxHashMap, FxHashSet};
@@ -31,40 +34,32 @@ pub(crate) fn extract_struct_from_enum_variant(
     ctx: &AssistContext,
 ) -> Option<()> {
     let variant = ctx.find_node_at_offset::<ast::Variant>()?;
-    let field_list = match variant.kind() {
-        ast::StructKind::Tuple(field_list) => field_list,
-        _ => return None,
-    };
-
-    // skip 1-tuple variants
-    if field_list.fields().count() == 1 {
-        return None;
-    }
+    let field_list = extract_field_list_if_applicable(&variant)?;
 
     let variant_name = variant.name()?;
     let variant_hir = ctx.sema.to_def(&variant)?;
-    if existing_struct_def(ctx.db(), &variant_name, &variant_hir) {
+    if existing_definition(ctx.db(), &variant_name, &variant_hir) {
         return None;
     }
+
     let enum_ast = variant.parent_enum();
-    let visibility = enum_ast.visibility();
     let enum_hir = ctx.sema.to_def(&enum_ast)?;
-    let variant_hir_name = variant_hir.name(ctx.db());
-    let enum_module_def = ModuleDef::from(enum_hir);
-    let current_module = enum_hir.module(ctx.db());
     let target = variant.syntax().text_range();
     acc.add(
         AssistId("extract_struct_from_enum_variant", AssistKind::RefactorRewrite),
         "Extract struct from enum variant",
         target,
         |builder| {
-            let definition = Definition::ModuleDef(ModuleDef::EnumVariant(variant_hir));
-            let res = definition.usages(&ctx.sema).all();
+            let variant_hir_name = variant_hir.name(ctx.db());
+            let enum_module_def = ModuleDef::from(enum_hir);
+            let usages =
+                Definition::ModuleDef(ModuleDef::EnumVariant(variant_hir)).usages(&ctx.sema).all();
 
             let mut visited_modules_set = FxHashSet::default();
+            let current_module = enum_hir.module(ctx.db());
             visited_modules_set.insert(current_module);
             let mut rewriters = FxHashMap::default();
-            for reference in res {
+            for reference in usages {
                 let rewriter = rewriters
                     .entry(reference.file_range.file_id)
                     .or_insert_with(SyntaxRewriter::default);
@@ -86,26 +81,49 @@ pub(crate) fn extract_struct_from_enum_variant(
                 builder.rewrite(rewriter);
             }
             builder.edit_file(ctx.frange.file_id);
-            update_variant(&mut rewriter, &variant_name, &field_list);
+            update_variant(&mut rewriter, &variant);
             extract_struct_def(
                 &mut rewriter,
                 &enum_ast,
                 variant_name.clone(),
                 &field_list,
                 &variant.parent_enum().syntax().clone().into(),
-                visibility,
+                enum_ast.visibility(),
             );
             builder.rewrite(rewriter);
         },
     )
 }
 
-fn existing_struct_def(db: &RootDatabase, variant_name: &ast::Name, variant: &EnumVariant) -> bool {
+fn extract_field_list_if_applicable(
+    variant: &ast::Variant,
+) -> Option<Either<ast::RecordFieldList, ast::TupleFieldList>> {
+    match variant.kind() {
+        ast::StructKind::Record(field_list) if field_list.fields().next().is_some() => {
+            Some(Either::Left(field_list))
+        }
+        ast::StructKind::Tuple(field_list) if field_list.fields().count() > 1 => {
+            Some(Either::Right(field_list))
+        }
+        _ => None,
+    }
+}
+
+fn existing_definition(db: &RootDatabase, variant_name: &ast::Name, variant: &EnumVariant) -> bool {
     variant
         .parent_enum(db)
         .module(db)
         .scope(db, None)
         .into_iter()
+        .filter(|(_, def)| match def {
+            // only check type-namespace
+            hir::ScopeDef::ModuleDef(def) => matches!(def,
+                ModuleDef::Module(_) | ModuleDef::Adt(_) |
+                ModuleDef::EnumVariant(_) | ModuleDef::Trait(_) |
+                ModuleDef::TypeAlias(_) | ModuleDef::BuiltinType(_)
+            ),
+            _ => false,
+        })
         .any(|(name, _)| name == variant_name.as_name())
 }
 
@@ -133,19 +151,29 @@ fn extract_struct_def(
     rewriter: &mut SyntaxRewriter,
     enum_: &ast::Enum,
     variant_name: ast::Name,
-    variant_list: &ast::TupleFieldList,
+    field_list: &Either<ast::RecordFieldList, ast::TupleFieldList>,
     start_offset: &SyntaxElement,
     visibility: Option<ast::Visibility>,
 ) -> Option<()> {
-    let variant_list = make::tuple_field_list(
-        variant_list
-            .fields()
-            .flat_map(|field| Some(make::tuple_field(Some(make::visibility_pub()), field.ty()?))),
-    );
+    let pub_vis = Some(make::visibility_pub());
+    let field_list = match field_list {
+        Either::Left(field_list) => {
+            make::record_field_list(field_list.fields().flat_map(|field| {
+                Some(make::record_field(pub_vis.clone(), field.name()?, field.ty()?))
+            }))
+            .into()
+        }
+        Either::Right(field_list) => make::tuple_field_list(
+            field_list
+                .fields()
+                .flat_map(|field| Some(make::tuple_field(pub_vis.clone(), field.ty()?))),
+        )
+        .into(),
+    };
 
     rewriter.insert_before(
         start_offset,
-        make::struct_(visibility, variant_name, None, variant_list.into()).syntax(),
+        make::struct_(visibility, variant_name, None, field_list).syntax(),
     );
     rewriter.insert_before(start_offset, &make::tokens::blank_line());
 
@@ -156,15 +184,14 @@ fn extract_struct_def(
     Some(())
 }
 
-fn update_variant(
-    rewriter: &mut SyntaxRewriter,
-    variant_name: &ast::Name,
-    field_list: &ast::TupleFieldList,
-) -> Option<()> {
-    let (l, r): (SyntaxElement, SyntaxElement) =
-        (field_list.l_paren_token()?.into(), field_list.r_paren_token()?.into());
-    let replacement = vec![l, variant_name.syntax().clone().into(), r];
-    rewriter.replace_with_many(field_list.syntax(), replacement);
+fn update_variant(rewriter: &mut SyntaxRewriter, variant: &ast::Variant) -> Option<()> {
+    let name = variant.name()?;
+    let tuple_field = make::tuple_field(None, make::ty(name.text()));
+    let replacement = make::variant(
+        name,
+        Some(ast::FieldList::TupleFieldList(make::tuple_field_list(iter::once(tuple_field)))),
+    );
+    rewriter.replace(variant.syntax(), replacement.syntax());
     Some(())
 }
 
@@ -211,7 +238,7 @@ mod tests {
     use super::*;
 
     #[test]
-    fn test_extract_struct_several_fields() {
+    fn test_extract_struct_several_fields_tuple() {
         check_assist(
             extract_struct_from_enum_variant,
             "enum A { <|>One(u32, u32) }",
@@ -222,6 +249,41 @@ enum A { One(One) }"#,
     }
 
     #[test]
+    fn test_extract_struct_several_fields_named() {
+        check_assist(
+            extract_struct_from_enum_variant,
+            "enum A { <|>One { foo: u32, bar: u32 } }",
+            r#"struct One{ pub foo: u32, pub bar: u32 }
+
+enum A { One(One) }"#,
+        );
+    }
+
+    #[test]
+    fn test_extract_struct_one_field_named() {
+        check_assist(
+            extract_struct_from_enum_variant,
+            "enum A { <|>One { foo: u32 } }",
+            r#"struct One{ pub foo: u32 }
+
+enum A { One(One) }"#,
+        );
+    }
+
+    #[test]
+    fn test_extract_enum_variant_name_value_namespace() {
+        check_assist(
+            extract_struct_from_enum_variant,
+            r#"const One: () = ();
+enum A { <|>One(u32, u32) }"#,
+            r#"const One: () = ();
+struct One(pub u32, pub u32);
+
+enum A { One(One) }"#,
+        );
+    }
+
+    #[test]
     fn test_extract_struct_pub_visibility() {
         check_assist(
             extract_struct_from_enum_variant,
@@ -298,7 +360,7 @@ fn another_fn() {
     fn test_extract_enum_not_applicable_if_struct_exists() {
         check_not_applicable(
             r#"struct One;
-        enum A { <|>One(u8) }"#,
+        enum A { <|>One(u8, u32) }"#,
         );
     }
 
@@ -306,4 +368,14 @@ fn another_fn() {
     fn test_extract_not_applicable_one_field() {
         check_not_applicable(r"enum A { <|>One(u32) }");
     }
+
+    #[test]
+    fn test_extract_not_applicable_no_field_tuple() {
+        check_not_applicable(r"enum A { <|>None() }");
+    }
+
+    #[test]
+    fn test_extract_not_applicable_no_field_named() {
+        check_not_applicable(r"enum A { <|>None {} }");
+    }
 }