Add explicit enum discriminant assist

Add assist for adding explicit discriminants to all variants of an enum.

3 files changed, 221 insertions, 5 deletions

diff --git a/src/tools/rust-analyzer/crates/hir-ty/src/consteval.rs b/src/tools/rust-analyzer/crates/hir-ty/src/consteval.rs
index 8b6cde975f6..3f60834e02d 100644
--- a/src/tools/rust-analyzer/crates/hir-ty/src/consteval.rs
+++ b/src/tools/rust-analyzer/crates/hir-ty/src/consteval.rs
@@ -11,7 +11,7 @@ use hir_def::{
     ConstBlockLoc, EnumVariantId, GeneralConstId, StaticId,
 };
 use hir_expand::Lookup;
-use stdx::never;
+use stdx::{never, IsNoneOr};
 use triomphe::Arc;
 
 use crate::{
@@ -169,15 +169,23 @@ pub fn usize_const(db: &dyn HirDatabase, value: Option<u128>, krate: CrateId) ->
 }
 
 pub fn try_const_usize(db: &dyn HirDatabase, c: &Const) -> Option<u128> {
+    try_const_usize_sign_extend(db, c, false)
+}
+
+pub fn try_const_usize_sign_extend(
+    db: &dyn HirDatabase,
+    c: &Const,
+    is_signed: bool,
+) -> Option<u128> {
     match &c.data(Interner).value {
         chalk_ir::ConstValue::BoundVar(_) => None,
         chalk_ir::ConstValue::InferenceVar(_) => None,
         chalk_ir::ConstValue::Placeholder(_) => None,
         chalk_ir::ConstValue::Concrete(c) => match &c.interned {
-            ConstScalar::Bytes(it, _) => Some(u128::from_le_bytes(pad16(it, false))),
+            ConstScalar::Bytes(it, _) => Some(u128::from_le_bytes(pad16(it, is_signed))),
             ConstScalar::UnevaluatedConst(c, subst) => {
                 let ec = db.const_eval(*c, subst.clone(), None).ok()?;
-                try_const_usize(db, &ec)
+                try_const_usize_sign_extend(db, &ec, is_signed)
             }
             _ => None,
         },
@@ -256,8 +264,8 @@ pub(crate) fn const_eval_discriminant_variant(
 ) -> Result<i128, ConstEvalError> {
     let def = variant_id.into();
     let body = db.body(def);
+    let loc = variant_id.lookup(db.upcast());
     if body.exprs[body.body_expr] == Expr::Missing {
-        let loc = variant_id.lookup(db.upcast());
         let prev_idx = loc.index.checked_sub(1);
         let value = match prev_idx {
             Some(prev_idx) => {
@@ -269,13 +277,17 @@ pub(crate) fn const_eval_discriminant_variant(
         };
         return Ok(value);
     }
+
+    let repr = db.enum_data(loc.parent).repr;
+    let is_signed = repr.and_then(|repr| repr.int).is_none_or(|int| int.is_signed());
+
     let mir_body = db.monomorphized_mir_body(
         def,
         Substitution::empty(Interner),
         db.trait_environment_for_body(def),
     )?;
     let c = interpret_mir(db, mir_body, false, None).0?;
-    let c = try_const_usize(db, &c).unwrap() as i128;
+    let c = try_const_usize_sign_extend(db, &c, is_signed).unwrap() as i128;
     Ok(c)
 }
 
diff --git a/src/tools/rust-analyzer/crates/ide-assists/src/handlers/explicit_enum_discriminant.rs b/src/tools/rust-analyzer/crates/ide-assists/src/handlers/explicit_enum_discriminant.rs
new file mode 100644
index 00000000000..b5b7da69f75
--- /dev/null
+++ b/src/tools/rust-analyzer/crates/ide-assists/src/handlers/explicit_enum_discriminant.rs
@@ -0,0 +1,202 @@
+use hir::Semantics;
+use ide_db::{
+    assists::{AssistId, AssistKind},
+    source_change::SourceChangeBuilder,
+    RootDatabase,
+};
+use syntax::{ast, AstNode};
+
+use crate::{AssistContext, Assists};
+
+// Assist: explicit_enum_discriminant
+//
+// Adds explicit discriminant to all enum variants.
+//
+// ```
+// enum TheEnum$0 {
+//     Foo,
+//     Bar,
+//     Baz = 42,
+//     Quux,
+// }
+// ```
+// ->
+// ```
+// enum TheEnum {
+//     Foo = 0,
+//     Bar = 1,
+//     Baz = 42,
+//     Quux = 43,
+// }
+// ```
+pub(crate) fn explicit_enum_discriminant(acc: &mut Assists, ctx: &AssistContext<'_>) -> Option<()> {
+    let enum_node = ctx.find_node_at_offset::<ast::Enum>()?;
+    let enum_def = ctx.sema.to_def(&enum_node)?;
+
+    let is_data_carrying = enum_def.is_data_carrying(ctx.db());
+    let has_primitive_repr = enum_def.repr(ctx.db()).and_then(|repr| repr.int).is_some();
+
+    // Data carrying enums without a primitive repr have no stable discriminants.
+    if is_data_carrying && !has_primitive_repr {
+        return None;
+    }
+
+    let variant_list = enum_node.variant_list()?;
+
+    // Don't offer the assist if the enum has no variants or if all variants already have an
+    // explicit discriminant.
+    if variant_list.variants().all(|variant_node| variant_node.expr().is_some()) {
+        return None;
+    }
+
+    acc.add(
+        AssistId("explicit_enum_discriminant", AssistKind::RefactorRewrite),
+        "Add explicit enum discriminants",
+        enum_node.syntax().text_range(),
+        |builder| {
+            for variant_node in variant_list.variants() {
+                add_variant_discriminant(&ctx.sema, builder, &variant_node);
+            }
+        },
+    );
+
+    Some(())
+}
+
+fn add_variant_discriminant(
+    sema: &Semantics<'_, RootDatabase>,
+    builder: &mut SourceChangeBuilder,
+    variant_node: &ast::Variant,
+) {
+    if variant_node.expr().is_some() {
+        return;
+    }
+
+    let Some(variant_def) = sema.to_def(variant_node) else {
+        return;
+    };
+    let Ok(discriminant) = variant_def.eval(sema.db) else {
+        return;
+    };
+
+    let variant_range = variant_node.syntax().text_range();
+
+    builder.insert(variant_range.end(), format!(" = {discriminant}"));
+}
+
+#[cfg(test)]
+mod tests {
+    use crate::tests::{check_assist, check_assist_not_applicable};
+
+    use super::explicit_enum_discriminant;
+
+    #[test]
+    fn non_primitive_repr_non_data_bearing_add_discriminant() {
+        check_assist(
+            explicit_enum_discriminant,
+            r#"
+enum TheEnum$0 {
+    Foo,
+    Bar,
+    Baz = 42,
+    Quux,
+}
+"#,
+            r#"
+enum TheEnum {
+    Foo = 0,
+    Bar = 1,
+    Baz = 42,
+    Quux = 43,
+}
+"#,
+        );
+    }
+
+    #[test]
+    fn primitive_repr_data_bearing_add_discriminant() {
+        check_assist(
+            explicit_enum_discriminant,
+            r#"
+#[repr(u8)]
+$0enum TheEnum {
+    Foo { x: u32 },
+    Bar,
+    Baz(String),
+    Quux,
+}
+"#,
+            r#"
+#[repr(u8)]
+enum TheEnum {
+    Foo { x: u32 } = 0,
+    Bar = 1,
+    Baz(String) = 2,
+    Quux = 3,
+}
+"#,
+        );
+    }
+
+    #[test]
+    fn non_primitive_repr_data_bearing_not_applicable() {
+        check_assist_not_applicable(
+            explicit_enum_discriminant,
+            r#"
+enum TheEnum$0 {
+    Foo,
+    Bar(u16),
+    Baz,
+}
+"#,
+        );
+    }
+
+    #[test]
+    fn primitive_repr_non_data_bearing_add_discriminant() {
+        check_assist(
+            explicit_enum_discriminant,
+            r#"
+#[repr(i64)]
+enum TheEnum {
+    Foo = 1 << 63,
+    Bar,
+    Baz$0 = 0x7fff_ffff_ffff_fffe,
+    Quux,
+}
+"#,
+            r#"
+#[repr(i64)]
+enum TheEnum {
+    Foo = 1 << 63,
+    Bar = -9223372036854775807,
+    Baz = 0x7fff_ffff_ffff_fffe,
+    Quux = 9223372036854775807,
+}
+"#,
+        );
+    }
+
+    #[test]
+    fn discriminants_already_explicit_not_applicable() {
+        check_assist_not_applicable(
+            explicit_enum_discriminant,
+            r#"
+enum TheEnum$0 {
+    Foo = 0,
+    Bar = 4,
+}
+"#,
+        );
+    }
+
+    #[test]
+    fn empty_enum_not_applicable() {
+        check_assist_not_applicable(
+            explicit_enum_discriminant,
+            r#"
+enum TheEnum$0 {}
+"#,
+        );
+    }
+}
diff --git a/src/tools/rust-analyzer/crates/ide-assists/src/lib.rs b/src/tools/rust-analyzer/crates/ide-assists/src/lib.rs
index c88cb3d5eaf..b2ccd1fde81 100644
--- a/src/tools/rust-analyzer/crates/ide-assists/src/lib.rs
+++ b/src/tools/rust-analyzer/crates/ide-assists/src/lib.rs
@@ -136,6 +136,7 @@ mod handlers {
     mod destructure_tuple_binding;
     mod desugar_doc_comment;
     mod expand_glob_import;
+    mod explicit_enum_discriminant;
     mod extract_expressions_from_format_string;
     mod extract_function;
     mod extract_module;
@@ -266,6 +267,7 @@ mod handlers {
             destructure_tuple_binding::destructure_tuple_binding,
             destructure_struct_binding::destructure_struct_binding,
             expand_glob_import::expand_glob_import,
+            explicit_enum_discriminant::explicit_enum_discriminant,
             extract_expressions_from_format_string::extract_expressions_from_format_string,
             extract_struct_from_enum_variant::extract_struct_from_enum_variant,
             extract_type_alias::extract_type_alias,