Add generic & safe (but unstable) implementation using std::simd

Author: marmeladema

Cargo.toml

@@ -19,3 +19,6 @@ cfg-if = "1"
 
 [profile.release]
 debug = true
+
+[features]
+stdsimd = []

src/aarch64.rs

@@ -26,6 +26,7 @@ static MD: [u8; 16] = [
 
 impl Vector for uint8x16_t {
     const LANES: usize = 16;
+    type Mask = Self;
 
     #[inline]
     unsafe fn splat(a: u8) -> Self {
@@ -58,6 +59,7 @@ impl Vector for uint8x16_t {
 
 impl Vector for uint8x8_t {
     const LANES: usize = 8;
+    type Mask = Self;
 
     #[inline]
     unsafe fn splat(a: u8) -> Self {
@@ -96,6 +98,7 @@ struct uint8x4_t(uint8x8_t);
 
 impl Vector for uint8x4_t {
     const LANES: usize = 4;
+    type Mask = Self;
 
     #[inline]
     unsafe fn splat(a: u8) -> Self {
@@ -136,6 +139,7 @@ struct uint8x2_t(uint8x8_t);
 
 impl Vector for uint8x2_t {
     const LANES: usize = 2;
+    type Mask = Self;
 
     #[inline]
     unsafe fn splat(a: u8) -> Self {

src/lib.rs

@@ -11,11 +11,16 @@
     allow(stable_features),
     feature(aarch64_target_feature)
 )]
+#![cfg_attr(feature = "stdsimd", feature(portable_simd))]
 
 /// Substring search implementations using aarch64 architecture features.
 #[cfg(target_arch = "aarch64")]
 pub mod aarch64;
 
+/// Substring search implementations using generic stdsimd features.
+#[cfg(feature = "stdsimd")]
+pub mod stdsimd;
+
 /// Substring search implementations using x86 architecture features.
 #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
 pub mod x86;
@@ -151,21 +156,24 @@ impl MemchrSearcher {
 trait Vector: Copy {
     const LANES: usize;
 
+    type Mask;
+
     unsafe fn splat(a: u8) -> Self;
 
     unsafe fn load(a: *const u8) -> Self;
 
-    unsafe fn lanes_eq(a: Self, b: Self) -> Self;
+    unsafe fn lanes_eq(a: Self, b: Self) -> Self::Mask;
 
-    unsafe fn bitwise_and(a: Self, b: Self) -> Self;
+    unsafe fn bitwise_and(a: Self::Mask, b: Self::Mask) -> Self::Mask;
 
-    unsafe fn to_bitmask(a: Self) -> i32;
+    unsafe fn to_bitmask(a: Self::Mask) -> i32;
 }
 
 /// 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.
+#[derive(Debug)]
 struct VectorHash<V: Vector> {
     first: V,
     last: V,
@@ -174,8 +182,8 @@ struct VectorHash<V: Vector> {
 impl<V: Vector> VectorHash<V> {
     unsafe fn new(first: u8, last: u8) -> Self {
         Self {
-            first: Vector::splat(first),
-            last: Vector::splat(last),
+            first: V::splat(first),
+            last: V::splat(last),
         }
     }
 }
@@ -206,14 +214,14 @@ trait Searcher<N: NeedleWithSize + ?Sized> {
         start: *const u8,
         mask: i32,
     ) -> bool {
-        let first = Vector::load(start);
-        let last = Vector::load(start.add(self.position()));
+        let first = V::load(start);
+        let last = V::load(start.add(self.position()));
 
-        let eq_first = Vector::lanes_eq(hash.first, first);
-        let eq_last = Vector::lanes_eq(hash.last, last);
+        let eq_first = V::lanes_eq(hash.first, first);
+        let eq_last = V::lanes_eq(hash.last, last);
 
-        let eq = Vector::bitwise_and(eq_first, eq_last);
-        let mut eq = (Vector::to_bitmask(eq) & mask) as u32;
+        let eq = V::bitwise_and(eq_first, eq_last);
+        let mut eq = (V::to_bitmask(eq) & mask) as u32;
 
         let start = start as usize - haystack.as_ptr() as usize;
         let chunk = haystack.as_ptr().add(start + 1);
@@ -379,10 +387,19 @@ mod tests {
 
                     let searcher = unsafe { NeonSearcher::with_position(needle, position) };
                     assert_eq!(unsafe { searcher.search_in(haystack) }, result);
-                } else {
+                } else if #[cfg(not(feature = "stdsimd"))] {
                     compile_error!("Unsupported architecture");
                 }
             }
+
+            cfg_if::cfg_if! {
+                if #[cfg(feature = "stdsimd")] {
+                    use crate::stdsimd::StdSimdSearcher;
+
+                    let searcher = StdSimdSearcher::with_position(needle, position);
+                    assert_eq!(searcher.search_in(haystack), result);
+                }
+            }
         }
 
         result

src/stdsimd.rs

@@ -0,0 +1,170 @@
+#![allow(clippy::missing_safety_doc)]
+
+use crate::{Needle, NeedleWithSize, Searcher, Vector, VectorHash};
+#[cfg(feature = "stdsimd")]
+use std::simd::*;
+
+trait ToFixedBitMask: Sized {
+    fn to_fixed_bitmask(self) -> u32;
+}
+
+impl<const LANES: usize> ToFixedBitMask for Mask<i8, LANES>
+where
+    LaneCount<LANES>: SupportedLaneCount,
+    Self: ToBitMask,
+    <Self as ToBitMask>::BitMask: Into<u32>,
+{
+    #[inline]
+    fn to_fixed_bitmask(self) -> u32 {
+        self.to_bitmask().into()
+    }
+}
+
+impl<const LANES: usize> Vector for Simd<u8, LANES>
+where
+    LaneCount<LANES>: SupportedLaneCount,
+    Mask<i8, LANES>: ToFixedBitMask,
+{
+    const LANES: usize = LANES;
+    type Mask = Mask<i8, LANES>;
+
+    #[inline]
+    unsafe fn splat(a: u8) -> Self {
+        Simd::splat(a as u8)
+    }
+
+    #[inline]
+    unsafe fn load(a: *const u8) -> Self {
+        std::ptr::read_unaligned(a as *const Self)
+    }
+
+    #[inline]
+    unsafe fn lanes_eq(a: Self, b: Self) -> Self::Mask {
+        a.lanes_eq(b)
+    }
+
+    #[inline]
+    unsafe fn bitwise_and(a: Self::Mask, b: Self::Mask) -> Self::Mask {
+        a & b
+    }
+
+    #[inline]
+    unsafe fn to_bitmask(a: Self::Mask) -> i32 {
+        std::mem::transmute(a.to_fixed_bitmask())
+    }
+}
+
+type Simd2 = Simd<u8, 2>;
+type Simd4 = Simd<u8, 4>;
+type Simd8 = Simd<u8, 8>;
+type Simd16 = Simd<u8, 16>;
+type Simd32 = Simd<u8, 32>;
+
+fn from_hash<const N1: usize, const N2: usize>(
+    hash: &VectorHash<Simd<u8, N1>>,
+) -> VectorHash<Simd<u8, N2>>
+where
+    LaneCount<N1>: SupportedLaneCount,
+    Mask<i8, N1>: ToFixedBitMask,
+    LaneCount<N2>: SupportedLaneCount,
+    Mask<i8, N2>: ToFixedBitMask,
+{
+    VectorHash {
+        first: Simd::splat(hash.first.as_array()[0]),
+        last: Simd::splat(hash.last.as_array()[0]),
+    }
+}
+
+/// Searcher for portable simd.
+pub struct StdSimdSearcher<N: Needle> {
+    needle: N,
+    position: usize,
+    simd32_hash: VectorHash<Simd32>,
+}
+
+impl<N: Needle> Searcher<N> for StdSimdSearcher<N> {
+    fn needle(&self) -> &N {
+        &self.needle
+    }
+
+    fn position(&self) -> usize {
+        self.position
+    }
+}
+
+impl<N: Needle> StdSimdSearcher<N> {
+    /// Creates a new searcher for `needle`. By default, `position` is set to
+    /// the last character in the needle.
+    ///
+    /// # Panics
+    ///
+    /// Panics if `needle` is empty or if the associated `SIZE` constant does
+    /// not correspond to the actual size of `needle`.
+    pub fn new(needle: N) -> Self {
+        // Wrapping prevents panicking on unsigned integer underflow when
+        // `needle` is empty.
+        let position = needle.size().wrapping_sub(1);
+        Self::with_position(needle, position)
+    }
+
+    /// Same as `new` but allows additionally specifying the `position` to use.
+    ///
+    /// # Panics
+    ///
+    /// Panics if `needle` is empty, if `position` is not a valid index for
+    /// `needle` or if the associated `SIZE` constant does not correspond to the
+    /// actual size of `needle`.
+    #[inline]
+    pub fn with_position(needle: N, position: usize) -> Self {
+        // Implicitly checks that the needle is not empty because position is an
+        // unsized integer.
+        assert!(position < needle.size());
+
+        let bytes = needle.as_bytes();
+        if let Some(size) = N::SIZE {
+            assert_eq!(size, bytes.len());
+        }
+
+        let simd32_hash = unsafe { VectorHash::new(bytes[0], bytes[position]) };
+
+        Self {
+            position,
+            simd32_hash,
+            needle,
+        }
+    }
+
+    /// Inlined version of `search_in` for hot call sites.
+    #[inline]
+    pub fn inlined_search_in(&self, haystack: &[u8]) -> bool {
+        if haystack.len() <= self.needle.size() {
+            return haystack == self.needle.as_bytes();
+        }
+
+        let end = haystack.len() - self.needle.size() + 1;
+
+        if end < Simd2::LANES {
+            unreachable!();
+        } else if end < Simd4::LANES {
+            let hash = from_hash::<32, 2>(&self.simd32_hash);
+            println!("hash: {:?}", hash);
+            unsafe { self.vector_search_in_default_version(haystack, end, &hash) }
+        } else if end < Simd8::LANES {
+            let hash = from_hash::<32, 4>(&self.simd32_hash);
+            unsafe { self.vector_search_in_default_version(haystack, end, &hash) }
+        } else if end < Simd16::LANES {
+            let hash = from_hash::<32, 8>(&self.simd32_hash);
+            unsafe { self.vector_search_in_default_version(haystack, end, &hash) }
+        } else if end < Simd32::LANES {
+            let hash = from_hash::<32, 16>(&self.simd32_hash);
+            unsafe { self.vector_search_in_default_version(haystack, end, &hash) }
+        } else {
+            unsafe { self.vector_search_in_default_version(haystack, end, &self.simd32_hash) }
+        }
+    }
+
+    /// Performs a substring search for the `needle` within `haystack`.
+    pub fn search_in(&self, haystack: &[u8]) -> bool {
+        self.inlined_search_in(haystack)
+    }
+}

src/wasm32.rs

@@ -6,6 +6,7 @@ use std::arch::wasm32::*;
 
 impl Vector for v128 {
     const LANES: usize = 16;
+    type Mask = Self;
 
     #[inline]
     #[target_feature(enable = "simd128")]
@@ -45,6 +46,7 @@ struct v64(v128);
 
 impl Vector for v64 {
     const LANES: usize = 8;
+    type Mask = Self;
 
     #[inline]
     #[target_feature(enable = "simd128")]
@@ -90,6 +92,7 @@ struct v32(v128);
 
 impl Vector for v32 {
     const LANES: usize = 4;
+    type Mask = Self;
 
     #[inline]
     #[target_feature(enable = "simd128")]
@@ -135,6 +138,7 @@ struct v16(v128);
 
 impl Vector for v16 {
     const LANES: usize = 2;
+    type Mask = Self;
 
     #[inline]
     #[target_feature(enable = "simd128")]

src/x86.rs

@@ -30,6 +30,7 @@ struct __m16i(__m128i);
 
 impl Vector for __m16i {
     const LANES: usize = 2;
+    type Mask = Self;
 
     #[inline]
     #[target_feature(enable = "avx2")]
@@ -76,6 +77,7 @@ struct __m32i(__m128i);
 
 impl Vector for __m32i {
     const LANES: usize = 4;
+    type Mask = Self;
 
     #[inline]
     #[target_feature(enable = "avx2")]
@@ -122,6 +124,7 @@ struct __m64i(__m128i);
 
 impl Vector for __m64i {
     const LANES: usize = 8;
+    type Mask = Self;
 
     #[inline]
     #[target_feature(enable = "avx2")]
@@ -163,6 +166,7 @@ impl From<__m128i> for __m64i {
 
 impl Vector for __m128i {
     const LANES: usize = 16;
+    type Mask = Self;
 
     #[inline]
     #[target_feature(enable = "avx2")]
@@ -197,6 +201,7 @@ impl Vector for __m128i {
 
 impl Vector for __m256i {
     const LANES: usize = 32;
+    type Mask = Self;
 
     #[inline]
     #[target_feature(enable = "avx2")]