From 4a254c00506abdbb660e9c71d34b5b836b86da8d Mon Sep 17 00:00:00 2001
From: Shotaro Yamada <sinkuu@sinkuu.xyz>
Date: Wed, 14 Mar 2018 18:11:42 +0900
Subject: [PATCH] Escape stringified expression

Payload of `Literal` token must be escaped.
Also print printable non-ASCII characters.
---
 src/libsyntax_ext/assert.rs        | 68 +++++++++++++++++++++++++++---
 src/libsyntax_ext/lib.rs           |  1 +
 src/test/run-pass/assert-escape.rs | 13 ++++++
 3 files changed, 77 insertions(+), 5 deletions(-)
 create mode 100644 src/test/run-pass/assert-escape.rs

diff --git a/src/libsyntax_ext/assert.rs b/src/libsyntax_ext/assert.rs
index 7962ec26c37ad..8b29e6adeb9eb 100644
--- a/src/libsyntax_ext/assert.rs
+++ b/src/libsyntax_ext/assert.rs
@@ -15,7 +15,7 @@ use syntax::ext::build::AstBuilder;
 use syntax::parse::token;
 use syntax::print::pprust;
 use syntax::tokenstream::{TokenStream, TokenTree};
-use syntax_pos::Span;
+use syntax_pos::{Span, DUMMY_SP};
 
 pub fn expand_assert<'cx>(
     cx: &'cx mut ExtCtxt,
@@ -41,10 +41,18 @@ pub fn expand_assert<'cx>(
         tts: if let Some(ts) = custom_msg_args {
             ts.into()
         } else {
-            let panic_str = format!("assertion failed: {}", pprust::expr_to_string(&cond_expr));
-            TokenStream::from(token::Literal(
-                token::Lit::Str_(Name::intern(&panic_str)),
-                None,
+            // `expr_to_string` escapes the string literals with `.escape_default()`
+            // which escapes all non-ASCII characters with `\u`.
+            let escaped_expr = escape_format_string(&unescape_printable_unicode(
+                &pprust::expr_to_string(&cond_expr),
+            ));
+
+            TokenStream::from(TokenTree::Token(
+                DUMMY_SP,
+                token::Literal(
+                    token::Lit::Str_(Name::intern(&format!("assertion failed: {}", escaped_expr))),
+                    None,
+                ),
             )).into()
         },
     };
@@ -62,3 +70,53 @@ pub fn expand_assert<'cx>(
     );
     MacEager::expr(if_expr)
 }
+
+/// Escapes a string for use as a formatting string.
+fn escape_format_string(s: &str) -> String {
+    let mut res = String::with_capacity(s.len());
+    for c in s.chars() {
+        res.extend(c.escape_debug());
+        match c {
+            '{' | '}' => res.push(c),
+            _ => {}
+        }
+    }
+    res
+}
+
+#[test]
+fn test_escape_format_string() {
+    assert!(escape_format_string(r"foo{}\") == r"foo{{}}\\");
+}
+
+/// Unescapes the escaped unicodes (`\u{...}`) that are printable.
+fn unescape_printable_unicode(mut s: &str) -> String {
+    use std::{char, u32};
+
+    let mut res = String::with_capacity(s.len());
+
+    loop {
+        if let Some(start) = s.find(r"\u{") {
+            res.push_str(&s[0..start]);
+            s = &s[start..];
+            s.find('}')
+                .and_then(|end| {
+                    let v = u32::from_str_radix(&s[3..end], 16).ok()?;
+                    let c = char::from_u32(v)?;
+                    // Escape unprintable characters.
+                    res.extend(c.escape_debug());
+                    s = &s[end + 1..];
+                    Some(())
+                })
+                .expect("lexer should have rejected invalid escape sequences");
+        } else {
+            res.push_str(s);
+            return res;
+        }
+    }
+}
+
+#[test]
+fn test_unescape_printable_unicode() {
+    assert!(unescape_printable_unicode(r"\u{2603}\n\u{0}") == r"☃\n\u{0}");
+}
diff --git a/src/libsyntax_ext/lib.rs b/src/libsyntax_ext/lib.rs
index a01878530b2ac..1a0d22a82c4c7 100644
--- a/src/libsyntax_ext/lib.rs
+++ b/src/libsyntax_ext/lib.rs
@@ -17,6 +17,7 @@
 
 #![feature(proc_macro_internals)]
 #![feature(decl_macro)]
+#![feature(str_escape)]
 
 extern crate fmt_macros;
 #[macro_use]
diff --git a/src/test/run-pass/assert-escape.rs b/src/test/run-pass/assert-escape.rs
new file mode 100644
index 0000000000000..d340806c3577d
--- /dev/null
+++ b/src/test/run-pass/assert-escape.rs
@@ -0,0 +1,13 @@
+// Copyright 2018 The Rust Project Developers. See the COPYRIGHT
+// file at the top-level directory of this distribution and at
+// http://rust-lang.org/COPYRIGHT.
+//
+// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
+// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
+// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
+// option. This file may not be copied, modified, or distributed
+// except according to those terms.
+
+fn main() {
+    assert!(r#"☃\backslash"#.contains("\\"));
+}