From 581e1bf0d7b05cafdc60c1f132ba99c5ca338bc6 Mon Sep 17 00:00:00 2001
From: yipinliu <ypliu18@gmail.com>
Date: Sat, 12 Mar 2022 19:00:17 +0800
Subject: [PATCH] Highlight escape sequences in byte strings

---
 crates/ide/src/syntax_highlighting.rs         | 29 +++++-------
 crates/ide/src/syntax_highlighting/escape.rs  | 25 ++++++++++
 .../test_data/highlight_byte_string.html      | 47 +++++++++++++++++++
 crates/ide/src/syntax_highlighting/tests.rs   | 14 ++++++
 4 files changed, 97 insertions(+), 18 deletions(-)
 create mode 100644 crates/ide/src/syntax_highlighting/escape.rs
 create mode 100644 crates/ide/src/syntax_highlighting/test_data/highlight_byte_string.html

diff --git a/crates/ide/src/syntax_highlighting.rs b/crates/ide/src/syntax_highlighting.rs
index 0119ba5cb779..1bd95c4c437a 100644
--- a/crates/ide/src/syntax_highlighting.rs
+++ b/crates/ide/src/syntax_highlighting.rs
@@ -7,6 +7,7 @@ mod highlight;
 mod format;
 mod macro_;
 mod inject;
+mod escape;
 
 mod html;
 #[cfg(test)]
@@ -16,7 +17,7 @@ use hir::{InFile, Name, Semantics};
 use ide_db::RootDatabase;
 use rustc_hash::FxHashMap;
 use syntax::{
-    ast::{self, IsString},
+    ast::{self},
     AstNode, AstToken, NodeOrToken,
     SyntaxKind::*,
     SyntaxNode, TextRange, WalkEvent, T,
@@ -30,6 +31,7 @@ use crate::{
     FileId, HlMod, HlTag,
 };
 
+use crate::syntax_highlighting::escape::highlight_escape_string;
 pub(crate) use html::highlight_as_html;
 
 #[derive(Debug, Clone, Copy)]
@@ -371,29 +373,20 @@ fn traverse(
         // string highlight injections, note this does not use the descended element as proc-macros
         // can rewrite string literals which invalidates our indices
         if let (Some(token), Some(descended_token)) = (token, descended_element.as_token()) {
-            let string = ast::String::cast(token);
-            let string_to_highlight = ast::String::cast(descended_token.clone());
-            if let Some((string, expanded_string)) = string.zip(string_to_highlight) {
+            if ast::String::can_cast(token.kind()) && ast::String::can_cast(descended_token.kind())
+            {
+                let string = ast::String::cast(token).unwrap();
+                let expanded_string = ast::String::cast(descended_token.clone()).unwrap();
                 if string.is_raw() {
                     if inject::ra_fixture(hl, sema, &string, &expanded_string).is_some() {
                         continue;
                     }
                 }
                 highlight_format_string(hl, &string, &expanded_string, range);
-                // Highlight escape sequences
-                string.escaped_char_ranges(&mut |piece_range, char| {
-                    if char.is_err() {
-                        return;
-                    }
-
-                    if string.text()[piece_range.start().into()..].starts_with('\\') {
-                        hl.add(HlRange {
-                            range: piece_range + range.start(),
-                            highlight: HlTag::EscapeSequence.into(),
-                            binding_hash: None,
-                        });
-                    }
-                });
+                highlight_escape_string(hl, &string, range.start());
+            } else if ast::ByteString::can_cast(token.kind()) {
+                let byte_string = ast::ByteString::cast(token).unwrap();
+                highlight_escape_string(hl, &byte_string, range.start());
             }
         }
 
diff --git a/crates/ide/src/syntax_highlighting/escape.rs b/crates/ide/src/syntax_highlighting/escape.rs
new file mode 100644
index 000000000000..1b53edb7dc47
--- /dev/null
+++ b/crates/ide/src/syntax_highlighting/escape.rs
@@ -0,0 +1,25 @@
+use crate::syntax_highlighting::highlights::Highlights;
+use crate::{HlRange, HlTag};
+use syntax::ast::IsString;
+use syntax::TextSize;
+
+/// Highlight escape sequences
+pub(super) fn highlight_escape_string<T: IsString>(
+    stack: &mut Highlights,
+    string: &T,
+    start: TextSize,
+) {
+    string.escaped_char_ranges(&mut |piece_range, char| {
+        if char.is_err() {
+            return;
+        }
+
+        if string.text()[piece_range.start().into()..].starts_with('\\') {
+            stack.add(HlRange {
+                range: piece_range + start,
+                highlight: HlTag::EscapeSequence.into(),
+                binding_hash: None,
+            });
+        }
+    });
+}
diff --git a/crates/ide/src/syntax_highlighting/test_data/highlight_byte_string.html b/crates/ide/src/syntax_highlighting/test_data/highlight_byte_string.html
new file mode 100644
index 000000000000..113b73af0c4f
--- /dev/null
+++ b/crates/ide/src/syntax_highlighting/test_data/highlight_byte_string.html
@@ -0,0 +1,47 @@
+
+<style>
+body                { margin: 0; }
+pre                 { color: #DCDCCC; background: #3F3F3F; font-size: 22px; padding: 0.4em; }
+
+.lifetime           { color: #DFAF8F; font-style: italic; }
+.label              { color: #DFAF8F; font-style: italic; }
+.comment            { color: #7F9F7F; }
+.documentation      { color: #629755; }
+.intra_doc_link     { font-style: italic; }
+.injected           { opacity: 0.65 ; }
+.struct, .enum      { color: #7CB8BB; }
+.enum_variant       { color: #BDE0F3; }
+.string_literal     { color: #CC9393; }
+.field              { color: #94BFF3; }
+.function           { color: #93E0E3; }
+.function.unsafe    { color: #BC8383; }
+.trait.unsafe       { color: #BC8383; }
+.operator.unsafe    { color: #BC8383; }
+.mutable.unsafe     { color: #BC8383; text-decoration: underline; }
+.keyword.unsafe     { color: #BC8383; font-weight: bold; }
+.parameter          { color: #94BFF3; }
+.text               { color: #DCDCCC; }
+.type               { color: #7CB8BB; }
+.builtin_type       { color: #8CD0D3; }
+.type_param         { color: #DFAF8F; }
+.attribute          { color: #94BFF3; }
+.numeric_literal    { color: #BFEBBF; }
+.bool_literal       { color: #BFE6EB; }
+.macro              { color: #94BFF3; }
+.derive             { color: #94BFF3; font-style: italic; }
+.module             { color: #AFD8AF; }
+.value_param        { color: #DCDCCC; }
+.variable           { color: #DCDCCC; }
+.format_specifier   { color: #CC696B; }
+.mutable            { text-decoration: underline; }
+.escape_sequence    { color: #94BFF3; }
+.keyword            { color: #F0DFAF; font-weight: bold; }
+.control            { font-style: italic; }
+.reference          { font-style: italic; font-weight: bold; }
+
+.unresolved_reference { color: #FC5555; text-decoration: wavy underline; }
+</style>
+<pre><code><span class="keyword">fn</span> <span class="function declaration">main</span><span class="parenthesis">(</span><span class="parenthesis">)</span> <span class="brace">{</span>
+    <span class="keyword">let</span> <span class="punctuation">_</span> <span class="operator">=</span> <span class="string_literal">"</span><span class="escape_sequence">\x28</span><span class="escape_sequence">\x28</span><span class="escape_sequence">\x00</span><span class="escape_sequence">\x63</span><span class="escape_sequence">\n</span><span class="string_literal">"</span><span class="semicolon">;</span>
+    <span class="keyword">let</span> <span class="punctuation">_</span> <span class="operator">=</span> <span class="string_literal">b"</span><span class="escape_sequence">\x28</span><span class="escape_sequence">\x28</span><span class="escape_sequence">\x00</span><span class="escape_sequence">\x63</span><span class="escape_sequence">\n</span><span class="string_literal">"</span><span class="semicolon">;</span>
+<span class="brace">}</span></code></pre>
\ No newline at end of file
diff --git a/crates/ide/src/syntax_highlighting/tests.rs b/crates/ide/src/syntax_highlighting/tests.rs
index efd03c28237c..f9c2404c358f 100644
--- a/crates/ide/src/syntax_highlighting/tests.rs
+++ b/crates/ide/src/syntax_highlighting/tests.rs
@@ -476,6 +476,20 @@ fn main() {
     );
 }
 
+#[test]
+fn test_byte_string_highlight() {
+    check_highlighting(
+        r#"
+fn main() {
+    let _ = "\x28\x28\x00\x63\n";
+    let _ = b"\x28\x28\x00\x63\n";
+}
+"#,
+        expect_file!["./test_data/highlight_byte_string.html"],
+        false,
+    );
+}
+
 #[test]
 fn test_unsafe_highlighting() {
     check_highlighting(