fast-find 0.1.0

Find fast finds bytes fast
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139

#[cfg(target_arch = "x86")]
use std::arch::x86::*;
#[cfg(target_arch = "x86_64")]
use std::arch::x86_64::*;
use std::mem::transmute;

use std::is_x86_feature_detected;

use crate::fallback;

pub use crate::fallback::find_in_4;

#[inline]
pub fn find_in_16(needle: u8, haystack: &[u8; 16], len: usize) -> Option<usize> {
    if is_x86_feature_detected!("avx") {
        let bitfield = unsafe { avx_eq_16(needle, haystack.as_ptr()) };
        find_in_bitfield(bitfield, len)
    } else {
        fallback::find_in_16(needle, haystack, len)
    }
}

#[inline]
pub fn find_in_32(needle: u8, haystack: &[u8; 32], len: usize) -> Option<usize> {
    let bitfield = if is_x86_feature_detected!("avx2") {
        unsafe { avx2_eq_32(needle, haystack.as_ptr()) }
    } else if is_x86_feature_detected!("avx") {
        unsafe { avx_eq_32(needle, haystack.as_ptr()) }
    } else {
        return fallback::find_in_32(needle, haystack, len);
    };

    find_in_bitfield(bitfield, len)
}

/// Finds the index of the last 1 in the bitfield, up-to len
///
/// Combines with avx_find_in methods
#[inline]
fn find_in_bitfield(bitfield: i32, len: usize) -> Option<usize> {
    let idx = bitfield.trailing_zeros() as usize;
    if len > idx {
        Some(idx)
    } else {
        None
    }
}

/// Returns a bitfield, where all elements in haystack that are equal to needle = 1, else 0
///
/// Safety: requires avx, 16 byte haystack
#[inline]
unsafe fn avx_eq_16(needle: u8, haystack: *const u8) -> i32 {
    let haystack = _mm_loadu_si128(haystack as *const __m128i);
    let needle_vec = _mm_set1_epi8(transmute::<u8, i8>(needle));
    let eq = _mm_cmpeq_epi8(needle_vec, haystack);
    _mm_movemask_epi8(eq)
}

/// Returns a bitfield, where all elements in haystack that are equal to needle = 1, else 0
///
/// Safety: requires avx, 32 byte haystack
#[inline]
unsafe fn avx_eq_32(needle: u8, haystack: *const u8) -> i32 {
    avx_eq_16(needle, haystack) | (avx_eq_16(needle, haystack.offset(16)) << 16)
}

/// Returns a bitfield, where all elements in haystack that are equal to needle = 1, else 0
///
/// Safety: requires avx2, 32 byte haystack
#[inline]
unsafe fn avx2_eq_32(needle: u8, haystack: *const u8) -> i32 {
    let haystack = _mm256_loadu_si256(haystack as *const __m256i);
    let needle_vec = _mm256_set1_epi8(transmute::<u8, i8>(needle));
    let eq = _mm256_cmpeq_epi8(needle_vec, haystack);
    _mm256_movemask_epi8(eq)
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_avx_find_16() {
        fn test_all(needle: u8, haystack: &[u8; 16], len: usize, expected: Option<usize>) {
            assert_eq!(
                find_in_bitfield(unsafe { avx_eq_16(needle, haystack.as_ptr()) }, len),
                expected
            );
            assert_eq!(find_in_16(needle, haystack, len), expected);
        }

        let mut array = [0; 16];

        test_all(4, &array, 16, None);

        array[4] = 4;
        array[6] = 10;

        test_all(4, &array, 16, Some(4));
        test_all(5, &array, 16, None);
        test_all(10, &array, 16, Some(6));

        test_all(4, &array, 5, Some(4));
        test_all(5, &array, 5, None);
        test_all(10, &array, 5, None);
    }

    #[test]
    fn test_avx_find_32() {
        fn test_all(needle: u8, haystack: &[u8; 32], len: usize, expected: Option<usize>) {
            assert_eq!(
                find_in_bitfield(unsafe { avx2_eq_32(needle, haystack.as_ptr()) }, len),
                expected
            );
            assert_eq!(
                find_in_bitfield(unsafe { avx_eq_32(needle, haystack.as_ptr()) }, len),
                expected
            );
            assert_eq!(find_in_32(needle, haystack, len), expected);
        }

        let mut array = [0; 32];

        test_all(4, &array, 32, None);

        array[4] = 4;
        array[6] = 10;
        array[20] = 5;

        test_all(4, &array, 32, Some(4));
        test_all(5, &array, 32, Some(20));
        test_all(10, &array, 32, Some(6));

        test_all(4, &array, 5, Some(4));
        test_all(5, &array, 5, None);
        test_all(10, &array, 5, None);
    }
}