Auto merge of #17985 - riverbl:explicit-enum-discriminant, r=Veykril
Add explicit enum discriminant assist
Add assist for adding explicit discriminants to all variants of an enum.
Closes #17798.
diff --git a/crates/hir-ty/src/consteval.rs b/crates/hir-ty/src/consteval.rs
index 8b6cde9..e5c493a 100644
--- a/crates/hir-ty/src/consteval.rs
+++ b/crates/hir-ty/src/consteval.rs
@@ -11,7 +11,7 @@
ConstBlockLoc, EnumVariantId, GeneralConstId, StaticId,
};
use hir_expand::Lookup;
-use stdx::never;
+use stdx::{never, IsNoneOr};
use triomphe::Arc;
use crate::{
@@ -184,6 +184,22 @@
}
}
+pub fn try_const_isize(db: &dyn HirDatabase, c: &Const) -> Option<i128> {
+ 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(i128::from_le_bytes(pad16(it, true))),
+ ConstScalar::UnevaluatedConst(c, subst) => {
+ let ec = db.const_eval(*c, subst.clone(), None).ok()?;
+ try_const_isize(db, &ec)
+ }
+ _ => None,
+ },
+ }
+}
+
pub(crate) fn const_eval_recover(
_: &dyn HirDatabase,
_: &Cycle,
@@ -256,8 +272,8 @@
) -> 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 +285,21 @@
};
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 = if is_signed {
+ try_const_isize(db, &c).unwrap()
+ } else {
+ try_const_usize(db, &c).unwrap() as i128
+ };
Ok(c)
}
diff --git a/crates/ide-assists/src/handlers/explicit_enum_discriminant.rs b/crates/ide-assists/src/handlers/explicit_enum_discriminant.rs
new file mode 100644
index 0000000..fafc344
--- /dev/null
+++ b/crates/ide-assists/src/handlers/explicit_enum_discriminant.rs
@@ -0,0 +1,206 @@
+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,
+ FooBar = -5,
+ FooBaz,
+}
+"#,
+ r#"
+enum TheEnum {
+ Foo = 0,
+ Bar = 1,
+ Baz = 42,
+ Quux = 43,
+ FooBar = -5,
+ FooBaz = -4,
+}
+"#,
+ );
+ }
+
+ #[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/crates/ide-assists/src/lib.rs b/crates/ide-assists/src/lib.rs
index c88cb3d..b2ccd1f 100644
--- a/crates/ide-assists/src/lib.rs
+++ b/crates/ide-assists/src/lib.rs
@@ -136,6 +136,7 @@
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 @@
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,
diff --git a/crates/ide-assists/src/tests/generated.rs b/crates/ide-assists/src/tests/generated.rs
index dce7bbf..48e12a8 100644
--- a/crates/ide-assists/src/tests/generated.rs
+++ b/crates/ide-assists/src/tests/generated.rs
@@ -910,6 +910,29 @@
}
#[test]
+fn doctest_explicit_enum_discriminant() {
+ check_doc_test(
+ "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 doctest_extract_expressions_from_format_string() {
check_doc_test(
"extract_expressions_from_format_string",
