Fix not applicable c-str and byte-str for raw_string

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

@@ -4,7 +4,7 @@ use syntax::{AstToken, TextRange, TextSize, ast, ast::IsString};
 
 use crate::{
     AssistContext, AssistId, Assists,
-    utils::{required_hashes, string_suffix},
+    utils::{required_hashes, string_prefix, string_suffix},
 };
 
 // Assist: make_raw_string
@@ -23,8 +23,7 @@ use crate::{
 // }
 // ```
 pub(crate) fn make_raw_string(acc: &mut Assists, ctx: &AssistContext<'_>) -> Option<()> {
-    // FIXME: This should support byte and c strings as well.
-    let token = ctx.find_token_at_offset::<ast::String>()?;
+    let token = ctx.find_token_at_offset::<ast::AnyString>()?;
     if token.is_raw() {
         return None;
     }
@@ -37,14 +36,15 @@ pub(crate) fn make_raw_string(acc: &mut Assists, ctx: &AssistContext<'_>) -> Opt
         |edit| {
             let hashes = "#".repeat(required_hashes(&value).max(1));
             let range = token.syntax().text_range();
+            let raw_prefix = token.raw_prefix();
             let suffix = string_suffix(token.text()).unwrap_or_default();
             let range = TextRange::new(range.start(), range.end() - TextSize::of(suffix));
             if matches!(value, Cow::Borrowed(_)) {
                 // Avoid replacing the whole string to better position the cursor.
-                edit.insert(range.start(), format!("r{hashes}"));
+                edit.insert(range.start(), format!("{raw_prefix}{hashes}"));
                 edit.insert(range.end(), hashes);
             } else {
-                edit.replace(range, format!("r{hashes}\"{value}\"{hashes}"));
+                edit.replace(range, format!("{raw_prefix}{hashes}\"{value}\"{hashes}"));
             }
         },
     )
@@ -66,7 +66,7 @@ pub(crate) fn make_raw_string(acc: &mut Assists, ctx: &AssistContext<'_>) -> Opt
 // }
 // ```
 pub(crate) fn make_usual_string(acc: &mut Assists, ctx: &AssistContext<'_>) -> Option<()> {
-    let token = ctx.find_token_at_offset::<ast::String>()?;
+    let token = ctx.find_token_at_offset::<ast::AnyString>()?;
     if !token.is_raw() {
         return None;
     }
@@ -80,18 +80,22 @@ pub(crate) fn make_usual_string(acc: &mut Assists, ctx: &AssistContext<'_>) -> O
             // parse inside string to escape `"`
             let escaped = value.escape_default().to_string();
             let suffix = string_suffix(token.text()).unwrap_or_default();
+            let prefix = string_prefix(token.text()).map_or("", |s| s.trim_end_matches('r'));
             if let Some(offsets) = token.quote_offsets()
                 && token.text()[offsets.contents - token.syntax().text_range().start()] == escaped
             {
+                let start_quote = offsets.quotes.0;
+                let start_quote =
+                    TextRange::new(start_quote.start() + TextSize::of(prefix), start_quote.end());
                 let end_quote = offsets.quotes.1;
                 let end_quote =
                     TextRange::new(end_quote.start(), end_quote.end() - TextSize::of(suffix));
-                edit.replace(offsets.quotes.0, "\"");
                 edit.replace(end_quote, "\"");
+                edit.replace(start_quote, "\"");
                 return;
             }
 
-            edit.replace(token.syntax().text_range(), format!("\"{escaped}\"{suffix}"));
+            edit.replace(token.syntax().text_range(), format!("{prefix}\"{escaped}\"{suffix}"));
         },
     )
 }
@@ -112,15 +116,15 @@ pub(crate) fn make_usual_string(acc: &mut Assists, ctx: &AssistContext<'_>) -> O
 // }
 // ```
 pub(crate) fn add_hash(acc: &mut Assists, ctx: &AssistContext<'_>) -> Option<()> {
-    let token = ctx.find_token_at_offset::<ast::String>()?;
+    let token = ctx.find_token_at_offset::<ast::AnyString>()?;
     if !token.is_raw() {
         return None;
     }
     let text_range = token.syntax().text_range();
     let target = text_range;
     acc.add(AssistId::refactor("add_hash"), "Add #", target, |edit| {
         let suffix = string_suffix(token.text()).unwrap_or_default();
-        edit.insert(text_range.start() + TextSize::of('r'), "#");
+        edit.insert(text_range.start() + TextSize::of(token.raw_prefix()), "#");
         edit.insert(text_range.end() - TextSize::of(suffix), "#");
     })
 }
@@ -141,17 +145,15 @@ pub(crate) fn add_hash(acc: &mut Assists, ctx: &AssistContext<'_>) -> Option<()>
 // }
 // ```
 pub(crate) fn remove_hash(acc: &mut Assists, ctx: &AssistContext<'_>) -> Option<()> {
-    let token = ctx.find_token_at_offset::<ast::String>()?;
+    let token = ctx.find_token_at_offset::<ast::AnyString>()?;
     if !token.is_raw() {
         return None;
     }
 
     let text = token.text();
-    if !text.starts_with("r#") && text.ends_with('#') {
-        return None;
-    }
 
-    let existing_hashes = text.chars().skip(1).take_while(|&it| it == '#').count();
+    let existing_hashes =
+        text.chars().skip(token.raw_prefix().len()).take_while(|&it| it == '#').count();
 
     let text_range = token.syntax().text_range();
     let internal_text = &text[token.text_range_between_quotes()? - text_range.start()];
@@ -163,7 +165,10 @@ pub(crate) fn remove_hash(acc: &mut Assists, ctx: &AssistContext<'_>) -> Option<
 
     acc.add(AssistId::refactor_rewrite("remove_hash"), "Remove #", text_range, |edit| {
         let suffix = string_suffix(text).unwrap_or_default();
-        edit.delete(TextRange::at(text_range.start() + TextSize::of('r'), TextSize::of('#')));
+        edit.delete(TextRange::at(
+            text_range.start() + TextSize::of(token.raw_prefix()),
+            TextSize::of('#'),
+        ));
         edit.delete(
             TextRange::new(text_range.end() - TextSize::of('#'), text_range.end())
                 - TextSize::of(suffix),
@@ -224,6 +229,42 @@ string"#;
         )
     }
 
+    #[test]
+    fn make_raw_byte_string_works() {
+        check_assist(
+            make_raw_string,
+            r#"
+fn f() {
+    let s = $0b"random\nstring";
+}
+"#,
+            r##"
+fn f() {
+    let s = br#"random
+string"#;
+}
+"##,
+        )
+    }
+
+    #[test]
+    fn make_raw_c_string_works() {
+        check_assist(
+            make_raw_string,
+            r#"
+fn f() {
+    let s = $0c"random\nstring";
+}
+"#,
+            r##"
+fn f() {
+    let s = cr#"random
+string"#;
+}
+"##,
+        )
+    }
+
     #[test]
     fn make_raw_string_hashes_inside_works() {
         check_assist(
@@ -348,6 +389,23 @@ string"###;
         )
     }
 
+    #[test]
+    fn add_hash_works_for_c_str() {
+        check_assist(
+            add_hash,
+            r#"
+            fn f() {
+                let s = $0cr"random string";
+            }
+            "#,
+            r##"
+            fn f() {
+                let s = cr#"random string"#;
+            }
+            "##,
+        )
+    }
+
     #[test]
     fn add_hash_has_suffix_works() {
         check_assist(
@@ -433,6 +491,15 @@ string"###;
         )
     }
 
+    #[test]
+    fn remove_hash_works_for_c_str() {
+        check_assist(
+            remove_hash,
+            r##"fn f() { let s = $0cr#"random string"#; }"##,
+            r#"fn f() { let s = cr"random string"; }"#,
+        )
+    }
+
     #[test]
     fn remove_hash_has_suffix_works() {
         check_assist(
@@ -529,6 +596,23 @@ string"###;
         )
     }
 
+    #[test]
+    fn make_usual_string_for_c_str() {
+        check_assist(
+            make_usual_string,
+            r##"
+            fn f() {
+                let s = $0cr#"random string"#;
+            }
+            "##,
+            r#"
+            fn f() {
+                let s = c"random string";
+            }
+            "#,
+        )
+    }
+
     #[test]
     fn make_usual_string_has_suffix_works() {
         check_assist(

crates/ide-assists/src/utils.rs

@@ -1057,6 +1057,21 @@ fn test_string_suffix() {
     assert_eq!(Some("i32"), string_suffix(r##"r#""#i32"##));
 }
 
+/// Calculate the string literal prefix length
+pub(crate) fn string_prefix(s: &str) -> Option<&str> {
+    s.split_once(['"', '\'', '#']).map(|(prefix, _)| prefix)
+}
+#[test]
+fn test_string_prefix() {
+    assert_eq!(Some(""), string_prefix(r#""abc""#));
+    assert_eq!(Some(""), string_prefix(r#""""#));
+    assert_eq!(Some(""), string_prefix(r#"""suffix"#));
+    assert_eq!(Some("c"), string_prefix(r#"c"""#));
+    assert_eq!(Some("r"), string_prefix(r#"r"""#));
+    assert_eq!(Some("cr"), string_prefix(r#"cr"""#));
+    assert_eq!(Some("r"), string_prefix(r##"r#""#"##));
+}
+
 /// Replaces the record expression, handling field shorthands including inside macros.
 pub(crate) fn replace_record_field_expr(
     ctx: &AssistContext<'_>,

crates/syntax/src/ast.rs

@@ -29,7 +29,9 @@ pub use self::{
         SlicePatComponents, StructKind, TypeBoundKind, TypeOrConstParam, VisibilityKind,
     },
     operators::{ArithOp, BinaryOp, CmpOp, LogicOp, Ordering, RangeOp, UnaryOp},
-    token_ext::{CommentKind, CommentPlacement, CommentShape, IsString, QuoteOffsets, Radix},
+    token_ext::{
+        AnyString, CommentKind, CommentPlacement, CommentShape, IsString, QuoteOffsets, Radix,
+    },
     traits::{
         AttrDocCommentIter, DocCommentIter, HasArgList, HasAttrs, HasDocComments, HasGenericArgs,
         HasGenericParams, HasLoopBody, HasModuleItem, HasName, HasTypeBounds, HasVisibility,