Merge pull request #19622 from A4-Tacks/raw-string-suffix

Fix ide-assists raw_string suffix fail

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

@@ -2,7 +2,10 @@ use std::borrow::Cow;
 
 use syntax::{AstToken, TextRange, TextSize, ast, ast::IsString};
 
-use crate::{AssistContext, AssistId, Assists, utils::required_hashes};
+use crate::{
+    AssistContext, AssistId, Assists,
+    utils::{required_hashes, string_suffix},
+};
 
 // Assist: make_raw_string
 //
@@ -33,12 +36,15 @@ pub(crate) fn make_raw_string(acc: &mut Assists, ctx: &AssistContext<'_>) -> Opt
         target,
         |edit| {
             let hashes = "#".repeat(required_hashes(&value).max(1));
+            let range = token.syntax().text_range();
+            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(token.syntax().text_range().start(), format!("r{hashes}"));
+                edit.insert(range.start(), format!("r{hashes}"));
-                edit.insert(token.syntax().text_range().end(), hashes);
+                edit.insert(range.end(), hashes);
             } else {
-                edit.replace(token.syntax().text_range(), format!("r{hashes}\"{value}\"{hashes}"));
+                edit.replace(range, format!("r{hashes}\"{value}\"{hashes}"));
             }
         },
     )
@@ -73,15 +79,19 @@ pub(crate) fn make_usual_string(acc: &mut Assists, ctx: &AssistContext<'_>) -> O
         |edit| {
             // parse inside string to escape `"`
             let escaped = value.escape_default().to_string();
+            let suffix = string_suffix(token.text()).unwrap_or_default();
             if let Some(offsets) = token.quote_offsets() {
                 if token.text()[offsets.contents - token.syntax().text_range().start()] == escaped {
+                    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(offsets.quotes.1, "\"");
+                    edit.replace(end_quote, "\"");
                     return;
                 }
             }
 
-            edit.replace(token.syntax().text_range(), format!("\"{escaped}\""));
+            edit.replace(token.syntax().text_range(), format!("\"{escaped}\"{suffix}"));
         },
     )
 }
@@ -109,8 +119,9 @@ pub(crate) fn add_hash(acc: &mut Assists, ctx: &AssistContext<'_>) -> Option<()>
     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.end(), "#");
+        edit.insert(text_range.end() - TextSize::of(suffix), "#");
     })
 }
 
@@ -151,8 +162,12 @@ 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::new(text_range.end() - TextSize::of('#'), text_range.end()));
+        edit.delete(
+            TextRange::new(text_range.end() - TextSize::of('#'), text_range.end())
+                - TextSize::of(suffix),
+        );
     })
 }
 
@@ -262,6 +277,23 @@ string"###;
         )
     }
 
+    #[test]
+    fn make_raw_string_has_suffix() {
+        check_assist(
+            make_raw_string,
+            r#"
+            fn f() {
+                let s = $0"random string"i32;
+            }
+            "#,
+            r##"
+            fn f() {
+                let s = r#"random string"#i32;
+            }
+            "##,
+        )
+    }
+
     #[test]
     fn make_raw_string_not_works_on_partial_string() {
         check_assist_not_applicable(
@@ -316,6 +348,23 @@ string"###;
         )
     }
 
+    #[test]
+    fn add_hash_has_suffix_works() {
+        check_assist(
+            add_hash,
+            r#"
+            fn f() {
+                let s = $0r"random string"i32;
+            }
+            "#,
+            r##"
+            fn f() {
+                let s = r#"random string"#i32;
+            }
+            "##,
+        )
+    }
+
     #[test]
     fn add_more_hash_works() {
         check_assist(
@@ -333,6 +382,23 @@ string"###;
         )
     }
 
+    #[test]
+    fn add_more_hash_has_suffix_works() {
+        check_assist(
+            add_hash,
+            r##"
+            fn f() {
+                let s = $0r#"random"string"#i32;
+            }
+            "##,
+            r###"
+            fn f() {
+                let s = r##"random"string"##i32;
+            }
+            "###,
+        )
+    }
+
     #[test]
     fn add_hash_not_works() {
         check_assist_not_applicable(
@@ -367,6 +433,15 @@ string"###;
         )
     }
 
+    #[test]
+    fn remove_hash_has_suffix_works() {
+        check_assist(
+            remove_hash,
+            r##"fn f() { let s = $0r#"random string"#i32; }"##,
+            r#"fn f() { let s = r"random string"i32; }"#,
+        )
+    }
+
     #[test]
     fn cant_remove_required_hash() {
         cov_mark::check!(cant_remove_required_hash);
@@ -397,6 +472,23 @@ string"###;
         )
     }
 
+    #[test]
+    fn remove_more_hash_has_suffix_works() {
+        check_assist(
+            remove_hash,
+            r###"
+            fn f() {
+                let s = $0r##"random string"##i32;
+            }
+            "###,
+            r##"
+            fn f() {
+                let s = r#"random string"#i32;
+            }
+            "##,
+        )
+    }
+
     #[test]
     fn remove_hash_does_not_work() {
         check_assist_not_applicable(remove_hash, r#"fn f() { let s = $0"random string"; }"#);
@@ -437,6 +529,23 @@ string"###;
         )
     }
 
+    #[test]
+    fn make_usual_string_has_suffix_works() {
+        check_assist(
+            make_usual_string,
+            r##"
+            fn f() {
+                let s = $0r#"random string"#i32;
+            }
+            "##,
+            r#"
+            fn f() {
+                let s = "random string"i32;
+            }
+            "#,
+        )
+    }
+
     #[test]
     fn make_usual_string_with_quote_works() {
         check_assist(
@@ -471,6 +580,23 @@ string"###;
         )
     }
 
+    #[test]
+    fn make_usual_string_more_hash_has_suffix_works() {
+        check_assist(
+            make_usual_string,
+            r###"
+            fn f() {
+                let s = $0r##"random string"##i32;
+            }
+            "###,
+            r##"
+            fn f() {
+                let s = "random string"i32;
+            }
+            "##,
+        )
+    }
+
     #[test]
     fn make_usual_string_not_works() {
         check_assist_not_applicable(

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

@@ -5,7 +5,7 @@ use syntax::{
     ast::IsString,
 };
 
-use crate::{AssistContext, AssistId, Assists};
+use crate::{AssistContext, AssistId, Assists, utils::string_suffix};
 
 // Assist: replace_string_with_char
 //
@@ -38,9 +38,11 @@ pub(crate) fn replace_string_with_char(acc: &mut Assists, ctx: &AssistContext<'_
         target,
         |edit| {
             let (left, right) = quote_offsets.quotes;
+            let suffix = TextSize::of(string_suffix(token.text()).unwrap_or_default());
+            let right = TextRange::new(right.start(), right.end() - suffix);
             edit.replace(left, '\'');
             edit.replace(right, '\'');
-            if value == "'" {
+            if token.text_without_quotes() == "'" {
                 edit.insert(left.end(), '\\');
             }
         },
@@ -71,12 +73,14 @@ pub(crate) fn replace_char_with_string(acc: &mut Assists, ctx: &AssistContext<'_
         "Replace char with string",
         target,
         |edit| {
-            if token.text() == "'\"'" {
+            let suffix = string_suffix(token.text()).unwrap_or_default();
-                edit.replace(token.text_range(), r#""\"""#);
+            if token.text().starts_with("'\"'") {
+                edit.replace(token.text_range(), format!(r#""\""{suffix}"#));
             } else {
                 let len = TextSize::of('\'');
+                let suffix = TextSize::of(suffix);
                 edit.replace(TextRange::at(target.start(), len), '"');
-                edit.replace(TextRange::at(target.end() - len, len), '"');
+                edit.replace(TextRange::at(target.end() - suffix - len, len), '"');
             }
         },
     )
@@ -105,6 +109,23 @@ fn f() {
         )
     }
 
+    #[test]
+    fn replace_string_with_char_has_suffix() {
+        check_assist(
+            replace_string_with_char,
+            r#"
+fn f() {
+    let s = "$0c"i32;
+}
+"#,
+            r##"
+fn f() {
+    let s = 'c'i32;
+}
+"##,
+        )
+    }
+
     #[test]
     fn replace_string_with_char_assist_with_multi_byte_char() {
         check_assist(
@@ -287,6 +308,40 @@ fn f() {
         )
     }
 
+    #[test]
+    fn replace_char_with_string_quote_has_suffix() {
+        check_assist(
+            replace_char_with_string,
+            r#"
+fn f() {
+    find($0'"'i32);
+}
+"#,
+            r#"
+fn f() {
+    find("\""i32);
+}
+"#,
+        )
+    }
+
+    #[test]
+    fn replace_char_with_string_escaped_quote_has_suffix() {
+        check_assist(
+            replace_char_with_string,
+            r#"
+fn f() {
+    find($0'\"'i32);
+}
+"#,
+            r#"
+fn f() {
+    find("\""i32);
+}
+"#,
+        )
+    }
+
     #[test]
     fn replace_string_with_char_quote() {
         check_assist(
@@ -300,6 +355,91 @@ fn f() {
 fn f() {
     find('\'');
 }
+"#,
+        )
+    }
+
+    #[test]
+    fn replace_string_with_escaped_char_quote() {
+        check_assist(
+            replace_string_with_char,
+            r#"
+fn f() {
+    find($0"\'");
+}
+"#,
+            r#"
+fn f() {
+    find('\'');
+}
+"#,
+        )
+    }
+
+    #[test]
+    fn replace_string_with_char_quote_has_suffix() {
+        check_assist(
+            replace_string_with_char,
+            r#"
+fn f() {
+    find($0"'"i32);
+}
+"#,
+            r#"
+fn f() {
+    find('\''i32);
+}
+"#,
+        )
+    }
+
+    #[test]
+    fn replace_string_with_escaped_char_quote_has_suffix() {
+        check_assist(
+            replace_string_with_char,
+            r#"
+fn f() {
+    find($0"\'"i32);
+}
+"#,
+            r#"
+fn f() {
+    find('\''i32);
+}
+"#,
+        )
+    }
+
+    #[test]
+    fn replace_raw_string_with_char_quote() {
+        check_assist(
+            replace_string_with_char,
+            r#"
+fn f() {
+    find($0r"'");
+}
+"#,
+            r#"
+fn f() {
+    find('\'');
+}
+"#,
+        )
+    }
+
+    #[test]
+    fn replace_string_with_code_escaped_char_quote() {
+        check_assist(
+            replace_string_with_char,
+            r#"
+fn f() {
+    find($0"\x27");
+}
+"#,
+            r#"
+fn f() {
+    find('\x27');
+}
 "#,
         )
     }

crates/ide-assists/src/utils.rs

@@ -1026,6 +1026,20 @@ fn test_required_hashes() {
     assert_eq!(5, required_hashes("#ab\"##\"####c"));
 }
 
+/// Calculate the string literal suffix length
+pub(crate) fn string_suffix(s: &str) -> Option<&str> {
+    s.rfind(['"', '\'', '#']).map(|i| &s[i + 1..])
+}
+#[test]
+fn test_string_suffix() {
+    assert_eq!(Some(""), string_suffix(r#""abc""#));
+    assert_eq!(Some(""), string_suffix(r#""""#));
+    assert_eq!(Some("a"), string_suffix(r#"""a"#));
+    assert_eq!(Some("i32"), string_suffix(r#"""i32"#));
+    assert_eq!(Some("i32"), string_suffix(r#"r""i32"#));
+    assert_eq!(Some("i32"), string_suffix(r##"r#""#i32"##));
+}
+
 /// Replaces the record expression, handling field shorthands including inside macros.
 pub(crate) fn replace_record_field_expr(
     ctx: &AssistContext<'_>,