Add documentation for AVX2 and memchr searchers

Author: zakcutner

src/avx2/mod.rs

@@ -8,6 +8,8 @@ pub use self::{original::*, rust::*};
 use crate::{bits, memchr::MemchrSearcher, memcmp};
 use std::{arch::x86_64::*, mem};
 
+/// Rolling hash for the simple Rabin-Karp implementation. As a hashing function, the XOR of all the
+/// bytes is computed.
 #[derive(Clone, Copy, Default, PartialEq)]
 struct ScalarHash(usize);
 
@@ -33,6 +35,8 @@ impl ScalarHash {
     }
 }
 
+/// Represents an SIMD register type that is x86-specific (but could be used more generically) in
+/// order to share functionality between SSE2, AVX2 and possibly future implementations.
 trait Vector: Copy {
     unsafe fn set1_epi8(a: i8) -> Self;
 
@@ -109,6 +113,9 @@ impl Vector for __m256i {
     }
 }
 
+/// Hash of the first and "last" bytes in the needle for use with the SIMD algorithm implemented by
+/// `Avx2Searcher::vector_search_in`. As explained, any byte can be chosen to represent the "last"
+/// byte of the hash to prevent worst-case attacks.
 struct VectorHash<V: Vector> {
     first: V,
     last: V,
@@ -126,6 +133,30 @@ impl<V: Vector> VectorHash<V> {
 
 macro_rules! avx2_searcher {
     ($name:ident, $size:literal, $memcmp:path) => {
+        /// Single-substring searcher using an AVX2 algorithm based on the "Generic SIMD" algorithm
+        /// [presented by Wojciech Muła](http://0x80.pl/articles/simd-strfind.html).
+        ///
+        /// It is similar to the Rabin-Karp algorithm, except that the hash is not rolling and is
+        /// calculated for several lanes at once. It begins by picking the first byte in the needle
+        /// and checking at which positions in the haystack it occurs. Any position where it does
+        /// not can be immediately discounted as a potential match.
+        ///
+        /// We then repeat this idea with a second byte in the needle (where the haystack is
+        /// suitably offset) and take a bitwise AND to further limit the possible positions the
+        /// needle can match in. Any remaining positions are fully evaluated using an equality
+        /// comparison with the needle.
+        ///
+        /// Originally, the algorithm always used the last byte for this second byte. Whilst this is
+        /// often the most efficient option, it is vulnerable to a worst-case attack and so this
+        /// implementation instead allows any byte (including a random one) to be chosen.
+        ///
+        /// In the case where the needle is not a multiple of the number of SIMD lanes, the last
+        /// chunk is made up of a partial overlap with the penultimate chunk to avoid reading random
+        /// memory, differing from the original implementation. In this case, a mask is used to
+        /// prevent performing an equality comparison on the same position twice.
+        ///
+        /// When the haystack is too short for an AVX2 register, a similar SSE2 fallback is used
+        /// instead. Finally, for very short haystacks there is a scalar Rabin-Karp implementation.
         pub struct $name {
             needle: Box<[u8]>,
             position: usize,
@@ -135,12 +166,15 @@ macro_rules! avx2_searcher {
         }
 
         impl $name {
+            /// Creates a new searcher for `needle`. By default, `position` is set to the last
+            /// character in the needle.
             #[target_feature(enable = "avx2")]
             pub unsafe fn new(needle: Box<[u8]>) -> Self {
                 let position = needle.len() - 1;
                 Self::with_position(needle, position)
             }
 
+            /// Same as `new` but allows additionally specifying the `position` to use.
             #[target_feature(enable = "avx2")]
             pub unsafe fn with_position(needle: Box<[u8]>, position: usize) -> Self {
                 assert!(!needle.is_empty());
@@ -270,6 +304,7 @@ macro_rules! avx2_searcher {
                 self.vector_search_in(haystack, &self.avx2_hash, Self::sse2_search_in)
             }
 
+            /// Inlined version of `search_in` for hot call sites.
             #[inline]
             #[target_feature(enable = "avx2")]
             pub unsafe fn inlined_search_in(&self, haystack: &[u8]) -> bool {
@@ -280,6 +315,7 @@ macro_rules! avx2_searcher {
                 self.avx2_search_in(haystack)
             }
 
+            /// Performs a substring search for the `needle` within `haystack`.
             #[inline]
             pub fn search_in(&self, haystack: &[u8]) -> bool {
                 unsafe { self.inlined_search_in(haystack) }
@@ -302,6 +338,12 @@ avx2_searcher!(Avx2Searcher11, 11, memcmp::memcmp10);
 avx2_searcher!(Avx2Searcher12, 12, memcmp::memcmp11);
 avx2_searcher!(Avx2Searcher13, 13, memcmp::memcmp12);
 
+/// Single-substring searcher based on `Avx2Searcher` but with dynamic algorithm selection.
+///
+/// It has specialized cases for zero-length needles, which are found in all haystacks, and
+/// one-length needles, which uses `MemchrSearcher`. For needles up to a length of thirteen it uses
+/// specialized versions of `Avx2Searcher`, finally falling back to the generic version of
+/// `Avx2Searcher` for longer needles.
 pub enum DynamicAvx2Searcher {
     N0,
     N1(MemchrSearcher),
@@ -321,12 +363,15 @@ pub enum DynamicAvx2Searcher {
 }
 
 impl DynamicAvx2Searcher {
+    /// Creates a new searcher for `needle`. By default, `position` is set to the last character in
+    /// the needle.
     #[target_feature(enable = "avx2")]
     pub unsafe fn new(needle: Box<[u8]>) -> Self {
         let position = needle.len() - 1;
         Self::with_position(needle, position)
     }
 
+    /// Same as `new` but allows additionally specifying the `position` to use.
     #[target_feature(enable = "avx2")]
     pub unsafe fn with_position(needle: Box<[u8]>, position: usize) -> Self {
         assert!(!needle.is_empty());
@@ -351,6 +396,7 @@ impl DynamicAvx2Searcher {
         }
     }
 
+    /// Inlined version of `search_in` for hot call sites.
     #[inline]
     #[target_feature(enable = "avx2")]
     pub unsafe fn inlined_search_in(&self, haystack: &[u8]) -> bool {
@@ -373,6 +419,7 @@ impl DynamicAvx2Searcher {
         }
     }
 
+    /// Performs a substring search for the `needle` within `haystack`.
     #[inline]
     pub fn search_in(&self, haystack: &[u8]) -> bool {
         unsafe { self.inlined_search_in(haystack) }

src/memchr.rs

@@ -1,12 +1,15 @@
 use memchr::memchr;
 
+/// Single-byte searcher using `memchr` for faster matching.
 pub struct MemchrSearcher(u8);
 
 impl MemchrSearcher {
+    /// Creates a new searcher for `needle`.
     pub fn new(needle: u8) -> Self {
         Self(needle)
     }
 
+    /// Inlined version of `search_in` for hot call sites.
     #[inline]
     pub fn inlined_search_in(&self, haystack: &[u8]) -> bool {
         if haystack.is_empty() {
@@ -16,6 +19,7 @@ impl MemchrSearcher {
         memchr(self.0, haystack).is_some()
     }
 
+    /// Performs a substring search for the `needle` within `haystack`.
     pub fn search_in(&self, haystack: &[u8]) -> bool {
         self.inlined_search_in(haystack)
     }