Use TrustedLen for Vec's FromIterator and Extend

diff --git a/src/libcollections/lib.rs b/src/libcollections/lib.rs
index 990de541b6783db047df209c0c013f7c5502bca7..23d6edd6d794e88c86cc7b2cfc2d0267365a2ee6 100644 (file)
--- a/src/libcollections/lib.rs
+++ b/src/libcollections/lib.rs
@@ -50,6 +50,7 @@
 #![feature(specialization)]
 #![feature(staged_api)]
 #![feature(step_by)]
+#![feature(trusted_len)]
 #![feature(unicode)]
 #![feature(unique)]
 #![cfg_attr(test, feature(rand, test))]

diff --git a/src/libcollections/vec.rs b/src/libcollections/vec.rs
index f3d78c20a4d6be2432e810c2f5fc62415b80a864..cd628a39af8b83a63ebf3272234a56a0bd6ee876 100644 (file)
--- a/src/libcollections/vec.rs
+++ b/src/libcollections/vec.rs
@@ -68,7 +68,7 @@
 use core::fmt;
 use core::hash::{self, Hash};
 use core::intrinsics::{arith_offset, assume};
-use core::iter::{FromIterator, FusedIterator};
+use core::iter::{FromIterator, FusedIterator, TrustedLen};
 use core::mem;
 use core::ops::{Index, IndexMut};
 use core::ops;
@@ -1589,6 +1589,18 @@ fn spec_extend(&mut self, ref mut other: Vec<T>) {
     }
 }
 
+trait IsTrustedLen : Iterator {
+    fn trusted_len(&self) -> Option<usize> { None }
+}
+impl<I> IsTrustedLen for I where I: Iterator { }
+
+impl<I> IsTrustedLen for I where I: TrustedLen
+{
+    fn trusted_len(&self) -> Option<usize> {
+        self.size_hint().1
+    }
+}
+
 impl<T> Vec<T> {
     fn extend_desugared<I: Iterator<Item = T>>(&mut self, mut iterator: I) {
         // This function should be the moral equivalent of:
@@ -1596,16 +1608,30 @@ fn extend_desugared<I: Iterator<Item = T>>(&mut self, mut iterator: I) {
         //      for item in iterator {
         //          self.push(item);
         //      }
-        while let Some(element) = iterator.next() {
-            let len = self.len();
-            if len == self.capacity() {
-                let (lower, _) = iterator.size_hint();
-                self.reserve(lower.saturating_add(1));
-            }
+        if let Some(additional) = iterator.trusted_len() {
+            self.reserve(additional);
             unsafe {
-                ptr::write(self.get_unchecked_mut(len), element);
-                // NB can't overflow since we would have had to alloc the address space
-                self.set_len(len + 1);
+                let mut ptr = self.as_mut_ptr().offset(self.len() as isize);
+                let mut local_len = SetLenOnDrop::new(&mut self.len);
+                for element in iterator {
+                    ptr::write(ptr, element);
+                    ptr = ptr.offset(1);
+                    // NB can't overflow since we would have had to alloc the address space
+                    local_len.increment_len(1);
+                }
+            }
+        } else {
+            while let Some(element) = iterator.next() {
+                let len = self.len();
+                if len == self.capacity() {
+                    let (lower, _) = iterator.size_hint();
+                    self.reserve(lower.saturating_add(1));
+                }
+                unsafe {
+                    ptr::write(self.get_unchecked_mut(len), element);
+                    // NB can't overflow since we would have had to alloc the address space
+                    self.set_len(len + 1);
+                }
             }
         }
     }