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.const_arg().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.const_arg()
76        && let Some(expr) = expr.expr()
77    {
78        *radix = expr_radix(&expr).unwrap_or(*radix);
79        return;
80    }
81
82    let Some(variant_def) = sema.to_def(variant_node) else {
83        return;
84    };
85    let Ok(discriminant) = variant_def.eval(sema.db) else {
86        return;
87    };
88
89    let variant_range = variant_node.syntax().text_range();
90
91    let (group_size, prefix, text) = match radix {
92        Radix::Binary => (4, "0b", format!("{discriminant:b}")),
93        Radix::Octal => (3, "0o", format!("{discriminant:o}")),
94        Radix::Decimal => (6, "", discriminant.to_string()),
95        Radix::Hexadecimal => (4, "0x", format!("{discriminant:x}")),
96    };
97    let pretty_num = add_group_separators(&text, group_size);
98    builder.insert(variant_range.end(), format!(" = {prefix}{pretty_num}"));
99}
100
101fn expr_radix(expr: &ast::Expr) -> Option<Radix> {
102    if let ast::Expr::Literal(lit) = expr
103        && let ast::LiteralKind::IntNumber(num) = lit.kind()
104    {
105        Some(num.radix())
106    } else {
107        None
108    }
109}
110
111#[cfg(test)]
112mod tests {
113    use crate::tests::{check_assist, check_assist_not_applicable};
114
115    use super::add_explicit_enum_discriminant;
116
117    #[test]
118    fn non_primitive_repr_non_data_bearing_add_discriminant() {
119        check_assist(
120            add_explicit_enum_discriminant,
121            r#"
122enum TheEnum$0 {
123    Foo,
124    Bar,
125    Baz = 42,
126    Quux,
127    FooBar = -5,
128    FooBaz,
129}
130"#,
131            r#"
132enum TheEnum {
133    Foo = 0,
134    Bar = 1,
135    Baz = 42,
136    Quux = 43,
137    FooBar = -5,
138    FooBaz = -4,
139}
140"#,
141        );
142    }
143
144    #[test]
145    fn primitive_repr_data_bearing_add_discriminant() {
146        check_assist(
147            add_explicit_enum_discriminant,
148            r#"
149#[repr(u8)]
150$0enum TheEnum {
151    Foo { x: u32 },
152    Bar,
153    Baz(String),
154    Quux,
155}
156"#,
157            r#"
158#[repr(u8)]
159enum TheEnum {
160    Foo { x: u32 } = 0,
161    Bar = 1,
162    Baz(String) = 2,
163    Quux = 3,
164}
165"#,
166        );
167    }
168
169    #[test]
170    fn non_primitive_repr_data_bearing_not_applicable() {
171        check_assist_not_applicable(
172            add_explicit_enum_discriminant,
173            r#"
174enum TheEnum$0 {
175    Foo,
176    Bar(u16),
177    Baz,
178}
179"#,
180        );
181    }
182
183    #[test]
184    fn primitive_repr_non_data_bearing_add_discriminant() {
185        check_assist(
186            add_explicit_enum_discriminant,
187            r#"
188#[repr(i64)]
189enum TheEnum {
190    Foo = 1 << 63,
191    Bar,
192    Baz$0 = 0x7fff_ffff_ffff_fffe,
193    Quux,
194}
195"#,
196            r#"
197#[repr(i64)]
198enum TheEnum {
199    Foo = 1 << 63,
200    Bar = -9_223372_036854_775807,
201    Baz = 0x7fff_ffff_ffff_fffe,
202    Quux = 0x7fff_ffff_ffff_ffff,
203}
204"#,
205        );
206    }
207
208    #[test]
209    fn discriminants_already_explicit_not_applicable() {
210        check_assist_not_applicable(
211            add_explicit_enum_discriminant,
212            r#"
213enum TheEnum$0 {
214    Foo = 0,
215    Bar = 4,
216}
217"#,
218        );
219    }
220
221    #[test]
222    fn empty_enum_not_applicable() {
223        check_assist_not_applicable(
224            add_explicit_enum_discriminant,
225            r#"
226enum TheEnum$0 {}
227"#,
228        );
229    }
230}
ide_assists/handlers/add_explicit_enum_discriminant.rs

ide_assists/handlers/
add_explicit_enum_discriminant.rs
