optimize joining and concatenation for slices

for both Vec<T> and String
- eliminates the boolean first flag in fn join()

for String only
- eliminates repeated bounds checks in join(), concat()
- adds fast paths for small string separators up to a len of 4 bytes

2 files changed, 113 insertions, 47 deletions

diff --git a/src/liballoc/slice.rs b/src/liballoc/slice.rs
index 161493f3892..82578c3206f 100644
--- a/src/liballoc/slice.rs
+++ b/src/liballoc/slice.rs
@@ -566,18 +566,18 @@ impl<T: Clone, V: Borrow<[T]>> SliceConcatExt<T> for [V] {
     }
 
     fn join(&self, sep: &T) -> Vec<T> {
-        let size = self.iter().fold(0, |acc, v| acc + v.borrow().len());
-        let mut result = Vec::with_capacity(size + self.len());
-        let mut first = true;
-        for v in self {
-            if first {
-                first = false
-            } else {
-                result.push(sep.clone())
+        let mut iter = self.iter();
+        iter.next().map_or(vec![], |first| {
+            let size = self.iter().fold(0, |acc, v| acc + v.borrow().len());
+            let mut result = Vec::with_capacity(size + self.len());
+            result.extend_from_slice(first.borrow());
+
+            for v in iter {
+                result.push(sep.clone());
+                result.extend_from_slice(v.borrow())
             }
-            result.extend_from_slice(v.borrow())
-        }
-        result
+            result
+        })
     }
 
     fn connect(&self, sep: &T) -> Vec<T> {
diff --git a/src/liballoc/str.rs b/src/liballoc/str.rs
index 823e56b64e3..45daabf86ab 100644
--- a/src/liballoc/str.rs
+++ b/src/liballoc/str.rs
@@ -86,52 +86,118 @@ impl<S: Borrow<str>> SliceConcatExt<str> for [S] {
     type Output = String;
 
     fn concat(&self) -> String {
-        if self.is_empty() {
-            return String::new();
-        }
-
-        // `len` calculation may overflow but push_str will check boundaries
-        let len = self.iter().map(|s| s.borrow().len()).sum();
-        let mut result = String::with_capacity(len);
-
-        for s in self {
-            result.push_str(s.borrow())
-        }
-
-        result
+        self.join("")
     }
 
     fn join(&self, sep: &str) -> String {
-        if self.is_empty() {
-            return String::new();
+        unsafe {
+            String::from_utf8_unchecked( join_generic_copy(self, sep.as_bytes()) )
         }
+    }
 
-        // concat is faster
-        if sep.is_empty() {
-            return self.concat();
-        }
+    fn connect(&self, sep: &str) -> String {
+        self.join(sep)
+    }
+}
 
-        // this is wrong without the guarantee that `self` is non-empty
-        // `len` calculation may overflow but push_str but will check boundaries
-        let len = sep.len() * (self.len() - 1) +
-                  self.iter().map(|s| s.borrow().len()).sum::<usize>();
-        let mut result = String::with_capacity(len);
-        let mut first = true;
+macro_rules! spezialize_for_lengths {
+    ($separator:expr, $target:expr, $iter:expr; $($num:expr),*) => {
+        let mut target = $target;
+        let iter = $iter;
+        let sep_len = $separator.len();
+        let sep_bytes = $separator;
+        match $separator.len() {
+            $(
+                // loops with hardcoded sizes run much faster
+                // specialize the cases with small separator lengths
+                $num => {
+                    for s in iter {
+                        target.get_unchecked_mut(..$num)
+                            .copy_from_slice(sep_bytes);
+
+                        let s_bytes = s.borrow().as_ref();
+                        let offset = s_bytes.len();
+                        target = {target}.get_unchecked_mut($num..);
+                        target.get_unchecked_mut(..offset)
+                            .copy_from_slice(s_bytes);
+                        target = {target}.get_unchecked_mut(offset..);
+                    }
+                },
+            )*
+            0 => {
+                // concat, same principle without the separator
+                for s in iter {
+                    let s_bytes = s.borrow().as_ref();
+                    let offset = s_bytes.len();
+                    target.get_unchecked_mut(..offset)
+                        .copy_from_slice(s_bytes);
+                    target = {target}.get_unchecked_mut(offset..);
+                }
+            },
+            _ => {
+                // arbitrary non-zero size fallback
+                for s in iter {
+                    target.get_unchecked_mut(..sep_len)
+                        .copy_from_slice(sep_bytes);
+
+                    let s_bytes = s.borrow().as_ref();
+                    let offset = s_bytes.len();
+                    target = {target}.get_unchecked_mut(sep_len..);
+                    target.get_unchecked_mut(..offset)
+                        .copy_from_slice(s_bytes);
+                    target = {target}.get_unchecked_mut(offset..);
+                }
+            }
+        }
+    };
+}
 
-        for s in self {
-            if first {
-                first = false;
-            } else {
-                result.push_str(sep);
+// Optimized join implementation that works for both Vec<T> (T: Copy) and String's inner vec
+// Currently (2018-05-13) there is a bug with type inference and specialization (see issue #36262)
+// For this reason SliceConcatExt<T> is not specialized for T: Copy and SliceConcatExt<str> is the
+// only user of this function. It is left in place for the time when that is fixed.
+//
+// the bounds for String-join are S: Borrow<str> and for Vec-join Borrow<[T]>
+// [T] and str both impl AsRef<[T]> for some T
+// => s.borrow().as_ref() and we always have slices
+fn join_generic_copy<B, T, S>(slice: &[S], sep: &[T]) -> Vec<T>
+where
+    T: Copy,
+    B: AsRef<[T]> + ?Sized,
+    S: Borrow<B>,
+{
+    let sep_len = sep.len();
+    let mut iter = slice.iter();
+    iter.next().map_or(vec![], |first| {
+        // this is wrong without the guarantee that `slice` is non-empty
+        // if the `len` calculation overflows, we'll panic
+        // we would have run out of memory anyway and the rest of the function requires
+        // the entire String pre-allocated for safety
+        //
+        // this is the exact len of the resulting String
+        let len =  sep_len.checked_mul(slice.len() - 1).and_then(|n| {
+            slice.iter().map(|s| s.borrow().as_ref().len()).try_fold(n, usize::checked_add)
+        }).expect("attempt to join into collection with len > usize::MAX");
+
+        // crucial for safety
+        let mut result = Vec::with_capacity(len);
+
+        unsafe {
+            result.extend_from_slice(first.borrow().as_ref());
+
+            {
+                let pos = result.len();
+                let target = result.get_unchecked_mut(pos..len);
+
+                // copy separator and strs over without bounds checks
+                // generate loops with hardcoded offsets for small separators
+                // massive improvements possible (~ x2)
+                spezialize_for_lengths!(sep, target, iter; 1, 2, 3, 4);
             }
-            result.push_str(s.borrow());
+            result.set_len(len);
         }
         result
-    }
-
-    fn connect(&self, sep: &str) -> String {
-        self.join(sep)
-    }
+    })
 }
 
 #[stable(feature = "rust1", since = "1.0.0")]