crates/ide_assists/src/handlers/generate_from_impl_for_enum.rs

   1 use ide_db::helpers::FamousDefs;
   2 use ide_db::RootDatabase;
   3 use syntax::ast::{self, AstNode, NameOwner};
   4 use test_utils::mark;
   5
   6 use crate::{utils::generate_trait_impl_text, AssistContext, AssistId, AssistKind, Assists};
   7
   8 // Assist: generate_from_impl_for_enum
   9 //
  10 // Adds a From impl for an enum variant with one tuple field.
  11 //
  12 // ```
  13 // enum A { $0One(u32) }
  14 // ```
  15 // ->
  16 // ```
  17 // enum A { One(u32) }
  18 //
  19 // impl From<u32> for A {
  20 //     fn from(v: u32) -> Self {
  21 //         Self::One(v)
  22 //     }
  23 // }
  24 // ```
  25 pub(crate) fn generate_from_impl_for_enum(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
  26     let variant = ctx.find_node_at_offset::<ast::Variant>()?;
  27     let variant_name = variant.name()?;
  28     let enum_ = ast::Adt::Enum(variant.parent_enum());
  29     let (field_name, field_type) = match variant.kind() {
  30         ast::StructKind::Tuple(field_list) => {
  31             if field_list.fields().count() != 1 {
  32                 return None;
  33             }
  34             (None, field_list.fields().next()?.ty()?)
  35         }
  36         ast::StructKind::Record(field_list) => {
  37             if field_list.fields().count() != 1 {
  38                 return None;
  39             }
  40             let field = field_list.fields().next()?;
  41             (Some(field.name()?), field.ty()?)
  42         }
  43         ast::StructKind::Unit => return None,
  44     };
  45
  46     if existing_from_impl(&ctx.sema, &variant).is_some() {
  47         mark::hit!(test_add_from_impl_already_exists);
  48         return None;
  49     }
  50
  51     let target = variant.syntax().text_range();
  52     acc.add(
  53         AssistId("generate_from_impl_for_enum", AssistKind::Generate),
  54         "Generate `From` impl for this enum variant",
  55         target,
  56         |edit| {
  57             let start_offset = variant.parent_enum().syntax().text_range().end();
  58             let from_trait = format!("From<{}>", field_type.syntax());
  59             let impl_code = if let Some(name) = field_name {
  60                 format!(
  61                     r#"    fn from({0}: {1}) -> Self {{
  62         Self::{2} {{ {0} }}
  63     }}"#,
  64                     name.text(),
  65                     field_type.syntax(),
  66                     variant_name,
  67                 )
  68             } else {
  69                 format!(
  70                     r#"    fn from(v: {}) -> Self {{
  71         Self::{}(v)
  72     }}"#,
  73                     field_type.syntax(),
  74                     variant_name,
  75                 )
  76             };
  77             let from_impl = generate_trait_impl_text(&enum_, &from_trait, &impl_code);
  78             edit.insert(start_offset, from_impl);
  79         },
  80     )
  81 }
  82
  83 fn existing_from_impl(
  84     sema: &'_ hir::Semantics<'_, RootDatabase>,
  85     variant: &ast::Variant,
  86 ) -> Option<()> {
  87     let variant = sema.to_def(variant)?;
  88     let enum_ = variant.parent_enum(sema.db);
  89     let krate = enum_.module(sema.db).krate();
  90
  91     let from_trait = FamousDefs(sema, Some(krate)).core_convert_From()?;
  92
  93     let enum_type = enum_.ty(sema.db);
  94
  95     let wrapped_type = variant.fields(sema.db).get(0)?.signature_ty(sema.db);
  96
  97     if enum_type.impls_trait(sema.db, from_trait, &[wrapped_type]) {
  98         Some(())
  99     } else {
 100         None
 101     }
 102 }
 103
 104 #[cfg(test)]
 105 mod tests {
 106     use test_utils::mark;
 107
 108     use crate::tests::{check_assist, check_assist_not_applicable};
 109
 110     use super::*;
 111
 112     #[test]
 113     fn test_generate_from_impl_for_enum() {
 114         check_assist(
 115             generate_from_impl_for_enum,
 116             "enum A { $0One(u32) }",
 117             r#"enum A { One(u32) }
 118
 119 impl From<u32> for A {
 120     fn from(v: u32) -> Self {
 121         Self::One(v)
 122     }
 123 }"#,
 124         );
 125     }
 126
 127     #[test]
 128     fn test_generate_from_impl_for_enum_complicated_path() {
 129         check_assist(
 130             generate_from_impl_for_enum,
 131             r#"enum A { $0One(foo::bar::baz::Boo) }"#,
 132             r#"enum A { One(foo::bar::baz::Boo) }
 133
 134 impl From<foo::bar::baz::Boo> for A {
 135     fn from(v: foo::bar::baz::Boo) -> Self {
 136         Self::One(v)
 137     }
 138 }"#,
 139         );
 140     }
 141
 142     fn check_not_applicable(ra_fixture: &str) {
 143         let fixture =
 144             format!("//- /main.rs crate:main deps:core\n{}\n{}", ra_fixture, FamousDefs::FIXTURE);
 145         check_assist_not_applicable(generate_from_impl_for_enum, &fixture)
 146     }
 147
 148     #[test]
 149     fn test_add_from_impl_no_element() {
 150         check_not_applicable("enum A { $0One }");
 151     }
 152
 153     #[test]
 154     fn test_add_from_impl_more_than_one_element_in_tuple() {
 155         check_not_applicable("enum A { $0One(u32, String) }");
 156     }
 157
 158     #[test]
 159     fn test_add_from_impl_struct_variant() {
 160         check_assist(
 161             generate_from_impl_for_enum,
 162             "enum A { $0One { x: u32 } }",
 163             r#"enum A { One { x: u32 } }
 164
 165 impl From<u32> for A {
 166     fn from(x: u32) -> Self {
 167         Self::One { x }
 168     }
 169 }"#,
 170         );
 171     }
 172
 173     #[test]
 174     fn test_add_from_impl_already_exists() {
 175         mark::check!(test_add_from_impl_already_exists);
 176         check_not_applicable(
 177             r#"
 178 enum A { $0One(u32), }
 179
 180 impl From<u32> for A {
 181     fn from(v: u32) -> Self {
 182         Self::One(v)
 183     }
 184 }
 185 "#,
 186         );
 187     }
 188
 189     #[test]
 190     fn test_add_from_impl_different_variant_impl_exists() {
 191         check_assist(
 192             generate_from_impl_for_enum,
 193             r#"enum A { $0One(u32), Two(String), }
 194
 195 impl From<String> for A {
 196     fn from(v: String) -> Self {
 197         A::Two(v)
 198     }
 199 }
 200
 201 pub trait From<T> {
 202     fn from(T) -> Self;
 203 }"#,
 204             r#"enum A { One(u32), Two(String), }
 205
 206 impl From<u32> for A {
 207     fn from(v: u32) -> Self {
 208         Self::One(v)
 209     }
 210 }
 211
 212 impl From<String> for A {
 213     fn from(v: String) -> Self {
 214         A::Two(v)
 215     }
 216 }
 217
 218 pub trait From<T> {
 219     fn from(T) -> Self;
 220 }"#,
 221         );
 222     }
 223
 224     #[test]
 225     fn test_add_from_impl_static_str() {
 226         check_assist(
 227             generate_from_impl_for_enum,
 228             "enum A { $0One(&'static str) }",
 229             r#"enum A { One(&'static str) }
 230
 231 impl From<&'static str> for A {
 232     fn from(v: &'static str) -> Self {
 233         Self::One(v)
 234     }
 235 }"#,
 236         );
 237     }
 238
 239     #[test]
 240     fn test_add_from_impl_generic_enum() {
 241         check_assist(
 242             generate_from_impl_for_enum,
 243             "enum Generic<T, U: Clone> { $0One(T), Two(U) }",
 244             r#"enum Generic<T, U: Clone> { One(T), Two(U) }
 245
 246 impl<T, U: Clone> From<T> for Generic<T, U> {
 247     fn from(v: T) -> Self {
 248         Self::One(v)
 249     }
 250 }"#,
 251         );
 252     }
 253
 254     #[test]
 255     fn test_add_from_impl_with_lifetime() {
 256         check_assist(
 257             generate_from_impl_for_enum,
 258             "enum Generic<'a> { $0One(&'a i32) }",
 259             r#"enum Generic<'a> { One(&'a i32) }
 260
 261 impl<'a> From<&'a i32> for Generic<'a> {
 262     fn from(v: &'a i32) -> Self {
 263         Self::One(v)
 264     }
 265 }"#,
 266         );
 267     }
 268 }