Skip to content

Commit

Permalink
Auto merge of #17985 - riverbl:explicit-enum-discriminant, r=Veykril
Browse files Browse the repository at this point in the history
Add explicit enum discriminant assist

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

Closes #17798.
  • Loading branch information
bors committed Sep 1, 2024
2 parents 72fe274 + db1ccb8 commit bae4951
Show file tree
Hide file tree
Showing 4 changed files with 258 additions and 3 deletions.
30 changes: 27 additions & 3 deletions crates/hir-ty/src/consteval.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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::{
Expand Down Expand Up @@ -184,6 +184,22 @@ pub fn try_const_usize(db: &dyn HirDatabase, c: &Const) -> Option<u128> {
}
}

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,
Expand Down Expand Up @@ -256,8 +272,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) => {
Expand All @@ -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)
}

Expand Down
206 changes: 206 additions & 0 deletions crates/ide-assists/src/handlers/explicit_enum_discriminant.rs
Original file line number Diff line number Diff line change
@@ -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 {}
"#,
);
}
}
2 changes: 2 additions & 0 deletions crates/ide-assists/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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,
Expand Down
23 changes: 23 additions & 0 deletions crates/ide-assists/src/tests/generated.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down

0 comments on commit bae4951

Please sign in to comment.