extend_from_slice.rs

#![feature(test)]

extern crate test;

use test::Bencher;

use std::ptr;

trait VecExt<T> {
    fn x_extend_from_slice(&mut self, other: &[T]);
    fn x_extend<I>(&mut self, iter: I) where I: IntoIterator<Item=T>;
}

impl<T> VecExt<T> for Vec<T> where T: Clone {
    fn x_extend_from_slice(&mut self, other: &[T]) {
        self.reserve(other.len());

        // Unsafe code so this can be optimised to a memcpy (or something
        // similarly fast) when T is Copy. LLVM is easily confused, so any
        // extra operations during the loop can prevent this optimisation.
        unsafe {
            let len = self.len();
            let ptr = self.get_unchecked_mut(len) as *mut T;
            // Use SetLenOnDrop to work around bug where compiler
            // may not realize the store through `ptr` trough self.set_len()
            // don't alias.
            let mut local_len = SetLenOnDrop::new(self);

            for i in 0..other.len() {
                ptr::write(ptr.offset(i as isize), other.get_unchecked(i).clone());
                local_len.increment_len(1);
            }

            // len set by scope guard
        }
    }
    fn x_extend<I>(&mut self, iter: I) where I: IntoIterator<Item=T> {
        self.spec_extend(iter.into_iter())
    }
}

struct SetLenOnDrop<T> {
    len: *mut Vec<T>,
    local_len: usize,
}

impl<T> SetLenOnDrop<T> {
    #[inline]
    fn new(len: &mut Vec<T>) -> Self {
        SetLenOnDrop { local_len: len.len(), len: len }
    }

    #[inline]
    fn increment_len(&mut self, increment: usize) {
        self.local_len += increment;
    }
}

impl<T> Drop for SetLenOnDrop<T> {
    #[inline]
    fn drop(&mut self) {
        unsafe {
            (*self.len).set_len(self.local_len);
        }
    }
}


// Specialization trait used for Vec::from_iter and Vec::extend
trait SpecExtend<I> {
    fn from_iter(iter: I) -> Self;
    fn spec_extend(&mut self, iter: I);
}

impl<I, T> SpecExtend<I> for Vec<T>
    where I: Iterator<Item=T>,
{
    fn from_iter(iterator: I) -> Self {
        let mut vector = Vec::new();
        vector.spec_extend(iterator);
        vector
    }

    fn spec_extend(&mut self, iterator: I) {
        // This is the case for a TrustedLen iterator.
        let (low, high) = iterator.size_hint();
        if let Some(high_value) = high {
            debug_assert_eq!(low, high_value,
                             "TrustedLen iterator's size hint is not exact: {:?}",
                             (low, high));
        }
        // None => capacity overflow
        let additional = high.unwrap_or(!0);
        self.reserve(additional);
        unsafe {
            let mut ptr = self.as_mut_ptr().offset(self.len() as isize);
            let mut local_len = SetLenOnDrop::new(self);
            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);
            }
        }
    }
}


#[bench]
fn test_extend_1(b: &mut Bencher) {
    let mut v = vec![0u8; 0];
    let data = [1; 16];

    b.iter(|| {
        v.clear();
        v.x_extend_from_slice(&data);
        v.len()
    })
}

#[bench]
fn test_extend_2(b: &mut Bencher) {
    let mut v = vec![0u8; 0];
    let data = [1; 16];

    b.iter(|| {
        v.clear();
        v.x_extend(data.iter().cloned());
        v.len()
    })
}