brie-tree 0.1.2

A fast B+ Tree implementation that uses integer keys
Documentation
use super::{SimdSearch, exact_div_unchecked};
#[cfg(target_arch = "x86")]
use core::arch::x86::*;
#[cfg(target_arch = "x86_64")]
use core::arch::x86_64::*;

#[inline]
#[target_feature(enable = "avx2")]
unsafe fn search8(keys: *const i8, search: i8) -> usize {
    let search = _mm256_set1_epi8(search);
    let cmp = |offset: usize| {
        let vec = unsafe { _mm256_load_si256(keys.cast::<__m256i>().add(offset)) };
        _mm256_cmpgt_epi8(search, vec)
    };
    let a = cmp(0);
    let b = cmp(1);
    let c = cmp(2);
    let d = cmp(3);
    let mask1 = _mm256_movemask_epi8(a);
    let mask2 = _mm256_movemask_epi8(b);
    let mask3 = _mm256_movemask_epi8(c);
    let mask4 = _mm256_movemask_epi8(d);
    mask1.count_ones() as usize
        + mask2.count_ones() as usize
        + mask3.count_ones() as usize
        + mask4.count_ones() as usize
}

#[inline]
#[target_feature(enable = "avx2")]
unsafe fn search16(keys: *const i16, search: i16) -> usize {
    let search = _mm256_set1_epi16(search);
    let cmp = |offset: usize| {
        let vec = unsafe { _mm256_load_si256(keys.cast::<__m256i>().add(offset)) };
        _mm256_cmpgt_epi16(search, vec)
    };
    let a = cmp(0);
    let b = cmp(1);
    let c = cmp(2);
    let d = cmp(3);
    let ab = _mm256_packs_epi16(a, b);
    let cd = _mm256_packs_epi16(c, d);
    let mask1 = _mm256_movemask_epi8(ab);
    let mask2 = _mm256_movemask_epi8(cd);
    mask1.count_ones() as usize + mask2.count_ones() as usize
}

#[inline]
#[target_feature(enable = "avx2")]
unsafe fn search32(keys: *const i32, search: i32) -> usize {
    let search = _mm256_set1_epi32(search);
    let cmp = |offset: usize| {
        let vec = unsafe { _mm256_load_si256(keys.cast::<__m256i>().add(offset)) };
        _mm256_cmpgt_epi32(search, vec)
    };
    let a = cmp(0);
    let b = cmp(1);
    let c = cmp(2);
    let d = cmp(3);
    let ab = _mm256_blend_epi16(a, b, 0b01010101);
    let cd = _mm256_blend_epi16(c, d, 0b01010101);
    let abcd = _mm256_packs_epi16(ab, cd);
    let mask = _mm256_movemask_epi8(abcd);
    mask.count_ones() as usize
}

#[inline]
#[target_feature(enable = "avx2")]
unsafe fn search64(keys: *const i64, search: i64) -> usize {
    let search = _mm256_set1_epi64x(search);
    let cmp = |offset: usize| {
        let vec = unsafe { _mm256_load_si256(keys.cast::<__m256i>().add(offset)) };
        _mm256_cmpgt_epi64(search, vec)
    };
    let a = cmp(0);
    let b = cmp(1);
    let c = cmp(2);
    let d = cmp(3);
    let ab = _mm256_blend_epi32(a, b, 0b01010101);
    let cd = _mm256_blend_epi32(c, d, 0b01010101);
    let abcd = _mm256_blend_epi16(ab, cd, 0b01010101);
    let mask = _mm256_movemask_epi8(abcd);
    unsafe { exact_div_unchecked(mask.count_ones() as usize, 2) }
}

impl SimdSearch for u8 {
    const SIMD_WIDTH: usize = 128;
    const BIAS: Self = i8::MIN as Self;
    #[inline]
    unsafe fn simd_search(keys: *const Self, search: Self) -> usize {
        unsafe { search8(keys.cast(), search as i8) }
    }
}
impl SimdSearch for u16 {
    const SIMD_WIDTH: usize = 64;
    const BIAS: Self = i16::MIN as Self;
    #[inline]
    unsafe fn simd_search(keys: *const Self, search: Self) -> usize {
        unsafe { search16(keys.cast(), search as i16) }
    }
}
impl SimdSearch for u32 {
    const SIMD_WIDTH: usize = 32;
    const BIAS: Self = i32::MIN as Self;
    #[inline]
    unsafe fn simd_search(keys: *const Self, search: Self) -> usize {
        unsafe { search32(keys.cast(), search as i32) }
    }
}
impl SimdSearch for u64 {
    const SIMD_WIDTH: usize = 16;
    const BIAS: Self = i64::MIN as Self;
    #[inline]
    unsafe fn simd_search(keys: *const Self, search: Self) -> usize {
        unsafe { search64(keys.cast(), search as i64) }
    }
}
impl SimdSearch for u128 {}
impl SimdSearch for i8 {
    const SIMD_WIDTH: usize = 128;
    #[inline]
    unsafe fn simd_search(keys: *const Self, search: Self) -> usize {
        unsafe { search8(keys, search) }
    }
}
impl SimdSearch for i16 {
    const SIMD_WIDTH: usize = 64;
    #[inline]
    unsafe fn simd_search(keys: *const Self, search: Self) -> usize {
        unsafe { search16(keys, search) }
    }
}
impl SimdSearch for i32 {
    const SIMD_WIDTH: usize = 32;
    #[inline]
    unsafe fn simd_search(keys: *const Self, search: Self) -> usize {
        unsafe { search32(keys, search) }
    }
}
impl SimdSearch for i64 {
    const SIMD_WIDTH: usize = 16;
    #[inline]
    unsafe fn simd_search(keys: *const Self, search: Self) -> usize {
        unsafe { search64(keys, search) }
    }
}
impl SimdSearch for i128 {}