core and std: add format_arg!() to speed up write*!() and print*!()

Before this patch, `write!(wr, "foo")` was not as fast as
`wr.write_str("foo")` because the underlying write_fmt has some overhead
in order to work with multiple arguments.

Author: erickt

src/libcollections/macros.rs

@@ -33,8 +33,21 @@ macro_rules! vec {
 /// format!("hello {}", "world!");
 /// format!("x = {}, y = {y}", 10, y = 30);
 /// ```
+#[cfg(stage0)]
 #[macro_export]
 #[stable(feature = "rust1", since = "1.0.0")]
 macro_rules! format {
     ($($arg:tt)*) => ($crate::fmt::format(format_args!($($arg)*)))
 }
+
+#[cfg(not(stage0))]
+#[macro_export]
+#[stable(feature = "rust1", since = "1.0.0")]
+macro_rules! format {
+    ($fmt:expr) => {
+        format_arg!($fmt).to_string()
+    };
+    ($fmt:expr, $($arg:tt)+) => {
+        $crate::fmt::format(format_args!($fmt, $($arg)*))
+    }
+}

src/libcore/macros.rs

@@ -177,11 +177,23 @@ macro_rules! try {
 /// write!(&mut w, "test");
 /// write!(&mut w, "formatted {}", "arguments");
 /// ```
+#[cfg(stage0)]
 #[macro_export]
 macro_rules! write {
     ($dst:expr, $($arg:tt)*) => ((&mut *$dst).write_fmt(format_args!($($arg)*)))
 }
 
+#[cfg(not(stage0))]
+#[macro_export]
+macro_rules! write {
+    ($dst:expr, $fmt:expr) => {
+        (&mut *$dst).write_str(format_arg!($fmt))
+    };
+    ($dst:expr, $fmt:expr, $($arg:tt)+) => {
+        (&mut *$dst).write_fmt(format_args!($fmt, $($arg)*))
+    };
+}
+
 /// Equivalent to the `write!` macro, except that a newline is appended after
 /// the message is written.
 #[macro_export]

src/libstd/macros.rs

@@ -79,12 +79,25 @@ macro_rules! format {
 
 /// Equivalent to the `println!` macro except that a newline is not printed at
 /// the end of the message.
+#[cfg(stage0)]
 #[macro_export]
 #[stable(feature = "rust1", since = "1.0.0")]
 macro_rules! print {
     ($($arg:tt)*) => ($crate::old_io::stdio::print_args(format_args!($($arg)*)))
 }
 
+#[cfg(not(stage0))]
+#[macro_export]
+#[stable(feature = "rust1", since = "1.0.0")]
+macro_rules! print {
+    ($fmt:expr) => {
+        $crate::old_io::stdio::print(format_arg!($fmt))
+    };
+    ($fmt:expr, $($arg:tt)+) => {
+        $crate::old_io::stdio::print_args(format_args!($fmt, $($arg)*))
+    };
+}
+
 /// Macro for printing to a task's stdout handle.
 ///
 /// Each task can override its stdout handle via `std::old_io::stdio::set_stdout`.
@@ -97,12 +110,25 @@ macro_rules! print {
 /// println!("hello there!");
 /// println!("format {} arguments", "some");
 /// ```
+#[cfg(stage0)]
 #[macro_export]
 #[stable(feature = "rust1", since = "1.0.0")]
 macro_rules! println {
     ($($arg:tt)*) => ($crate::old_io::stdio::println_args(format_args!($($arg)*)))
 }
 
+#[cfg(not(stage0))]
+#[macro_export]
+#[stable(feature = "rust1", since = "1.0.0")]
+macro_rules! println {
+    ($fmt:expr) => {
+        $crate::old_io::stdio::println(format_arg!($fmt))
+    };
+    ($fmt:expr, $($arg:tt)*) => {
+        $crate::old_io::stdio::println_args(format_args!($fmt, $($arg)*))
+    };
+}
+
 /// Helper macro for unwrapping `Result` values while returning early with an
 /// error if the value of the expression is `Err`. For more information, see
 /// `std::io`.
@@ -185,6 +211,28 @@ macro_rules! log {
 /// into libsyntax itself.
 #[cfg(dox)]
 pub mod builtin {
+    /// Parse a `&'static str` with a compatible format to `format_args!()`.
+    ///
+    /// This macro produces a value of type `&'static str`, and is intended to
+    /// be used by the other formatting macros (`format!`, `write!`, `println!`)
+    /// are proxied through this one.
+    ///
+    /// For more information, see the documentation in `std::fmt`.
+    ///
+    /// # Example
+    ///
+    /// ```rust
+    /// use std::fmt;
+    ///
+    /// let s = format_arg!("hello");
+    /// assert_eq!(s, format!("hello"));
+    ///
+    /// ```
+    #[macro_export]
+    macro_rules! format_arg { ($fmt:expr) => ({
+        /* compiler built-in */
+    }) }
+
     /// The core macro for formatted string creation & output.
     ///
     /// This macro produces a value of type `fmt::Arguments`. This value can be

src/libsyntax/ext/base.rs

@@ -447,6 +447,9 @@ fn initial_syntax_expander_table(ecfg: &expand::ExpansionConfig) -> SyntaxEnv {
 
     let mut syntax_expanders = SyntaxEnv::new();
     syntax_expanders.insert(intern("macro_rules"), MacroRulesTT);
+    syntax_expanders.insert(intern("format_arg"),
+                            builtin_normal_expander(
+                                ext::format::expand_format_arg));
     syntax_expanders.insert(intern("format_args"),
                             builtin_normal_expander(
                                 ext::format::expand_format_args));

src/libsyntax/ext/format.rs

@@ -57,7 +57,7 @@ struct Context<'a, 'b:'a> {
     /// Collection of the compiled `rt::Argument` structures
     pieces: Vec<P<ast::Expr>>,
     /// Collection of string literals
-    str_pieces: Vec<P<ast::Expr>>,
+    str_pieces: Vec<(Span, token::InternedString)>,
     /// Stays `true` if all formatting parameters are default (as in "{}{}").
     all_pieces_simple: bool,
 
@@ -335,11 +335,11 @@ impl<'a, 'b> Context<'a, 'b> {
     }
 
     /// Translate the accumulated string literals to a literal expression
-    fn trans_literal_string(&mut self) -> P<ast::Expr> {
+    fn trans_string(&mut self) -> (Span, token::InternedString) {
         let sp = self.fmtsp;
         let s = token::intern_and_get_ident(&self.literal[]);
         self.literal.clear();
-        self.ecx.expr_str(sp, s)
+        (sp, s)
     }
 
     /// Translate a `parse::Piece` to a static `rt::Argument` or append
@@ -475,10 +475,15 @@ impl<'a, 'b> Context<'a, 'b> {
                 self.ecx.ty_ident(self.fmtsp, self.ecx.ident_of("str")),
                 Some(static_lifetime),
                 ast::MutImmutable);
+
+        let str_pieces = self.str_pieces.iter()
+            .map(|&(sp, ref s)| self.ecx.expr_str(sp, s.clone()))
+            .collect();
+
         let pieces = Context::static_array(self.ecx,
                                            "__STATIC_FMTSTR",
                                            piece_ty,
-                                           self.str_pieces);
+                                           str_pieces);
 
 
         // Right now there is a bug such that for the expression:
@@ -627,6 +632,43 @@ impl<'a, 'b> Context<'a, 'b> {
     }
 }
 
+pub fn expand_format_arg<'cx>(ecx: &'cx mut ExtCtxt, sp: Span,
+                              tts: &[ast::TokenTree])
+                              -> Box<base::MacResult+'cx> {
+
+    match parse_args(ecx, sp, tts) {
+        Some((efmt, args, order, names)) => {
+            if !args.is_empty() || !order.is_empty() || !names.is_empty() {
+                ecx.span_err(sp,
+                             &format!("requires no arguments, found {}",
+                                      args.len() + order.len()));
+
+                DummyResult::expr(sp)
+            } else {
+                MacExpr::new(expand_preparsed_format_arg(ecx, sp, efmt))
+            }
+        }
+        None => DummyResult::expr(sp)
+    }
+}
+
+/// Take the various parts of `format_arg!(efmt)`
+/// and construct the appropriate formatting expression.
+pub fn expand_preparsed_format_arg(ecx: &mut ExtCtxt, sp: Span,
+                                   efmt: P<ast::Expr>)
+                                   -> P<ast::Expr> {
+    match parse_format_args(ecx, sp, efmt, Vec::new(), Vec::new(), HashMap::new()) {
+        Some(cx) => {
+            let mut s = String::new();
+            for &(_, ref part) in cx.str_pieces.iter() {
+                s.push_str(part.as_slice());
+            }
+            cx.ecx.expr_str(sp, token::intern_and_get_ident(&s[]))
+        }
+        None => DummyResult::raw_expr(sp),
+    }
+}
+
 pub fn expand_format_args<'cx>(ecx: &'cx mut ExtCtxt, sp: Span,
                                tts: &[ast::TokenTree])
                                -> Box<base::MacResult+'cx> {
@@ -648,6 +690,18 @@ pub fn expand_preparsed_format_args(ecx: &mut ExtCtxt, sp: Span,
                                     name_ordering: Vec<String>,
                                     names: HashMap<String, P<ast::Expr>>)
                                     -> P<ast::Expr> {
+    match parse_format_args(ecx, sp, efmt, args, name_ordering, names) {
+        Some(cx) => cx.into_expr(),
+        None => DummyResult::raw_expr(sp),
+    }
+}
+
+fn parse_format_args<'a, 'b: 'a>(ecx: &'a mut ExtCtxt<'b>, sp: Span,
+                          efmt: P<ast::Expr>,
+                          args: Vec<P<ast::Expr>>,
+                          name_ordering: Vec<String>,
+                          names: HashMap<String, P<ast::Expr>>)
+                          -> Option<Context<'a, 'b>> {
     let arg_types: Vec<_> = (0..args.len()).map(|_| None).collect();
     let mut cx = Context {
         ecx: ecx,
@@ -670,7 +724,7 @@ pub fn expand_preparsed_format_args(ecx: &mut ExtCtxt, sp: Span,
                                    efmt,
                                    "format argument must be a string literal.") {
         Some((fmt, _)) => fmt,
-        None => return DummyResult::raw_expr(sp)
+        None => return None
     };
 
     let mut parser = parse::Parser::new(&fmt);
@@ -682,7 +736,7 @@ pub fn expand_preparsed_format_args(ecx: &mut ExtCtxt, sp: Span,
                 cx.verify_piece(&piece);
                 match cx.trans_piece(&piece) {
                     Some(piece) => {
-                        let s = cx.trans_literal_string();
+                        let s = cx.trans_string();
                         cx.str_pieces.push(s);
                         cx.pieces.push(piece);
                     }
@@ -695,10 +749,10 @@ pub fn expand_preparsed_format_args(ecx: &mut ExtCtxt, sp: Span,
     if !parser.errors.is_empty() {
         cx.ecx.span_err(cx.fmtsp, &format!("invalid format string: {}",
                                           parser.errors.remove(0))[]);
-        return DummyResult::raw_expr(sp);
+        return None;
     }
     if !cx.literal.is_empty() {
-        let s = cx.trans_literal_string();
+        let s = cx.trans_string();
         cx.str_pieces.push(s);
     }
 
@@ -714,5 +768,5 @@ pub fn expand_preparsed_format_args(ecx: &mut ExtCtxt, sp: Span,
         }
     }
 
-    cx.into_expr()
+    Some(cx)
 }

src/test/compile-fail/ifmt-bad-format-args.rs

@@ -9,6 +9,9 @@
 // except according to those terms.
 
 fn main() {
+    format_arg!(); //~ ERROR: requires at least a format string argument
+    format_arg!("{}", "bar"); //~ ERROR: requires no arguments, found 1.
+
     format_args!(); //~ ERROR: requires at least a format string argument
     format_args!(|| {}); //~ ERROR: must be a string literal
 }

src/test/run-pass/ifmt.rs

@@ -157,6 +157,7 @@ pub fn main() {
         format!("{}", a);
     }
 
+    test_format_arg();
     test_format_args();
 
     // test that trailing commas are acceptable
@@ -172,13 +173,17 @@ fn test_write() {
     write!(&mut buf, "{}", 3);
     {
         let w = &mut buf;
+        write!(w, "{{");
+        write!(w, "foo");
         write!(w, "{foo}", foo=4);
         write!(w, "{}", "hello");
+        writeln!(w, "line");
         writeln!(w, "{}", "line");
         writeln!(w, "{foo}", foo="bar");
+        writeln!(w, "}}");
     }
 
-    t!(buf, "34helloline\nbar\n");
+    t!(buf, "{{foo34helloline\nbar\n}}\n");
 }
 
 // Just make sure that the macros are defined, there's not really a lot that we
@@ -198,12 +203,18 @@ fn test_format_args() {
     let mut buf = String::new();
     {
         let w = &mut buf;
+        write!(w, "{{");
         write!(w, "{}", format_args!("{}", 1));
         write!(w, "{}", format_args!("test"));
         write!(w, "{}", format_args!("{test}", test=3));
+        write!(w, "}}");
     }
     let s = buf;
-    t!(s, "1test3");
+    t!(s, "{1test3}");
+
+    t!(format_args!("hello"), "hello");
+    t!(format_args!("{{"), "{");
+    t!(format_args!("}}"), "}");
 
     let s = fmt::format(format_args!("hello {}", "world"));
     t!(s, "hello world");