From e74b55b9da5a163e02bbcbf97e7ea24d48742943 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Micha=C5=82=20Krasnoborski?= <mkrdln@gmail.com>
Date: Fri, 27 Jan 2017 04:51:24 +0100
Subject: [PATCH] use `String::with_capacity` in `format!`

---
 src/libcollections/fmt.rs  |  3 ++-
 src/libcore/fmt/mod.rs     | 28 ++++++++++++++++++++++++++++
 src/libcoretest/fmt/mod.rs |  8 ++++++++
 src/libcoretest/lib.rs     |  1 +
 4 files changed, 39 insertions(+), 1 deletion(-)

diff --git a/src/libcollections/fmt.rs b/src/libcollections/fmt.rs
index 883417e9f4e..bd74848a01d 100644
--- a/src/libcollections/fmt.rs
+++ b/src/libcollections/fmt.rs
@@ -539,7 +539,8 @@ use string;
 /// [format!]: ../macro.format.html
 #[stable(feature = "rust1", since = "1.0.0")]
 pub fn format(args: Arguments) -> string::String {
-    let mut output = string::String::new();
+    let capacity = args.estimated_capacity();
+    let mut output = string::String::with_capacity(capacity);
     let _ = output.write_fmt(args);
     output
 }
diff --git a/src/libcore/fmt/mod.rs b/src/libcore/fmt/mod.rs
index 2ba7d6e8bd1..ed9c5d3337f 100644
--- a/src/libcore/fmt/mod.rs
+++ b/src/libcore/fmt/mod.rs
@@ -265,6 +265,34 @@ impl<'a> Arguments<'a> {
             args: args
         }
     }
+
+    /// Estimates the length of the formatted text.
+    ///
+    /// This is intended to be used for setting initial `String` capacity
+    /// when using `format!`. Note: this is neither the lower nor upper bound.
+    #[doc(hidden)] #[inline]
+    #[unstable(feature = "fmt_internals", reason = "internal to format_args!",
+               issue = "0")]
+    pub fn estimated_capacity(&self) -> usize {
+        // Using wrapping arithmetics in this function, because
+        // wrong result is highly unlikely and doesn't cause unsafety.
+        use ::num::Wrapping as W;
+
+        let pieces_length: W<usize> = self.pieces.iter()
+            .map(|x| W(x.len())).sum();
+
+        // If they are any arguments to format, the string will most likely
+        // double in size. So we're pre-doubling it here.
+        let multiplier = if self.args.is_empty() { W(1) } else { W(2) };
+
+        let capacity = multiplier * pieces_length;
+        if multiplier == W(2) && (W(1)..W(8)).contains(capacity) {
+            // Allocations smaller than 8 don't really make sense for String.
+            8
+        } else {
+            capacity.0
+        }
+    }
 }
 
 /// This structure represents a safely precompiled version of a format string
diff --git a/src/libcoretest/fmt/mod.rs b/src/libcoretest/fmt/mod.rs
index ed33596e1c2..71b3a440f7b 100644
--- a/src/libcoretest/fmt/mod.rs
+++ b/src/libcoretest/fmt/mod.rs
@@ -28,3 +28,11 @@ fn test_pointer_formats_data_pointer() {
     assert_eq!(format!("{:p}", s), format!("{:p}", s.as_ptr()));
     assert_eq!(format!("{:p}", b), format!("{:p}", b.as_ptr()));
 }
+
+#[test]
+fn test_estimated_capacity() {
+    assert_eq!(format_args!("{}", "").estimated_capacity(), 0);
+    assert_eq!(format_args!("Hello").estimated_capacity(), 5);
+    assert_eq!(format_args!("Hello, {}!", "").estimated_capacity(), 16);
+    assert_eq!(format_args!("{}, hello!", "World").estimated_capacity(), 16);
+}
diff --git a/src/libcoretest/lib.rs b/src/libcoretest/lib.rs
index ee47b510ee0..fed5b86c369 100644
--- a/src/libcoretest/lib.rs
+++ b/src/libcoretest/lib.rs
@@ -34,6 +34,7 @@
 #![feature(ordering_chaining)]
 #![feature(result_unwrap_or_default)]
 #![feature(ptr_unaligned)]
+#![feature(fmt_internals)]
 
 extern crate core;
 extern crate test;