crates/ide_assists/src/handlers/generate_default_from_enum_variant.rs

   1 use ide_db::helpers::FamousDefs;
   2 use ide_db::RootDatabase;
   3 use syntax::ast::{self, AstNode, NameOwner};
   4
   5 use crate::{AssistContext, AssistId, AssistKind, Assists};
   6
   7 // Assist: generate_default_from_enum_variant
   8 //
   9 // Adds a Default impl for an enum using a variant.
  10 //
  11 // ```
  12 // enum Version {
  13 //  Undefined,
  14 //  Minor$0,
  15 //  Major,
  16 // }
  17 // ```
  18 // ->
  19 // ```
  20 // enum Version {
  21 //  Undefined,
  22 //  Minor,
  23 //  Major,
  24 // }
  25 //
  26 // impl Default for Version {
  27 //     fn default() -> Self {
  28 //         Self::Minor
  29 //     }
  30 // }
  31 // ```
  32 pub(crate) fn generate_default_from_enum_variant(
  33     acc: &mut Assists,
  34     ctx: &AssistContext,
  35 ) -> Option<()> {
  36     let variant = ctx.find_node_at_offset::<ast::Variant>()?;
  37     let variant_name = variant.name()?;
  38     let enum_name = variant.parent_enum().name()?;
  39     if !matches!(variant.kind(), ast::StructKind::Unit) {
  40         cov_mark::hit!(test_gen_default_on_non_unit_variant_not_implemented);
  41         return None;
  42     }
  43
  44     if existing_default_impl(&ctx.sema, &variant).is_some() {
  45         cov_mark::hit!(test_gen_default_impl_already_exists);
  46         return None;
  47     }
  48
  49     let target = variant.syntax().text_range();
  50     acc.add(
  51         AssistId("generate_default_from_enum_variant", AssistKind::Generate),
  52         "Generate `Default` impl from this enum variant",
  53         target,
  54         |edit| {
  55             let start_offset = variant.parent_enum().syntax().text_range().end();
  56             let buf = format!(
  57                 r#"
  58
  59 impl Default for {0} {{
  60     fn default() -> Self {{
  61         Self::{1}
  62     }}
  63 }}"#,
  64                 enum_name, variant_name
  65             );
  66             edit.insert(start_offset, buf);
  67         },
  68     )
  69 }
  70
  71 fn existing_default_impl(
  72     sema: &'_ hir::Semantics<'_, RootDatabase>,
  73     variant: &ast::Variant,
  74 ) -> Option<()> {
  75     let variant = sema.to_def(variant)?;
  76     let enum_ = variant.parent_enum(sema.db);
  77     let krate = enum_.module(sema.db).krate();
  78
  79     let default_trait = FamousDefs(sema, Some(krate)).core_default_Default()?;
  80     let enum_type = enum_.ty(sema.db);
  81
  82     if enum_type.impls_trait(sema.db, default_trait, &[]) {
  83         Some(())
  84     } else {
  85         None
  86     }
  87 }
  88
  89 #[cfg(test)]
  90 mod tests {
  91     use crate::tests::{check_assist, check_assist_not_applicable};
  92
  93     use super::*;
  94
  95     fn check_not_applicable(ra_fixture: &str) {
  96         let fixture =
  97             format!("//- /main.rs crate:main deps:core\n{}\n{}", ra_fixture, FamousDefs::FIXTURE);
  98         check_assist_not_applicable(generate_default_from_enum_variant, &fixture)
  99     }
 100
 101     #[test]
 102     fn test_generate_default_from_variant() {
 103         check_assist(
 104             generate_default_from_enum_variant,
 105             r#"
 106 enum Variant {
 107     Undefined,
 108     Minor$0,
 109     Major,
 110 }"#,
 111             r#"enum Variant {
 112     Undefined,
 113     Minor,
 114     Major,
 115 }
 116
 117 impl Default for Variant {
 118     fn default() -> Self {
 119         Self::Minor
 120     }
 121 }"#,
 122         );
 123     }
 124
 125     #[test]
 126     fn test_generate_default_already_implemented() {
 127         cov_mark::check!(test_gen_default_impl_already_exists);
 128         check_not_applicable(
 129             r#"
 130 enum Variant {
 131     Undefined,
 132     Minor$0,
 133     Major,
 134 }
 135
 136 impl Default for Variant {
 137     fn default() -> Self {
 138         Self::Minor
 139     }
 140 }"#,
 141         );
 142     }
 143
 144     #[test]
 145     fn test_add_from_impl_no_element() {
 146         cov_mark::check!(test_gen_default_on_non_unit_variant_not_implemented);
 147         check_not_applicable(
 148             r#"
 149 enum Variant {
 150     Undefined,
 151     Minor(u32)$0,
 152     Major,
 153 }"#,
 154         );
 155     }
 156
 157     #[test]
 158     fn test_generate_default_from_variant_with_one_variant() {
 159         check_assist(
 160             generate_default_from_enum_variant,
 161             r#"enum Variant { Undefi$0ned }"#,
 162             r#"
 163 enum Variant { Undefined }
 164
 165 impl Default for Variant {
 166     fn default() -> Self {
 167         Self::Undefined
 168     }
 169 }"#,
 170         );
 171     }
 172 }