src/tools/rust-analyzer/crates/ide-assists/src/handlers/generate_from_impl_for_enum.rs

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