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