Merge pull request #286 from nyurik/litstr

Optimize performance for string literals in Display::fmt
dtolnay · Feb 11, 2024 · 097251d · 097251d
2 parents 0717de3 + cd79876
commit 097251d
Show file tree

Hide file tree

Showing 3 changed files with 79 additions and 4 deletions.
diff --git a/impl/src/attr.rs b/impl/src/attr.rs
@@ -18,6 +18,7 @@ pub struct Attrs<'a> {
 #[derive(Clone)]
 pub struct Display<'a> {
     pub original: &'a Attribute,
+    pub use_write_str: bool,
     pub fmt: LitStr,
     pub args: TokenStream,
     pub has_bonus_display: bool,
@@ -103,10 +104,14 @@ fn parse_error_attribute<'a>(attrs: &mut Attrs<'a>, attr: &'a Attribute) -> Resu
             return Ok(());
         }
 
+        let fmt = input.parse()?;
+        let args = parse_token_expr(input, false)?;
         let display = Display {
             original: attr,
-            fmt: input.parse()?,
-            args: parse_token_expr(input, false)?,
+            // This will be updated later if format_args are still required (i.e. has braces)
+            use_write_str: args.is_empty(),
+            fmt,
+            args,
             has_bonus_display: false,
             implied_bounds: Set::new(),
         };
@@ -196,8 +201,18 @@ impl ToTokens for Display<'_> {
     fn to_tokens(&self, tokens: &mut TokenStream) {
         let fmt = &self.fmt;
         let args = &self.args;
-        tokens.extend(quote! {
-            ::core::write!(__formatter, #fmt #args)
+
+        // Currently compiler is unable to generate as efficient code for
+        //    write!(f, "text")   as it does for   f.write_str("text"),
+        // so handle it here when the literal string has no braces/no args.
+        tokens.extend(if self.use_write_str {
+            quote! {
+                __formatter.write_str(#fmt)
+            }
+        } else {
+            quote! {
+                ::core::write!(__formatter, #fmt #args)
+            }
         });
     }
 }

diff --git a/impl/src/fmt.rs b/impl/src/fmt.rs
@@ -32,7 +32,12 @@ impl Display<'_> {
             }
         }
 
+        if self.use_write_str && fmt.contains('}') {
+            self.use_write_str = false;
+        }
+
         while let Some(brace) = read.find('{') {
+            self.use_write_str = false;
             out += &read[..brace + 1];
             read = &read[brace + 1..];
             if read.starts_with('{') {

diff --git a/tests/test_display.rs b/tests/test_display.rs
@@ -301,3 +301,58 @@ fn test_keyword() {
 
     assert("error: 1", Error);
 }
+
+#[test]
+fn test_str_special_chars() {
+    #[derive(Error, Debug)]
+    pub enum Error {
+        #[error("brace left {{")]
+        BraceLeft,
+        #[error("brace left 2 \x7B\x7B")]
+        BraceLeft2,
+        #[error("brace left 3 \u{7B}\u{7B}")]
+        BraceLeft3,
+        #[error("brace right }}")]
+        BraceRight,
+        #[error("brace right 2 \x7D\x7D")]
+        BraceRight2,
+        #[error("brace right 3 \u{7D}\u{7D}")]
+        BraceRight3,
+        #[error(
+            "new_\
+line"
+        )]
+        NewLine,
+        #[error("escape24 \u{78}")]
+        Escape24,
+    }
+
+    assert("brace left {", Error::BraceLeft);
+    assert("brace left 2 {", Error::BraceLeft2);
+    assert("brace left 3 {", Error::BraceLeft3);
+    assert("brace right }", Error::BraceRight);
+    assert("brace right 2 }", Error::BraceRight2);
+    assert("brace right 3 }", Error::BraceRight3);
+    assert("new_line", Error::NewLine);
+    assert("escape24 x", Error::Escape24);
+}
+
+#[test]
+fn test_raw_str() {
+    #[derive(Error, Debug)]
+    pub enum Error {
+        #[error(r#"raw brace left {{"#)]
+        BraceLeft,
+        #[error(r#"raw brace left 2 \x7B"#)]
+        BraceLeft2,
+        #[error(r#"raw brace right }}"#)]
+        BraceRight,
+        #[error(r#"raw brace right 2 \x7D"#)]
+        BraceRight2,
+    }
+
+    assert(r#"raw brace left {"#, Error::BraceLeft);
+    assert(r#"raw brace left 2 \x7B"#, Error::BraceLeft2);
+    assert(r#"raw brace right }"#, Error::BraceRight);
+    assert(r#"raw brace right 2 \x7D"#, Error::BraceRight2);
+}