Repository: rust-lang/rust-analyzer (rust)

Issue:
Assist to turn discriminants of enums explicit
That is
```rs
enum Foo {
    V,
    V2,
    V3 = 10,
    V4,
}
```
becomes
```rs
enum Foo {
    V = 0,
    V2 = 1,
    V3 = 10,
    V4 = 11,
}
```

Pull request description:
Add explicit enum discriminant assist
Add assist for adding explicit discriminants to all variants of an enum.

Closes #17798.

Patch:
diff --git a/crates/hir-ty/src/consteval.rs b/crates/hir-ty/src/consteval.rs
index 8b6cde975f..e5c493ac56 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<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,
@@ -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) => {
@@ -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 0000000000..fafc3448a8
--- /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 c88cb3d5ea..b2ccd1fde8 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/src/inlay_hints.rs b/crates/ide/src/inlay_hints.rs
index 2e49af4914..be99510af2 100644
--- a/crates/ide/src/inlay_hints.rs
+++ b/crates/ide/src/inlay_hints.rs
@@ -15,7 +15,7 @@ use span::{Edition, EditionedFileId};
 use stdx::never;
 use syntax::{
     ast::{self, AstNode, HasGenericParams},
-    format_smolstr, match_ast, SmolStr, SyntaxNode, TextRange, TextSize, WalkEvent,
+    format_smolstr, match_ast, NodeOrToken, SmolStr, SyntaxNode, TextRange, TextSize, WalkEvent,
 };
 use text_edit::TextEdit;
 
@@ -95,27 +95,26 @@ pub(crate) fn inlay_hints(
     let famous_defs = FamousDefs(&sema, scope.krate());
 
     let ctx = &mut InlayHintCtx::default();
-    let mut hints = |event| {
-        if let Some(node) = handle_event(ctx, event) {
-            hints(&mut acc, ctx, &famous_defs, config, file_id, node);
-        }
-    };
-    let mut preorder = file.preorder();
-    while let Some(event) = preorder.next() {
+    let hints = |node| hints(&mut acc, ctx, &famous_defs, config, file_id, node);
+    match range_limit {
         // FIXME: This can miss some hints that require the parent of the range to calculate
-        if matches!((&event, range_limit), (WalkEvent::Enter(node), Some(range)) if range.intersect(node.text_range()).is_none())
-        {
-            preorder.skip_subtree();
-            continue;
-        }
-        hints(event);
-    }
+        Some(range) => match file.covering_element(range) {
+            NodeOrToken::Token(_) => return acc,
+            NodeOrToken::Node(n) => n
+                .preorder()
+                .filter(|event| matches!(event, WalkEvent::Enter(node) if range.intersect(node.text_range()).is_some()))
+                .for_each(hints),
+        },
+        None => file.preorder().for_each(hints),
+    };
+
     acc
 }
 
 #[derive(Default)]
 struct InlayHintCtx {
     lifetime_stacks: Vec<Vec<SmolStr>>,
+    is_param_list: bool,
 }
 
 pub(crate) fn inlay_hints_resolve(
@@ -139,52 +138,20 @@ pub(crate) fn inlay_hints_resolve(
     let mut acc = Vec::new();
 
     let ctx = &mut InlayHintCtx::default();
-    let mut hints = |event| {
-        if let Some(node) = handle_event(ctx, event) {
-            hints(&mut acc, ctx, &famous_defs, config, file_id, node);
-        }
+    let hints = |node| hints(&mut acc, ctx, &famous_defs, config, file_id, node);
+
+    let mut res = file.clone();
+    let res = loop {
+        res = match res.child_or_token_at_range(resolve_range) {
+            Some(NodeOrToken::Node(n)) if n.text_range() == resolve_range => break n,
+            Some(NodeOrToken::Node(n)) => n,
+            _ => break res,
+        };
     };
-
-    let mut preorder = file.preorder();
-    while let Some(event) = preorder.next() {
-        // FIXME: This can miss some hints that require the parent of the range to calculate
-        if matches!(&event, WalkEvent::Enter(node) if resolve_range.intersect(node.text_range()).is_none())
-        {
-            preorder.skip_subtree();
-            continue;
-        }
-        hints(event);
-    }
+    res.preorder().for_each(hints);
     acc.into_iter().find(|hint| hasher(hint) == hash)
 }
 
-fn handle_event(ctx: &mut InlayHintCtx, node: WalkEvent<SyntaxNode>) -> Option<SyntaxNode> {
-    match node {
-        WalkEvent::Enter(node) => {
-            if let Some(node) = ast::AnyHasGenericParams::cast(node.clone()) {
-                let params = node
-                    .generic_param_list()
-                    .map(|it| {
-                        it.lifetime_params()
-                            .filter_map(|it| {
-                                it.lifetime().map(|it| format_smolstr!("{}", &it.text()[1..]))
-                            })
-                            .collect()
-                    })
-                    .unwrap_or_default();
-                ctx.lifetime_stacks.push(params);
-            }
-            Some(node)
-        }
-        WalkEvent::Leave(n) => {
-            if ast::AnyHasGenericParams::can_cast(n.kind()) {
-                ctx.lifetime_stacks.pop();
-            }
-            None
-        }
-    }
-}
-
 // FIXME: At some point when our hir infra is fleshed out enough we should flip this and traverse the
 // HIR instead of the syntax tree.
 fn hints(
@@ -193,8 +160,35 @@ fn hints(
     famous_defs @ FamousDefs(sema, _): &FamousDefs<'_, '_>,
     config: &InlayHintsConfig,
     file_id: EditionedFileId,
-    node: SyntaxNode,
+    node: WalkEvent<SyntaxNode>,
 ) {
+    let node = match node {
+        WalkEvent::Enter(node) => node,
+        WalkEvent::Leave(n) => {
+            if ast::AnyHasGenericParams::can_cast(n.kind()) {
+                ctx.lifetime_stacks.pop();
+                // pop
+            }
+            if ast::ParamList::can_cast(n.kind()) {
+                ctx.is_param_list = false;
+                // pop
+            }
+            return;
+        }
+    };
+
+    if let Some(node) = ast::AnyHasGenericParams::cast(node.clone()) {
+        let params = node
+            .generic_param_list()
+            .map(|it| {
+                it.lifetime_params()
+                    .filter_map(|it| it.lifetime().map(|it| format_smolstr!("{}", &it.text()[1..])))
+                    .collect()
+            })
+            .unwrap_or_default();
+        ctx.lifetime_stacks.push(params);
+    }
+
     closing_brace::hints(hints, sema, config, file_id, node.clone());
     if let Some(any_has_generic_args) = ast::AnyHasGenericArgs::cast(node.clone()) {
         generic_param::hints(hints, sema, config, any_has_generic_args);
@@ -248,6 +242,10 @@ fn hints(
                 ast::Type::PathType(path) => lifetime::fn_path_hints(hints, ctx, famous_defs, config, file_id, path),
                 _ => Some(()),
             },
+            ast::ParamList(_) => {
+                ctx.is_param_list = true;
+                Some(())
+            },
             _ => Some(()),
         }
     };
diff --git a/crates/ide/src/inlay_hints/lifetime.rs b/crates/ide/src/inlay_hints/lifetime.rs
index 653e3a6ef1..de46367067 100644
--- a/crates/ide/src/inlay_hints/lifetime.rs
+++ b/crates/ide/src/inlay_hints/lifetime.rs
@@ -532,7 +532,7 @@ fn fn_ptr(a: fn(&()) -> &fn(&()) -> &()) {}
                        //^^ for<'1>
                           //^'1
                                   //^'1
-fn fn_ptr2(a: for<'a> fn(&()) -> &()) {}
+fn fn_ptr2(a: for<'a> fn(&()) -> &())) {}
                //^'0, $
                        //^'0
                                //^'0


Test patch:
diff --git a/crates/ide-assists/src/tests/generated.rs b/crates/ide-assists/src/tests/generated.rs
index dce7bbf342..48e12a8107 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(