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, HasName};
   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)?.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             r#"
 114 //- minicore: from
 115 enum A { $0One(u32) }
 116 "#,
 117             r#"
 118 enum A { One(u32) }
 119
 120 impl From<u32> for A {
 121     fn from(v: u32) -> Self {
 122         Self::One(v)
 123     }
 124 }
 125 "#,
 126         );
 127     }
 128
 129     #[test]
 130     fn test_generate_from_impl_for_enum_complicated_path() {
 131         check_assist(
 132             generate_from_impl_for_enum,
 133             r#"
 134 //- minicore: from
 135 enum A { $0One(foo::bar::baz::Boo) }
 136 "#,
 137             r#"
 138 enum A { One(foo::bar::baz::Boo) }
 139
 140 impl From<foo::bar::baz::Boo> for A {
 141     fn from(v: foo::bar::baz::Boo) -> Self {
 142         Self::One(v)
 143     }
 144 }
 145 "#,
 146         );
 147     }
 148
 149     #[test]
 150     fn test_add_from_impl_no_element() {
 151         check_assist_not_applicable(
 152             generate_from_impl_for_enum,
 153             r#"
 154 //- minicore: from
 155 enum A { $0One }
 156 "#,
 157         );
 158     }
 159
 160     #[test]
 161     fn test_add_from_impl_more_than_one_element_in_tuple() {
 162         check_assist_not_applicable(
 163             generate_from_impl_for_enum,
 164             r#"
 165 //- minicore: from
 166 enum A { $0One(u32, String) }
 167 "#,
 168         );
 169     }
 170
 171     #[test]
 172     fn test_add_from_impl_struct_variant() {
 173         check_assist(
 174             generate_from_impl_for_enum,
 175             r#"
 176 //- minicore: from
 177 enum A { $0One { x: u32 } }
 178 "#,
 179             r#"
 180 enum A { One { x: u32 } }
 181
 182 impl From<u32> for A {
 183     fn from(x: u32) -> Self {
 184         Self::One { x }
 185     }
 186 }
 187 "#,
 188         );
 189     }
 190
 191     #[test]
 192     fn test_add_from_impl_already_exists() {
 193         cov_mark::check!(test_add_from_impl_already_exists);
 194         check_assist_not_applicable(
 195             generate_from_impl_for_enum,
 196             r#"
 197 //- minicore: from
 198 enum A { $0One(u32), }
 199
 200 impl From<u32> for A {
 201     fn from(v: u32) -> Self {
 202         Self::One(v)
 203     }
 204 }
 205 "#,
 206         );
 207     }
 208
 209     #[test]
 210     fn test_add_from_impl_different_variant_impl_exists() {
 211         check_assist(
 212             generate_from_impl_for_enum,
 213             r#"
 214 //- minicore: from
 215 enum A { $0One(u32), Two(String), }
 216
 217 impl From<String> for A {
 218     fn from(v: String) -> Self {
 219         A::Two(v)
 220     }
 221 }
 222
 223 pub trait From<T> {
 224     fn from(T) -> Self;
 225 }
 226 "#,
 227             r#"
 228 enum A { One(u32), Two(String), }
 229
 230 impl From<u32> for A {
 231     fn from(v: u32) -> Self {
 232         Self::One(v)
 233     }
 234 }
 235
 236 impl From<String> for A {
 237     fn from(v: String) -> Self {
 238         A::Two(v)
 239     }
 240 }
 241
 242 pub trait From<T> {
 243     fn from(T) -> Self;
 244 }
 245 "#,
 246         );
 247     }
 248
 249     #[test]
 250     fn test_add_from_impl_static_str() {
 251         check_assist(
 252             generate_from_impl_for_enum,
 253             r#"
 254 //- minicore: from
 255 enum A { $0One(&'static str) }
 256 "#,
 257             r#"
 258 enum A { One(&'static str) }
 259
 260 impl From<&'static str> for A {
 261     fn from(v: &'static str) -> Self {
 262         Self::One(v)
 263     }
 264 }
 265 "#,
 266         );
 267     }
 268
 269     #[test]
 270     fn test_add_from_impl_generic_enum() {
 271         check_assist(
 272             generate_from_impl_for_enum,
 273             r#"
 274 //- minicore: from
 275 enum Generic<T, U: Clone> { $0One(T), Two(U) }
 276 "#,
 277             r#"
 278 enum Generic<T, U: Clone> { One(T), Two(U) }
 279
 280 impl<T, U: Clone> From<T> for Generic<T, U> {
 281     fn from(v: T) -> Self {
 282         Self::One(v)
 283     }
 284 }
 285 "#,
 286         );
 287     }
 288
 289     #[test]
 290     fn test_add_from_impl_with_lifetime() {
 291         check_assist(
 292             generate_from_impl_for_enum,
 293             r#"
 294 //- minicore: from
 295 enum Generic<'a> { $0One(&'a i32) }
 296 "#,
 297             r#"
 298 enum Generic<'a> { One(&'a i32) }
 299
 300 impl<'a> From<&'a i32> for Generic<'a> {
 301     fn from(v: &'a i32) -> Self {
 302         Self::One(v)
 303     }
 304 }
 305 "#,
 306         );
 307     }
 308 }