ide_assists/handlers/
add_explicit_enum_discriminant.rs

1use hir::Semantics;
2use ide_db::{RootDatabase, assists::AssistId, source_change::SourceChangeBuilder};
3use syntax::{
4    AstNode,
5    ast::{self, Radix},
6};
7
8use crate::{AssistContext, Assists, utils::add_group_separators};
9
10// Assist: add_explicit_enum_discriminant
11//
12// Adds explicit discriminant to all enum variants.
13//
14// ```
15// enum TheEnum$0 {
16//     Foo,
17//     Bar,
18//     Baz = 42,
19//     Quux,
20// }
21// ```
22// ->
23// ```
24// enum TheEnum {
25//     Foo = 0,
26//     Bar = 1,
27//     Baz = 42,
28//     Quux = 43,
29// }
30// ```
31pub(crate) fn add_explicit_enum_discriminant(
32    acc: &mut Assists,
33    ctx: &AssistContext<'_>,
34) -> Option<()> {
35    let enum_node = ctx.find_node_at_offset::<ast::Enum>()?;
36    let enum_def = ctx.sema.to_def(&enum_node)?;
37
38    let is_data_carrying = enum_def.is_data_carrying(ctx.db());
39    let has_primitive_repr = enum_def.repr(ctx.db()).and_then(|repr| repr.int).is_some();
40
41    // Data carrying enums without a primitive repr have no stable discriminants.
42    if is_data_carrying && !has_primitive_repr {
43        return None;
44    }
45
46    let variant_list = enum_node.variant_list()?;
47
48    // Don't offer the assist if the enum has no variants or if all variants already have an
49    // explicit discriminant.
50    if variant_list.variants().all(|variant_node| variant_node.expr().is_some()) {
51        return None;
52    }
53
54    acc.add(
55        AssistId::refactor_rewrite("add_explicit_enum_discriminant"),
56        "Add explicit enum discriminants",
57        enum_node.syntax().text_range(),
58        |builder| {
59            let mut radix = Radix::Decimal;
60            for variant_node in variant_list.variants() {
61                add_variant_discriminant(&ctx.sema, builder, &variant_node, &mut radix);
62            }
63        },
64    );
65
66    Some(())
67}
68
69fn add_variant_discriminant(
70    sema: &Semantics<'_, RootDatabase>,
71    builder: &mut SourceChangeBuilder,
72    variant_node: &ast::Variant,
73    radix: &mut Radix,
74) {
75    if let Some(expr) = variant_node.expr() {
76        *radix = expr_radix(&expr).unwrap_or(*radix);
77        return;
78    }
79
80    let Some(variant_def) = sema.to_def(variant_node) else {
81        return;
82    };
83    let Ok(discriminant) = variant_def.eval(sema.db) else {
84        return;
85    };
86
87    let variant_range = variant_node.syntax().text_range();
88
89    let (group_size, prefix, text) = match radix {
90        Radix::Binary => (4, "0b", format!("{discriminant:b}")),
91        Radix::Octal => (3, "0o", format!("{discriminant:o}")),
92        Radix::Decimal => (6, "", discriminant.to_string()),
93        Radix::Hexadecimal => (4, "0x", format!("{discriminant:x}")),
94    };
95    let pretty_num = add_group_separators(&text, group_size);
96    builder.insert(variant_range.end(), format!(" = {prefix}{pretty_num}"));
97}
98
99fn expr_radix(expr: &ast::Expr) -> Option<Radix> {
100    if let ast::Expr::Literal(lit) = expr
101        && let ast::LiteralKind::IntNumber(num) = lit.kind()
102    {
103        Some(num.radix())
104    } else {
105        None
106    }
107}
108
109#[cfg(test)]
110mod tests {
111    use crate::tests::{check_assist, check_assist_not_applicable};
112
113    use super::add_explicit_enum_discriminant;
114
115    #[test]
116    fn non_primitive_repr_non_data_bearing_add_discriminant() {
117        check_assist(
118            add_explicit_enum_discriminant,
119            r#"
120enum TheEnum$0 {
121    Foo,
122    Bar,
123    Baz = 42,
124    Quux,
125    FooBar = -5,
126    FooBaz,
127}
128"#,
129            r#"
130enum TheEnum {
131    Foo = 0,
132    Bar = 1,
133    Baz = 42,
134    Quux = 43,
135    FooBar = -5,
136    FooBaz = -4,
137}
138"#,
139        );
140    }
141
142    #[test]
143    fn primitive_repr_data_bearing_add_discriminant() {
144        check_assist(
145            add_explicit_enum_discriminant,
146            r#"
147#[repr(u8)]
148$0enum TheEnum {
149    Foo { x: u32 },
150    Bar,
151    Baz(String),
152    Quux,
153}
154"#,
155            r#"
156#[repr(u8)]
157enum TheEnum {
158    Foo { x: u32 } = 0,
159    Bar = 1,
160    Baz(String) = 2,
161    Quux = 3,
162}
163"#,
164        );
165    }
166
167    #[test]
168    fn non_primitive_repr_data_bearing_not_applicable() {
169        check_assist_not_applicable(
170            add_explicit_enum_discriminant,
171            r#"
172enum TheEnum$0 {
173    Foo,
174    Bar(u16),
175    Baz,
176}
177"#,
178        );
179    }
180
181    #[test]
182    fn primitive_repr_non_data_bearing_add_discriminant() {
183        check_assist(
184            add_explicit_enum_discriminant,
185            r#"
186#[repr(i64)]
187enum TheEnum {
188    Foo = 1 << 63,
189    Bar,
190    Baz$0 = 0x7fff_ffff_ffff_fffe,
191    Quux,
192}
193"#,
194            r#"
195#[repr(i64)]
196enum TheEnum {
197    Foo = 1 << 63,
198    Bar = -9_223372_036854_775807,
199    Baz = 0x7fff_ffff_ffff_fffe,
200    Quux = 0x7fff_ffff_ffff_ffff,
201}
202"#,
203        );
204    }
205
206    #[test]
207    fn discriminants_already_explicit_not_applicable() {
208        check_assist_not_applicable(
209            add_explicit_enum_discriminant,
210            r#"
211enum TheEnum$0 {
212    Foo = 0,
213    Bar = 4,
214}
215"#,
216        );
217    }
218
219    #[test]
220    fn empty_enum_not_applicable() {
221        check_assist_not_applicable(
222            add_explicit_enum_discriminant,
223            r#"
224enum TheEnum$0 {}
225"#,
226        );
227    }
228}
ide_assists/handlers/add_explicit_enum_discriminant.rs

ide_assists/handlers/
add_explicit_enum_discriminant.rs
