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(not(feature = "std"))]
pub(crate) use self::nostd::Forward;
#[cfg(feature = "std")]
pub(crate) use self::std::Forward;

#[cfg(feature = "std")]
mod std {
    use core::arch::x86_64::{__m128i, __m256i};

    use crate::memmem::{genericsimd, NeedleInfo};

    /// An AVX accelerated vectorized substring search routine that only works
    /// on small needles.
    #[derive(Clone, Copy, Debug)]
    pub(crate) struct Forward(genericsimd::Forward);

    impl Forward {
        /// Create a new "generic simd" forward searcher. If one could not be
        /// created from the given inputs, then None is returned.
        pub(crate) fn new(
            ninfo: &NeedleInfo,
            needle: &[u8],
        ) -> Option<Forward> {
            if !cfg!(memchr_runtime_avx) || !is_x86_feature_detected!("avx2") {
                return None;
            }
            genericsimd::Forward::new(ninfo, needle).map(Forward)
        }

        /// Returns the minimum length of haystack that is needed for this
        /// searcher to work. Passing a haystack with a length smaller than
        /// this will cause `find` to panic.
        #[inline(always)]
        pub(crate) fn min_haystack_len(&self) -> usize {
            self.0.min_haystack_len::<__m128i>()
        }

        #[inline(always)]
        pub(crate) fn find(
            &self,
            haystack: &[u8],
            needle: &[u8],
        ) -> Option<usize> {
            // SAFETY: The only way a Forward value can exist is if the avx2
            // target feature is enabled. This is the only safety requirement
            // for calling the genericsimd searcher.
            unsafe { self.find_impl(haystack, needle) }
        }

        /// The implementation of find marked with the appropriate target
        /// feature.
        ///
        /// # Safety
        ///
        /// Callers must ensure that the avx2 CPU feature is enabled in the
        /// current environment.
        #[target_feature(enable = "avx2")]
        unsafe fn find_impl(
            &self,
            haystack: &[u8],
            needle: &[u8],
        ) -> Option<usize> {
            if haystack.len() < self.0.min_haystack_len::<__m256i>() {
                genericsimd::fwd_find::<__m128i>(&self.0, haystack, needle)
            } else {
                genericsimd::fwd_find::<__m256i>(&self.0, haystack, needle)
            }
        }
    }
}

// We still define the avx "forward" type on nostd to make caller code a bit
// simpler. This avoids needing a lot more conditional compilation.
#[cfg(not(feature = "std"))]
mod nostd {
    use crate::memmem::NeedleInfo;

    #[derive(Clone, Copy, Debug)]
    pub(crate) struct Forward(());

    impl Forward {
        pub(crate) fn new(
            ninfo: &NeedleInfo,
            needle: &[u8],
        ) -> Option<Forward> {
            None
        }

        pub(crate) fn min_haystack_len(&self) -> usize {
            unreachable!()
        }

        pub(crate) fn find(
            &self,
            haystack: &[u8],
            needle: &[u8],
        ) -> Option<usize> {
            unreachable!()
        }
    }
}

#[cfg(all(test, feature = "std", not(miri)))]
mod tests {
    use crate::memmem::{prefilter::PrefilterState, NeedleInfo};

    fn find(
        _: &mut PrefilterState,
        ninfo: &NeedleInfo,
        haystack: &[u8],
        needle: &[u8],
    ) -> Option<usize> {
        super::Forward::new(ninfo, needle).unwrap().find(haystack, needle)
    }

    #[test]
    fn prefilter_permutations() {
        use crate::memmem::prefilter::tests::PrefilterTest;

        if !is_x86_feature_detected!("avx2") {
            return;
        }
        // SAFETY: The safety of find only requires that the current CPU
        // support AVX2, which we checked above.
        unsafe {
            PrefilterTest::run_all_tests_filter(find, |t| {
                // This substring searcher only works on certain configs, so
                // filter our tests such that Forward::new will be guaranteed
                // to succeed. (And also remove tests with a haystack that is
                // too small.)
                let fwd = match super::Forward::new(&t.ninfo, &t.needle) {
                    None => return false,
                    Some(fwd) => fwd,
                };
                t.haystack.len() >= fwd.min_haystack_len()
            })
        }
    }
}