Add a push_light method and use likely/unlikely intrinsics

Cargo.toml

@@ -16,6 +16,8 @@ union = []
 default = ["std"]
 specialization = []
 may_dangle = []
+likely = []
+push_light = []
 
 [lib]
 name = "smallvec"

benches/bench.rs

@@ -9,10 +9,12 @@ use smallvec::{ExtendFromSlice, SmallVec};
 
 const VEC_SIZE: usize = 16;
 const SPILLED_SIZE: usize = 100;
+const LARGE_SIZE: usize = 10000;
 
 trait Vector<T>: for<'a> From<&'a [T]> + Extend<T> + ExtendFromSlice<T> {
     fn new() -> Self;
     fn push(&mut self, val: T);
+    fn push_light(&mut self, val: T);
     fn pop(&mut self) -> Option<T>;
     fn remove(&mut self, p: usize) -> T;
     fn insert(&mut self, n: usize, val: T);
@@ -29,6 +31,8 @@ impl<T: Copy> Vector<T> for Vec<T> {
         self.push(val)
     }
 
+    fn push_light(&mut self, val: T) {}
+
     fn pop(&mut self) -> Option<T> {
         self.pop()
     }
@@ -59,6 +63,14 @@ impl<T: Copy> Vector<T> for SmallVec<[T; VEC_SIZE]> {
         self.push(val)
     }
 
+    #[cfg(feature = "push_light")]
+    fn push_light(&mut self, val: T) {
+        self.push_light(val)
+    }
+
+    #[cfg(not(feature = "push_light"))]
+    fn push_light(&mut self, _: T) {}
+
     fn pop(&mut self) -> Option<T> {
         self.pop()
     }
@@ -93,8 +105,12 @@ macro_rules! make_benches {
 
 make_benches! {
     SmallVec<[u64; VEC_SIZE]> {
-        bench_push => gen_push(SPILLED_SIZE as _),
+        bench_push_light_small => gen_push_light(VEC_SIZE as _),
+        bench_push_light => gen_push_light(SPILLED_SIZE as _),
+        bench_push_light_large => gen_push_light(LARGE_SIZE as _),
         bench_push_small => gen_push(VEC_SIZE as _),
+        bench_push => gen_push(SPILLED_SIZE as _),
+        bench_push_large => gen_push(LARGE_SIZE as _),
         bench_insert => gen_insert(SPILLED_SIZE as _),
         bench_insert_small => gen_insert(VEC_SIZE as _),
         bench_remove => gen_remove(SPILLED_SIZE as _),
@@ -150,6 +166,21 @@ fn gen_push<V: Vector<u64>>(n: u64, b: &mut Bencher) {
     });
 }
 
+fn gen_push_light<V: Vector<u64>>(n: u64, b: &mut Bencher) {
+    #[inline(never)]
+    fn push_light_noinline<V: Vector<u64>>(vec: &mut V, x: u64) {
+        vec.push_light(x);
+    }
+
+    b.iter(|| {
+        let mut vec = V::new();
+        for x in 0..n {
+            push_light_noinline(&mut vec, x);
+        }
+        vec
+    });
+}
+
 fn gen_insert<V: Vector<u64>>(n: u64, b: &mut Bencher) {
     #[inline(never)]
     fn insert_noinline<V: Vector<u64>>(vec: &mut V, p: usize, x: u64) {

lib.rs

@@ -31,6 +31,9 @@
 #![cfg_attr(not(feature = "std"), no_std)]
 #![cfg_attr(not(feature = "std"), feature(alloc))]
 #![cfg_attr(feature = "union", feature(untagged_unions))]
+#![cfg_attr(feature = "likely", feature(stmt_expr_attributes))]
+#![cfg_attr(feature = "push_light", feature(nll))]
+#![cfg_attr(any(feature = "likely", feature = "push_light"), feature(core_intrinsics))]
 #![cfg_attr(feature = "specialization", feature(specialization))]
 #![cfg_attr(feature = "may_dangle", feature(dropck_eyepatch))]
 #![deny(missing_docs)]
@@ -145,6 +148,26 @@ macro_rules! debug_unreachable {
     }
 }
 
+#[cfg(not(feature = "likely"))]
+macro_rules! likely {
+    ($e:expr) => { $e }
+}
+
+#[cfg(feature = "likely")]
+macro_rules! likely {
+    ($e:expr) => { #[allow(unused_unsafe)] { unsafe { std::intrinsics::likely($e) }} }
+}
+
+#[cfg(not(feature = "likely"))]
+macro_rules! unlikely {
+    ($e:expr) => { $e }
+}
+
+#[cfg(feature = "likely")]
+macro_rules! unlikely {
+    ($e:expr) => { #[allow(unused_unsafe)] { unsafe { std::intrinsics::unlikely($e) }} }
+}
+
 /// Common operations implemented by both `Vec` and `SmallVec`.
 ///
 /// This can be used to write generic code that works with both `Vec` and `SmallVec`.
@@ -611,7 +634,7 @@ impl<A: Array> SmallVec<A> {
     pub fn push(&mut self, value: A::Item) {
         unsafe {
             let (_, &mut len, cap) = self.triple_mut();
-            if len == cap {
+            if unlikely!(len == cap) {
                 self.reserve(1);
             }
             let (ptr, len_ptr, _) = self.triple_mut();
@@ -620,6 +643,61 @@ impl<A: Array> SmallVec<A> {
         }
     }
 
+    /// Append an item to the vector. This is always inlined with a fast
+    /// path for when the vector doesn't need an heap allocation.
+    #[cfg(feature = "push_light")]
+    #[inline(always)]
+    pub fn push_light(&mut self, value: A::Item) {
+        unsafe {
+            if likely!(self.capacity < A::size()) {
+                let ptr = self.data.inline_mut().ptr_mut();
+                ptr::write(ptr.offset(self.capacity as isize), value);
+                self.capacity = self.capacity + 1;
+            } else {
+                self.push_light_cold(self.capacity, value);
+            }
+        }
+    }
+
+    // Slow path
+    #[cfg(feature = "push_light")]
+    #[inline(never)]
+    #[cold]
+    unsafe fn push_light_cold(&mut self, cap: usize, value: A::Item) {
+        std::intrinsics::assume(self.capacity == cap);
+        if likely!(cap != A::size()) {
+            debug_assert!(self.spilled());
+            let &mut (ptr, ref mut len_ptr) = self.data.heap_mut();
+
+            let len = *len_ptr;
+
+            if unlikely!(cap - len < 1) {
+                std::intrinsics::assume(self.capacity >= A::size());
+                if unlikely!(cap > (isize::max_value() >> 1) as usize) {
+                    panic!("size overflow")
+                }
+                let new_cap = cap << 1;
+                self.grow(new_cap);
+                let &mut (ptr, ref mut len_ptr) = self.data.heap_mut();
+                *len_ptr = len + 1;
+                ptr::write(ptr.offset(len as isize), value);
+            } else {
+                *len_ptr = len + 1;
+                ptr::write(ptr.offset(len as isize), value);
+            }
+        } else {
+            debug_assert!(self.len() == A::size());
+            let new_cap = A::size().checked_add(1).
+                and_then(usize::checked_next_power_of_two).
+                unwrap_or(usize::max_value());
+            self.grow(new_cap);
+            debug_assert!(self.spilled());
+            let &mut (ptr, ref mut len_ptr) = self.data.heap_mut();
+            *len_ptr = A::size() + 1;
+            ptr::write(ptr.offset(A::size() as isize), value);
+        }
+    }
+
     /// Remove an item from the end of the vector and return it, or None if empty.
     #[inline]
     pub fn pop(&mut self) -> Option<A::Item> {

testing/Cargo.toml

@@ -0,0 +1,14 @@
+[package]
+name = "testing"
+version = "0.1.0"
+authors = ["John Kåre Alsaker <[email protected]>"]
+edition = "2018"
+
+[lib]
+crate-type = ["dylib"]
+
+[dependencies]
+smallvec = { path = ".." }
+
+[profile.release]
+debug = true

testing/src/lib.rs

@@ -0,0 +1,81 @@
+#![feature(core_intrinsics)]
+
+use smallvec::{SmallVec, Array};
+use std::ptr;
+
+#[cfg(not(feature = "likely"))]
+macro_rules! likely {
+    ($e:expr) => { $e }
+}
+
+#[cfg(feature = "likely")]
+macro_rules! likely {
+    ($e:expr) => { #[allow(unused_unsafe)] { unsafe { std::intrinsics::likely($e) }} }
+}
+
+#[cfg(not(feature = "likely"))]
+macro_rules! unlikely {
+    ($e:expr) => { $e }
+}
+
+#[cfg(feature = "likely")]
+macro_rules! unlikely {
+    ($e:expr) => { #[allow(unused_unsafe)] { unsafe { std::intrinsics::unlikely($e) }} }
+}
+
+pub trait Test<A: Array> {
+    fn push_light(&mut self, v: A::Item);
+}
+
+#[inline(never)]
+#[cold]
+fn push_light_cold<A: Array>(vec: &mut SmallVec<A>, v: A::Item) {
+    if likely!(vec.spilled()) {
+        let len = vec.len();
+        if likely!(len < vec.capacity()) {
+            unsafe {
+                ptr::write(vec.as_mut_ptr().offset(len as isize), v);
+                vec.set_len(len + 1);
+            }
+        } else {
+            vec.push_old(v)
+        }
+    } else {
+        unsafe {
+            std::intrinsics::assume(vec.capacity() == A::size());
+        }
+        vec.push_old(v)
+    }
+}
+
+impl<A: Array> Test<A> for SmallVec<A> {
+    #[inline(always)]
+    fn push_light(&mut self, v: A::Item) {
+        let (free, len) = if !self.spilled() {
+            let len = self.len();
+            (len < A::size(), len)
+        } else {
+            (false, 0)
+        };
+        if likely!(free) {
+            unsafe {
+                ptr::write(self.as_mut_ptr().offset(len as isize), v);
+                std::intrinsics::assume(!self.spilled());
+                self.set_len(self.len() + 1);
+            }
+        } else {
+            push_light_cold(self, v);
+        }
+    }
+}
+
+// FIXME: Remove
+#[no_mangle]
+pub fn test1(a: &mut SmallVec<[u32; 8]>) {
+    a.push(5);
+}
+
+#[no_mangle]
+pub fn test2(a: &mut SmallVec<[u32; 8]>) {
+    a.push_light(5);
+}