lib.rs

/// An incomplete rusty simd implementation of memchr(3)
#![feature(stdsimd, exact_chunks)]
#![cfg_attr(test, feature(test))]
#![cfg_attr(feature = "cargo-clippy", allow(clippy_pedantic))]
#[cfg(test)]
extern crate memchr as other_memchr;
#[cfg(test)]
extern crate quickcheck;
#[cfg(test)]
extern crate test;

use std::arch::x86_64::*;
use std::simd::*;

/// Helper function to see if a needle is in a haystack.
///
/// Checks to see if a byte `n` exists in a byte array `haystack`.
/// It is a logic error for a caller to supply a `haystack` that is not of length 16.
/// Callers are expected to pad the haystack as necessary to satisfy the length requirement.
///
/// Returns an offset into `haystack` if a match was found.
/// If no match was found 16 will be returned.
#[inline(always)]
#[cfg(target_arch = "x86_64")]
#[cfg(target_feature = "sse4.2")]
pub fn find_byte(n: u8, haystack: &[u8]) -> usize {
    // We pad the needle to a &[u8; 16] because cmpistri implicitly discovers the length
    // and has a lower latency than explicitly providing the lengths
    let mut needle = [0; 16];
    needle[0] = n;
    unsafe {
        let needle = u8x16::load_aligned(&needle);
        let haystack = u8x16::load_aligned(haystack);
        _mm_cmpistri(needle.into_bits(), haystack.into_bits(), 0) as usize
    }
}

/// A simplistic implementation that tries to loosely keep the semantics of `memchr(3)`
/// except we do not consider strings to be null-terminated nor non-ascii. This is both to keep the
/// implementation simple (e.g. cmpistri (and the other string comparison instructions) expect null
/// bytes to represent "no value", instead of the value's termination)
/// and because rust does not have equivalent string semantics (excluding std::ffi::* constructs).
///
/// Checks to see if a byte `needle` is a substring of a byte array `haystack`.
/// In the largest deviation from `memchr(3)`, it is a logic error to provide a `haystack` with a
/// length that is not divisible by 16.
///
/// Returns `Option<usize>` with `Some<usize>` denoting a match as well as the offset into `haystack`,
/// `None` signifys that no match was found.
#[inline(always)]
#[cfg(target_arch = "x86_64")]
#[cfg(target_feature = "sse4.2")]
pub fn memchr(needle: u8, haystack: &[u8]) -> Option<usize> {
    if haystack.is_empty() {
        return None;
    }

    let mut result = 0;

    for (i, chunk) in haystack.exact_chunks(chunk_size).enumerate() {
        result = find_byte(needle, &chunk);
        if result != 16 {
            // calculate offset
            return Some((i << 4) + result as usize);
        }
    }

    None
}

#[cfg(test)]
mod tests {
    // borrowed from the memchr crate's tests
    use super::*;
    use other_memchr::memchr as libc_memchr;
    use quickcheck;
    use test::Bencher;

    #[test]
    fn returns_some_on_zero() {
        let data = vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16];
        for byte in 0..256u32 {
            let byte = byte as u8;
            let control = data.iter().position(|elt| *elt == byte);
            let predicate = memchr(byte, &data);
            println!("{:?}, {:?}", control, predicate);
            println!("{:?}, {:?}", byte, data);
            assert!(control == predicate);
        }
    }

    #[test]
    fn qc_never_fail() {
        fn prop(needle: u8, haystack: Vec<u8>) -> bool {
            memchr(needle, &haystack);
            true
        }
        quickcheck::quickcheck(prop as fn(u8, Vec<u8>) -> bool);
    }

    #[test]
    fn qc_correct() {
        fn prop(v: Vec<u8>) -> bool {
            if v.iter().any(|x| *x == 0) {
                return true;
            }
            for byte in 0..256u32 {
                let byte = byte as u8;
                let control = v.iter().position(|elt| *elt == byte);
                let predicate = memchr(byte, &v);
                if predicate != control {
                    return false;
                }
            }
            true
        }
        quickcheck::quickcheck(prop as fn(Vec<u8>) -> bool);
    }

    fn bench_data() -> Vec<u8> {
        let mut v: Vec<u8> = ::std::iter::repeat(b'z').take(10000000).collect();
        let len = v.len();
        v[len / 2] = b'a';
        v
    }

    #[bench]
    fn simd_memchr(b: &mut Bencher) {
        let haystack = bench_data();
        let needle = b'a';
        b.iter(|| {
            assert!(super::memchr(needle, &haystack).is_some());
        });
    }

    #[bench]
    fn iterator_memchr(b: &mut Bencher) {
        let haystack = bench_data();
        let needle = b'a';
        b.iter(|| {
            assert!(haystack.iter().position(|&b| b == needle).is_some());
        });
    }

    #[bench]
    fn control_memchr(b: &mut Bencher) {
        let haystack = bench_data();
        let needle = b'a';
        b.iter(|| {
            assert!(libc_memchr(needle, &haystack).is_some());
        })
    }
}