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}>");
  60             let impl_code = if let Some(name) = field_name {
  61                 format!(
  62                     r#"    fn from({name}: {field_type}) -> Self {{
  63         Self::{variant_name} {{ {name} }}
  64     }}"#
  65                 )
  66             } else {
  67                 format!(
  68                     r#"    fn from(v: {field_type}) -> Self {{
  69         Self::{variant_name}(v)
  70     }}"#
  71                 )
  72             };
  73             let from_impl = generate_trait_impl_text(&enum_, &from_trait, &impl_code);
  74             edit.insert(start_offset, from_impl);
  75         },
  76     )
  77 }
  78
  79 fn existing_from_impl(
  80     sema: &'_ hir::Semantics<'_, RootDatabase>,
  81     variant: &ast::Variant,
  82 ) -> Option<()> {
  83     let variant = sema.to_def(variant)?;
  84     let enum_ = variant.parent_enum(sema.db);
  85     let krate = enum_.module(sema.db).krate();
  86
  87     let from_trait = FamousDefs(sema, krate).core_convert_From()?;
  88
  89     let enum_type = enum_.ty(sema.db);
  90
  91     let wrapped_type = variant.fields(sema.db).get(0)?.ty(sema.db);
  92
  93     if enum_type.impls_trait(sema.db, from_trait, &[wrapped_type]) {
  94         Some(())
  95     } else {
  96         None
  97     }
  98 }
  99
 100 #[cfg(test)]
 101 mod tests {
 102     use crate::tests::{check_assist, check_assist_not_applicable};
 103
 104     use super::*;
 105
 106     #[test]
 107     fn test_generate_from_impl_for_enum() {
 108         check_assist(
 109             generate_from_impl_for_enum,
 110             r#"
 111 //- minicore: from
 112 enum A { $0One(u32) }
 113 "#,
 114             r#"
 115 enum A { One(u32) }
 116
 117 impl From<u32> for A {
 118     fn from(v: u32) -> Self {
 119         Self::One(v)
 120     }
 121 }
 122 "#,
 123         );
 124     }
 125
 126     #[test]
 127     fn test_generate_from_impl_for_enum_complicated_path() {
 128         check_assist(
 129             generate_from_impl_for_enum,
 130             r#"
 131 //- minicore: from
 132 enum A { $0One(foo::bar::baz::Boo) }
 133 "#,
 134             r#"
 135 enum A { One(foo::bar::baz::Boo) }
 136
 137 impl From<foo::bar::baz::Boo> for A {
 138     fn from(v: foo::bar::baz::Boo) -> Self {
 139         Self::One(v)
 140     }
 141 }
 142 "#,
 143         );
 144     }
 145
 146     #[test]
 147     fn test_add_from_impl_no_element() {
 148         check_assist_not_applicable(
 149             generate_from_impl_for_enum,
 150             r#"
 151 //- minicore: from
 152 enum A { $0One }
 153 "#,
 154         );
 155     }
 156
 157     #[test]
 158     fn test_add_from_impl_more_than_one_element_in_tuple() {
 159         check_assist_not_applicable(
 160             generate_from_impl_for_enum,
 161             r#"
 162 //- minicore: from
 163 enum A { $0One(u32, String) }
 164 "#,
 165         );
 166     }
 167
 168     #[test]
 169     fn test_add_from_impl_struct_variant() {
 170         check_assist(
 171             generate_from_impl_for_enum,
 172             r#"
 173 //- minicore: from
 174 enum A { $0One { x: u32 } }
 175 "#,
 176             r#"
 177 enum A { One { x: u32 } }
 178
 179 impl From<u32> for A {
 180     fn from(x: u32) -> Self {
 181         Self::One { x }
 182     }
 183 }
 184 "#,
 185         );
 186     }
 187
 188     #[test]
 189     fn test_add_from_impl_already_exists() {
 190         cov_mark::check!(test_add_from_impl_already_exists);
 191         check_assist_not_applicable(
 192             generate_from_impl_for_enum,
 193             r#"
 194 //- minicore: from
 195 enum A { $0One(u32), }
 196
 197 impl From<u32> for A {
 198     fn from(v: u32) -> Self {
 199         Self::One(v)
 200     }
 201 }
 202 "#,
 203         );
 204     }
 205
 206     #[test]
 207     fn test_add_from_impl_different_variant_impl_exists() {
 208         check_assist(
 209             generate_from_impl_for_enum,
 210             r#"
 211 //- minicore: from
 212 enum A { $0One(u32), Two(String), }
 213
 214 impl From<String> for A {
 215     fn from(v: String) -> Self {
 216         A::Two(v)
 217     }
 218 }
 219
 220 pub trait From<T> {
 221     fn from(T) -> Self;
 222 }
 223 "#,
 224             r#"
 225 enum A { One(u32), Two(String), }
 226
 227 impl From<u32> for A {
 228     fn from(v: u32) -> Self {
 229         Self::One(v)
 230     }
 231 }
 232
 233 impl From<String> for A {
 234     fn from(v: String) -> Self {
 235         A::Two(v)
 236     }
 237 }
 238
 239 pub trait From<T> {
 240     fn from(T) -> Self;
 241 }
 242 "#,
 243         );
 244     }
 245
 246     #[test]
 247     fn test_add_from_impl_static_str() {
 248         check_assist(
 249             generate_from_impl_for_enum,
 250             r#"
 251 //- minicore: from
 252 enum A { $0One(&'static str) }
 253 "#,
 254             r#"
 255 enum A { One(&'static str) }
 256
 257 impl From<&'static str> for A {
 258     fn from(v: &'static str) -> Self {
 259         Self::One(v)
 260     }
 261 }
 262 "#,
 263         );
 264     }
 265
 266     #[test]
 267     fn test_add_from_impl_generic_enum() {
 268         check_assist(
 269             generate_from_impl_for_enum,
 270             r#"
 271 //- minicore: from
 272 enum Generic<T, U: Clone> { $0One(T), Two(U) }
 273 "#,
 274             r#"
 275 enum Generic<T, U: Clone> { One(T), Two(U) }
 276
 277 impl<T, U: Clone> From<T> for Generic<T, U> {
 278     fn from(v: T) -> Self {
 279         Self::One(v)
 280     }
 281 }
 282 "#,
 283         );
 284     }
 285
 286     #[test]
 287     fn test_add_from_impl_with_lifetime() {
 288         check_assist(
 289             generate_from_impl_for_enum,
 290             r#"
 291 //- minicore: from
 292 enum Generic<'a> { $0One(&'a i32) }
 293 "#,
 294             r#"
 295 enum Generic<'a> { One(&'a i32) }
 296
 297 impl<'a> From<&'a i32> for Generic<'a> {
 298     fn from(v: &'a i32) -> Self {
 299         Self::One(v)
 300     }
 301 }
 302 "#,
 303         );
 304     }
 305 }