raw string assists work in macros

1 files changed, 58 insertions, 49 deletions

diff --git a/crates/ra_assists/src/assists/raw_string.rs b/crates/ra_assists/src/assists/raw_string.rs
index 2d2e31e51..9c996c902 100644
--- a/crates/ra_assists/src/assists/raw_string.rs
+++ b/crates/ra_assists/src/assists/raw_string.rs
@@ -1,17 +1,19 @@
 //! FIXME: write short doc here
 
 use hir::db::HirDatabase;
-use ra_syntax::{ast::AstNode, ast::Literal, TextRange, TextUnit};
+use ra_syntax::{
+    SyntaxKind::{RAW_STRING, STRING},
+    SyntaxToken, TextRange, TextUnit,
+};
 use rustc_lexer;
 
 use crate::{Assist, AssistCtx, AssistId};
 
 pub(crate) fn make_raw_string(mut ctx: AssistCtx<impl HirDatabase>) -> Option<Assist> {
-    let literal = ctx.node_at_offset::<Literal>()?;
-    if literal.token().kind() != ra_syntax::SyntaxKind::STRING {
+    let token = ctx.token_at_offset().right_biased()?;
+    if token.kind() != STRING {
         return None;
     }
-    let token = literal.token();
     let text = token.text().as_str();
     let usual_string_range = find_usual_string_range(text)?;
     let start_of_inside = usual_string_range.start().to_usize() + 1;
@@ -30,85 +32,52 @@ pub(crate) fn make_raw_string(mut ctx: AssistCtx<impl HirDatabase>) -> Option<As
         return None;
     }
     ctx.add_action(AssistId("make_raw_string"), "make raw string", |edit| {
-        edit.target(literal.syntax().text_range());
+        edit.target(token.text_range());
         let max_hash_streak = count_hashes(&unescaped);
         let mut hashes = String::with_capacity(max_hash_streak + 1);
         for _ in 0..hashes.capacity() {
             hashes.push('#');
         }
-        edit.replace(
-            literal.syntax().text_range(),
-            format!("r{}\"{}\"{}", hashes, unescaped, hashes),
-        );
+        edit.replace(token.text_range(), format!("r{}\"{}\"{}", hashes, unescaped, hashes));
     });
     ctx.build()
 }
 
-fn count_hashes(s: &str) -> usize {
-    let mut max_hash_streak = 0usize;
-    for idx in s.match_indices("\"#").map(|(i, _)| i) {
-        let (_, sub) = s.split_at(idx + 1);
-        let nb_hash = sub.chars().take_while(|c| *c == '#').count();
-        if nb_hash > max_hash_streak {
-            max_hash_streak = nb_hash;
-        }
-    }
-    max_hash_streak
-}
-
-fn find_usual_string_range(s: &str) -> Option<TextRange> {
-    Some(TextRange::from_to(
-        TextUnit::from(s.find('"')? as u32),
-        TextUnit::from(s.rfind('"')? as u32),
-    ))
-}
-
 pub(crate) fn make_usual_string(mut ctx: AssistCtx<impl HirDatabase>) -> Option<Assist> {
-    let literal = ctx.node_at_offset::<Literal>()?;
-    if literal.token().kind() != ra_syntax::SyntaxKind::RAW_STRING {
-        return None;
-    }
-    let token = literal.token();
+    let token = raw_string_token(&ctx)?;
     let text = token.text().as_str();
     let usual_string_range = find_usual_string_range(text)?;
     ctx.add_action(AssistId("make_usual_string"), "make usual string", |edit| {
-        edit.target(literal.syntax().text_range());
+        edit.target(token.text_range());
         // parse inside string to escape `"`
         let start_of_inside = usual_string_range.start().to_usize() + 1;
         let end_of_inside = usual_string_range.end().to_usize();
         let inside_str = &text[start_of_inside..end_of_inside];
         let escaped = inside_str.escape_default().to_string();
-        edit.replace(literal.syntax().text_range(), format!("\"{}\"", escaped));
+        edit.replace(token.text_range(), format!("\"{}\"", escaped));
     });
     ctx.build()
 }
 
 pub(crate) fn add_hash(mut ctx: AssistCtx<impl HirDatabase>) -> Option<Assist> {
-    let literal = ctx.node_at_offset::<Literal>()?;
-    if literal.token().kind() != ra_syntax::SyntaxKind::RAW_STRING {
-        return None;
-    }
+    let token = raw_string_token(&ctx)?;
     ctx.add_action(AssistId("add_hash"), "add hash to raw string", |edit| {
-        edit.target(literal.syntax().text_range());
-        edit.insert(literal.syntax().text_range().start() + TextUnit::of_char('r'), "#");
-        edit.insert(literal.syntax().text_range().end(), "#");
+        edit.target(token.text_range());
+        edit.insert(token.text_range().start() + TextUnit::of_char('r'), "#");
+        edit.insert(token.text_range().end(), "#");
     });
     ctx.build()
 }
 
 pub(crate) fn remove_hash(mut ctx: AssistCtx<impl HirDatabase>) -> Option<Assist> {
-    let literal = ctx.node_at_offset::<Literal>()?;
-    if literal.token().kind() != ra_syntax::SyntaxKind::RAW_STRING {
-        return None;
-    }
-    let token = literal.token();
+    let token = raw_string_token(&ctx)?;
     let text = token.text().as_str();
     if text.starts_with("r\"") {
         // no hash to remove
         return None;
     }
     ctx.add_action(AssistId("remove_hash"), "remove hash from raw string", |edit| {
-        edit.target(literal.syntax().text_range());
+        edit.target(token.text_range());
         let result = &text[2..text.len() - 1];
         let result = if result.starts_with("\"") {
             // no more hash, escape
@@ -117,11 +86,34 @@ pub(crate) fn remove_hash(mut ctx: AssistCtx<impl HirDatabase>) -> Option<Assist
         } else {
             result.to_owned()
         };
-        edit.replace(literal.syntax().text_range(), format!("r{}", result));
+        edit.replace(token.text_range(), format!("r{}", result));
     });
     ctx.build()
 }
 
+fn raw_string_token(ctx: &AssistCtx<impl HirDatabase>) -> Option<SyntaxToken> {
+    ctx.token_at_offset().right_biased().filter(|it| it.kind() == RAW_STRING)
+}
+
+fn count_hashes(s: &str) -> usize {
+    let mut max_hash_streak = 0usize;
+    for idx in s.match_indices("\"#").map(|(i, _)| i) {
+        let (_, sub) = s.split_at(idx + 1);
+        let nb_hash = sub.chars().take_while(|c| *c == '#').count();
+        if nb_hash > max_hash_streak {
+            max_hash_streak = nb_hash;
+        }
+    }
+    max_hash_streak
+}
+
+fn find_usual_string_range(s: &str) -> Option<TextRange> {
+    Some(TextRange::from_to(
+        TextUnit::from(s.find('"')? as u32),
+        TextUnit::from(s.rfind('"')? as u32),
+    ))
+}
+
 #[cfg(test)]
 mod test {
     use super::*;
@@ -159,6 +151,23 @@ string"#;
     }
 
     #[test]
+    fn make_raw_string_works_inside_macros() {
+        check_assist(
+            make_raw_string,
+            r#"
+            fn f() {
+                format!(<|>"x = {}", 92)
+            }
+            "#,
+            r##"
+            fn f() {
+                format!(<|>r#"x = {}"#, 92)
+            }
+            "##,
+        )
+    }
+
+    #[test]
     fn make_raw_string_hashes_inside_works() {
         check_assist(
             make_raw_string,