diff --git a/src/libcollections/lib.rs b/src/libcollections/lib.rs index 990de541b678..23d6edd6d794 100644 --- 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 f3d78c20a4d6..cd628a39af8b 100644 --- a/src/libcollections/vec.rs +++ b/src/libcollections/vec.rs @@ -68,7 +68,7 @@ use core::cmp::Ordering; 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 @@ impl SpecExtend> for Vec { } } +trait IsTrustedLen : Iterator { + fn trusted_len(&self) -> Option { None } +} +impl IsTrustedLen for I where I: Iterator { } + +impl IsTrustedLen for I where I: TrustedLen +{ + fn trusted_len(&self) -> Option { + self.size_hint().1 + } +} + impl Vec { fn extend_desugared>(&mut self, mut iterator: I) { // This function should be the moral equivalent of: @@ -1596,16 +1608,30 @@ impl Vec { // 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); + } } } }