Add pretty number for add_explicit_enum_discriminant

crates/ide-assists/src/handlers/add_explicit_enum_discriminant.rs

@@ -1,8 +1,11 @@
 use hir::Semantics;
 use ide_db::{RootDatabase, assists::AssistId, source_change::SourceChangeBuilder};
-use syntax::{AstNode, ast};
+use syntax::{
+    AstNode,
+    ast::{self, Radix},
+};
 
-use crate::{AssistContext, Assists};
+use crate::{AssistContext, Assists, utils::add_group_separators};
 
 // Assist: add_explicit_enum_discriminant
 //
@@ -53,8 +56,9 @@ pub(crate) fn add_explicit_enum_discriminant(
         "Add explicit enum discriminants",
         enum_node.syntax().text_range(),
         |builder| {
+            let mut radix = Radix::Decimal;
             for variant_node in variant_list.variants() {
-                add_variant_discriminant(&ctx.sema, builder, &variant_node);
+                add_variant_discriminant(&ctx.sema, builder, &variant_node, &mut radix);
             }
         },
     );
@@ -66,8 +70,10 @@ fn add_variant_discriminant(
     sema: &Semantics<'_, RootDatabase>,
     builder: &mut SourceChangeBuilder,
     variant_node: &ast::Variant,
+    radix: &mut Radix,
 ) {
-    if variant_node.expr().is_some() {
+    if let Some(expr) = variant_node.expr() {
+        *radix = expr_radix(&expr).unwrap_or(*radix);
         return;
     }
 
@@ -80,7 +86,24 @@ fn add_variant_discriminant(
 
     let variant_range = variant_node.syntax().text_range();
 
-    builder.insert(variant_range.end(), format!(" = {discriminant}"));
+    let (group_size, prefix, text) = match radix {
+        Radix::Binary => (4, "0b", format!("{discriminant:b}")),
+        Radix::Octal => (3, "0o", format!("{discriminant:o}")),
+        Radix::Decimal => (6, "", discriminant.to_string()),
+        Radix::Hexadecimal => (4, "0x", format!("{discriminant:x}")),
+    };
+    let pretty_num = add_group_separators(&text, group_size);
+    builder.insert(variant_range.end(), format!(" = {prefix}{pretty_num}"));
+}
+
+fn expr_radix(expr: &ast::Expr) -> Option<Radix> {
+    if let ast::Expr::Literal(lit) = expr
+        && let ast::LiteralKind::IntNumber(num) = lit.kind()
+    {
+        Some(num.radix())
+    } else {
+        None
+    }
 }
 
 #[cfg(test)]
@@ -172,9 +195,9 @@ enum TheEnum {
 #[repr(i64)]
 enum TheEnum {
     Foo = 1 << 63,
-    Bar = -9223372036854775807,
+    Bar = -9_223372_036854_775807,
     Baz = 0x7fff_ffff_ffff_fffe,
-    Quux = 9223372036854775807,
+    Quux = 0x7fff_ffff_ffff_ffff,
 }
 "#,
         );

crates/ide-assists/src/handlers/number_representation.rs

@@ -1,6 +1,6 @@
 use syntax::{AstToken, ast, ast::Radix};
 
-use crate::{AssistContext, AssistId, Assists, GroupLabel};
+use crate::{AssistContext, AssistId, Assists, GroupLabel, utils::add_group_separators};
 
 const MIN_NUMBER_OF_DIGITS_TO_FORMAT: usize = 5;
 
@@ -70,18 +70,6 @@ const fn group_size(r: Radix) -> usize {
     }
 }
 
-fn add_group_separators(s: &str, group_size: usize) -> String {
-    let mut chars = Vec::new();
-    for (i, ch) in s.chars().filter(|&ch| ch != '_').rev().enumerate() {
-        if i > 0 && i % group_size == 0 {
-            chars.push('_');
-        }
-        chars.push(ch);
-    }
-
-    chars.into_iter().rev().collect()
-}
-
 #[cfg(test)]
 mod tests {
     use crate::tests::{check_assist_by_label, check_assist_not_applicable, check_assist_target};

crates/ide-assists/src/utils.rs

@@ -1080,6 +1080,18 @@ fn test_string_prefix() {
     assert_eq!(Some("r"), string_prefix(r##"r#""#"##));
 }
 
+pub(crate) fn add_group_separators(s: &str, group_size: usize) -> String {
+    let mut chars = Vec::new();
+    for (i, ch) in s.chars().filter(|&ch| ch != '_').rev().enumerate() {
+        if i > 0 && i % group_size == 0 && ch != '-' {
+            chars.push('_');
+        }
+        chars.push(ch);
+    }
+
+    chars.into_iter().rev().collect()
+}
+
 /// Replaces the record expression, handling field shorthands including inside macros.
 pub(crate) fn replace_record_field_expr(
     ctx: &AssistContext<'_>,