diff --git a/crates/hir-ty/src/consteval.rs b/crates/hir-ty/src/consteval.rs index 8b6cde975f63..e5c493ac56e1 100644 --- a/crates/hir-ty/src/consteval.rs +++ b/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::{ @@ -184,6 +184,22 @@ pub fn try_const_usize(db: &dyn HirDatabase, c: &Const) -> Option { } } +pub fn try_const_isize(db: &dyn HirDatabase, c: &Const) -> Option { + 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 @@ pub(crate) fn const_eval_discriminant_variant( ) -> Result { 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 @@ 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 = 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 000000000000..fafc3448a87c --- /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::()?; + 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 c88cb3d5eaf0..b2ccd1fde815 100644 --- a/crates/ide-assists/src/lib.rs +++ b/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, diff --git a/crates/ide-assists/src/tests/generated.rs b/crates/ide-assists/src/tests/generated.rs index dce7bbf342f7..48e12a810735 100644 --- a/crates/ide-assists/src/tests/generated.rs +++ b/crates/ide-assists/src/tests/generated.rs @@ -909,6 +909,29 @@ fn qux(bar: Bar, baz: Baz) {} ) } +#[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(