Merge pull request #112 from bluss/extend-improvement

Improve .extend() performance (?)

benches/extend.rs

@@ -2,6 +2,8 @@
 extern crate arrayvec;
 #[macro_use] extern crate bencher;
 
+use std::io::Write;
+
 use arrayvec::ArrayVec;
 
 use bencher::Bencher;
@@ -12,8 +14,9 @@ fn extend_with_constant(b: &mut Bencher) {
     let cap = v.capacity();
     b.iter(|| {
         v.clear();
-        v.extend((0..cap).map(|_| 1));
-        v[0]
+        let constant = black_box(1);
+        v.extend((0..cap).map(move |_| constant));
+        v[511]
     });
     b.bytes = v.capacity() as u64;
 }
@@ -23,8 +26,9 @@ fn extend_with_range(b: &mut Bencher) {
     let cap = v.capacity();
     b.iter(|| {
         v.clear();
-        v.extend((0..cap).map(|x| x as _));
-        v[0]
+        let range = 0..cap;
+        v.extend(range.map(|x| black_box(x as _)));
+        v[511]
     });
     b.bytes = v.capacity() as u64;
 }
@@ -34,8 +38,20 @@ fn extend_with_slice(b: &mut Bencher) {
     let data = [1; 512];
     b.iter(|| {
         v.clear();
-        v.extend(black_box(data.iter()).cloned());
-        v[0]
+        let iter = data.iter().map(|&x| x);
+        v.extend(iter);
+        v[511]
+    });
+    b.bytes = v.capacity() as u64;
+}
+
+fn extend_with_write(b: &mut Bencher) {
+    let mut v = ArrayVec::<[u8; 512]>::new();
+    let data = [1; 512];
+    b.iter(|| {
+        v.clear();
+        v.write(&data[..]).ok();
+        v[511]
     });
     b.bytes = v.capacity() as u64;
 }
@@ -51,5 +67,12 @@ fn extend_from_slice(b: &mut Bencher) {
     b.bytes = v.capacity() as u64;
 }
 
-benchmark_group!(benches, extend_with_constant, extend_with_range, extend_with_slice, extend_from_slice);
+benchmark_group!(benches,
+                 extend_with_constant,
+                 extend_with_range,
+                 extend_with_slice,
+                 extend_with_write,
+                 extend_from_slice
+);
+
 benchmark_main!(benches);

src/lib.rs

@@ -933,6 +933,7 @@ impl<A: Array> Extend<A::Item> for ArrayVec<A> {
         unsafe {
             let len = self.len();
             let mut ptr = self.as_mut_ptr().offset(len as isize);
+            let end_ptr = ptr.offset(take as isize);
             // Keep the length in a separate variable, write it back on scope
             // exit. To help the compiler with alias analysis and stuff.
             // We update the length to handle panic in the iteration of the
@@ -944,10 +945,16 @@ impl<A: Array> Extend<A::Item> for ArrayVec<A> {
                     **self_len = Index::from(len);
                 }
             };
-            for elt in iter.into_iter().take(take) {
-                ptr::write(ptr, elt);
-                ptr = ptr.offset(1);
-                guard.data += 1;
+            let mut iter = iter.into_iter();
+            loop {
+                if ptr == end_ptr { break; }
+                if let Some(elt) = iter.next() {
+                    ptr::write(ptr, elt);
+                    ptr = ptr.offset(1);
+                    guard.data += 1;
+                } else {
+                    break;
+                }
             }
         }
     }