Add number representation assists

crates/ide_assists/src/handlers/number_representation.rs

@@ -0,0 +1,185 @@
+use syntax::{ast, ast::Radix, AstToken};
+
+use crate::{AssistContext, AssistId, AssistKind, Assists, GroupLabel};
+
+const MIN_NUMBER_OF_DIGITS_TO_FORMAT: usize = 5;
+
+// Assist: reformat_number_literal
+//
+// Adds or removes seprators from integer literal.
+//
+// ```
+// const _: i32 = 1012345$0;
+// ```
+// ->
+// ```
+// const _: i32 = 1_012_345;
+// ```
+pub(crate) fn reformat_number_literal(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
+    let literal = ctx.find_node_at_offset::<ast::Literal>()?;
+    let literal = match literal.kind() {
+        ast::LiteralKind::IntNumber(it) => it,
+        _ => return None,
+    };
+
+    let text = literal.text();
+    if text.contains('_') {
+        return remove_separators(acc, literal);
+    }
+
+    let value = literal.str_value();
+    if value.len() < MIN_NUMBER_OF_DIGITS_TO_FORMAT {
+        return None;
+    }
+
+    let radix = literal.radix();
+    let mut converted = literal.prefix().to_string();
+    converted.push_str(&add_group_separators(literal.str_value(), group_size(radix)));
+    if let Some(suffix) = literal.suffix() {
+        converted.push_str(suffix);
+    }
+
+    let group_id = GroupLabel("Reformat number literal".into());
+    let label = format!("Convert {} to {}", literal, converted);
+    let range = literal.syntax().text_range();
+    acc.add_group(
+        &group_id,
+        AssistId("reformat_number_literal", AssistKind::RefactorInline),
+        label,
+        range,
+        |builder| builder.replace(range, converted),
+    )
+}
+
+fn remove_separators(acc: &mut Assists, literal: ast::IntNumber) -> Option<()> {
+    let group_id = GroupLabel("Reformat number literal".into());
+    let range = literal.syntax().text_range();
+    acc.add_group(
+        &group_id,
+        AssistId("reformat_number_literal", AssistKind::RefactorInline),
+        "Remove digit seprators",
+        range,
+        |builder| builder.replace(range, literal.text().replace("_", "")),
+    )
+}
+
+const fn group_size(r: Radix) -> usize {
+    match r {
+        Radix::Binary => 4,
+        Radix::Octal => 3,
+        Radix::Decimal => 3,
+        Radix::Hexadecimal => 4,
+    }
+}
+
+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};
+
+    use super::*;
+
+    #[test]
+    fn group_separators() {
+        let cases = vec![
+            ("", 4, ""),
+            ("1", 4, "1"),
+            ("12", 4, "12"),
+            ("123", 4, "123"),
+            ("1234", 4, "1234"),
+            ("12345", 4, "1_2345"),
+            ("123456", 4, "12_3456"),
+            ("1234567", 4, "123_4567"),
+            ("12345678", 4, "1234_5678"),
+            ("123456789", 4, "1_2345_6789"),
+            ("1234567890", 4, "12_3456_7890"),
+            ("1_2_3_4_5_6_7_8_9_0_", 4, "12_3456_7890"),
+            ("1234567890", 3, "1_234_567_890"),
+            ("1234567890", 2, "12_34_56_78_90"),
+            ("1234567890", 1, "1_2_3_4_5_6_7_8_9_0"),
+        ];
+
+        for case in cases {
+            let (input, group_size, expected) = case;
+            assert_eq!(add_group_separators(input, group_size), expected)
+        }
+    }
+
+    #[test]
+    fn good_targets() {
+        let cases = vec![
+            ("const _: i32 = 0b11111$0", "0b11111"),
+            ("const _: i32 = 0o77777$0;", "0o77777"),
+            ("const _: i32 = 10000$0;", "10000"),
+            ("const _: i32 = 0xFFFFF$0;", "0xFFFFF"),
+            ("const _: i32 = 10000i32$0;", "10000i32"),
+            ("const _: i32 = 0b_10_0i32$0;", "0b_10_0i32"),
+        ];
+
+        for case in cases {
+            check_assist_target(reformat_number_literal, case.0, case.1);
+        }
+    }
+
+    #[test]
+    fn bad_targets() {
+        let cases = vec![
+            "const _: i32 = 0b111$0",
+            "const _: i32 = 0b1111$0",
+            "const _: i32 = 0o77$0;",
+            "const _: i32 = 0o777$0;",
+            "const _: i32 = 10$0;",
+            "const _: i32 = 999$0;",
+            "const _: i32 = 0xFF$0;",
+            "const _: i32 = 0xFFFF$0;",
+        ];
+
+        for case in cases {
+            check_assist_not_applicable(reformat_number_literal, case);
+        }
+    }
+
+    #[test]
+    fn labels() {
+        let cases = vec![
+            ("const _: i32 = 10000$0", "const _: i32 = 10_000", "Convert 10000 to 10_000"),
+            (
+                "const _: i32 = 0xFF0000$0;",
+                "const _: i32 = 0xFF_0000;",
+                "Convert 0xFF0000 to 0xFF_0000",
+            ),
+            (
+                "const _: i32 = 0b11111111$0;",
+                "const _: i32 = 0b1111_1111;",
+                "Convert 0b11111111 to 0b1111_1111",
+            ),
+            (
+                "const _: i32 = 0o377211$0;",
+                "const _: i32 = 0o377_211;",
+                "Convert 0o377211 to 0o377_211",
+            ),
+            (
+                "const _: i32 = 10000i32$0;",
+                "const _: i32 = 10_000i32;",
+                "Convert 10000i32 to 10_000i32",
+            ),
+            ("const _: i32 = 1_0_0_0_i32$0;", "const _: i32 = 1000i32;", "Remove digit seprators"),
+        ];
+
+        for case in cases {
+            let (before, after, label) = case;
+            check_assist_by_label(reformat_number_literal, before, after, label);
+        }
+    }
+}

crates/ide_assists/src/lib.rs

@@ -158,6 +158,7 @@ mod handlers {
     mod move_module_to_file;
     mod move_to_mod_rs;
     mod move_from_mod_rs;
+    mod number_representation;
     mod promote_local_to_const;
     mod pull_assignment_up;
     mod qualify_path;
@@ -241,6 +242,7 @@ mod handlers {
             move_module_to_file::move_module_to_file,
             move_to_mod_rs::move_to_mod_rs,
             move_from_mod_rs::move_from_mod_rs,
+            number_representation::reformat_number_literal,
             pull_assignment_up::pull_assignment_up,
             promote_local_to_const::promote_local_to_const,
             qualify_path::qualify_path,

crates/ide_assists/src/tests/generated.rs

@@ -1577,6 +1577,19 @@ pub mod std { pub mod collections { pub struct HashMap { } } }
     )
 }
 
+#[test]
+fn doctest_reformat_number_literal() {
+    check_doc_test(
+        "reformat_number_literal",
+        r#####"
+const _: i32 = 1012345$0;
+"#####,
+        r#####"
+const _: i32 = 1_012_345;
+"#####,
+    )
+}
+
 #[test]
 fn doctest_remove_dbg() {
     check_doc_test(

crates/syntax/src/ast/token_ext.rs

@@ -613,6 +613,8 @@ impl HasFormatSpecifier for ast::String {
     }
 }
 
+struct IntNumberParts<'a>(&'a str, &'a str, &'a str);
+
 impl ast::IntNumber {
     pub fn radix(&self) -> Radix {
         match self.text().get(..2).unwrap_or_default() {
@@ -623,41 +625,46 @@ impl ast::IntNumber {
         }
     }
 
-    pub fn value(&self) -> Option<u128> {
+    fn split_into_parts(&self) -> IntNumberParts {
-        let token = self.syntax();
-
-        let mut text = token.text();
-        if let Some(suffix) = self.suffix() {
-            text = &text[..text.len() - suffix.len()];
-        }
-
         let radix = self.radix();
-        text = &text[radix.prefix_len()..];
+        let (prefix, mut text) = self.text().split_at(radix.prefix_len());
 
-        let buf;
-        if text.contains('_') {
-            buf = text.replace('_', "");
-            text = buf.as_str();
-        };
-
-        let value = u128::from_str_radix(text, radix as u32).ok()?;
-        Some(value)
-    }
-
-    pub fn suffix(&self) -> Option<&str> {
-        let text = self.text();
-        let radix = self.radix();
-        let mut indices = text.char_indices();
-        if radix != Radix::Decimal {
-            indices.next()?;
-            indices.next()?;
-        }
         let is_suffix_start: fn(&(usize, char)) -> bool = match radix {
             Radix::Hexadecimal => |(_, c)| matches!(c, 'g'..='z' | 'G'..='Z'),
             _ => |(_, c)| c.is_ascii_alphabetic(),
         };
-        let (suffix_start, _) = indices.find(is_suffix_start)?;
+
-        Some(&text[suffix_start..])
+        let mut suffix = "";
+        if let Some((suffix_start, _)) = text.char_indices().find(is_suffix_start) {
+            let (text2, suffix2) = text.split_at(suffix_start);
+            text = text2;
+            suffix = suffix2;
+        };
+
+        IntNumberParts(prefix, text, suffix)
+    }
+
+    pub fn prefix(&self) -> &str {
+        self.split_into_parts().0
+    }
+
+    pub fn str_value(&self) -> &str {
+        self.split_into_parts().1
+    }
+
+    pub fn value(&self) -> Option<u128> {
+        let text = self.str_value().replace("_", "");
+        let value = u128::from_str_radix(&text, self.radix() as u32).ok()?;
+        Some(value)
+    }
+
+    pub fn suffix(&self) -> Option<&str> {
+        let suffix = self.split_into_parts().2;
+        if suffix.is_empty() {
+            None
+        } else {
+            Some(suffix)
+        }
     }
 }