Merge pull request #74 from bluss/improved-extend

Improve .extend() performance

Author: bluss

Cargo.toml

@@ -30,6 +30,11 @@ version = "1.0"
 
 [dev-dependencies]
 matches = { version = "0.1" }
+bencher = "0.1.4"
+
+[[bench]]
+name = "extend"
+harness = false
 
 [features]
 default = ["std"]

benches/extend.rs

@@ -0,0 +1,43 @@
+
+extern crate arrayvec;
+#[macro_use] extern crate bencher;
+
+use arrayvec::ArrayVec;
+
+use bencher::Bencher;
+
+fn extend_with_constant(b: &mut Bencher) {
+    let mut v = ArrayVec::<[u8; 512]>::new();
+    let cap = v.capacity();
+    b.iter(|| {
+        v.clear();
+        v.extend((0..cap).map(|_| 1));
+        v[0]
+    });
+    b.bytes = v.capacity() as u64;
+}
+
+fn extend_with_range(b: &mut Bencher) {
+    let mut v = ArrayVec::<[u8; 512]>::new();
+    let cap = v.capacity();
+    b.iter(|| {
+        v.clear();
+        v.extend((0..cap).map(|x| x as _));
+        v[0]
+    });
+    b.bytes = v.capacity() as u64;
+}
+
+fn extend_with_slice(b: &mut Bencher) {
+    let mut v = ArrayVec::<[u8; 512]>::new();
+    let data = [1; 512];
+    b.iter(|| {
+        v.clear();
+        v.extend(data.iter().cloned());
+        v[0]
+    });
+    b.bytes = v.capacity() as u64;
+}
+
+benchmark_group!(benches, extend_with_constant, extend_with_range, extend_with_slice);
+benchmark_main!(benches);

src/lib.rs

@@ -805,6 +805,21 @@ impl<'a, A: Array> Drop for Drain<'a, A>
     }
 }
 
+struct ScopeExitGuard<T, Data, F>
+    where F: FnMut(&Data, &mut T)
+{
+    value: T,
+    data: Data,
+    f: F,
+}
+
+impl<T, Data, F> Drop for ScopeExitGuard<T, Data, F>
+    where F: FnMut(&Data, &mut T)
+{
+    fn drop(&mut self) {
+        (self.f)(&self.data, &mut self.value)
+    }
+}
 
 
 
@@ -815,9 +830,26 @@ impl<'a, A: Array> Drop for Drain<'a, A>
 impl<A: Array> Extend<A::Item> for ArrayVec<A> {
     fn extend<T: IntoIterator<Item=A::Item>>(&mut self, iter: T) {
         let take = self.capacity() - self.len();
-        for elt in iter.into_iter().take(take) {
-            unsafe {
-                self.push_unchecked(elt);
+        unsafe {
+            let len = self.len();
+            let mut ptr = self.as_mut_ptr().offset(len 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
+            // user's iterator, without dropping any elements on the floor.
+            let mut guard = ScopeExitGuard {
+                value: self,
+                data: len,
+                f: |&len, self_| {
+                    unsafe {
+                        self_.set_len(len)
+                    }
+                }
+            };
+            for elt in iter.into_iter().take(take) {
+                ptr::write(ptr, elt);
+                ptr = ptr.offset(1);
+                guard.data += 1;
             }
         }
     }