Provide SIMD implementations of the by-unit is_foo checks in mem.

src/ascii.rs

@@ -24,6 +24,10 @@
 #[cfg(all(feature = "simd-accel", any(target_feature = "sse2", all(target_endian = "little", target_arch = "aarch64"))))]
 use simd_funcs::*;
 
+// `as` truncates, so works on 32-bit, too.
+pub const ASCII_MASK: usize = 0x80808080_80808080u64 as usize;
+pub const BASIC_LATIN_MASK: usize = 0xFF80FF80_FF80FF80u64 as usize;
+
 #[allow(unused_macros)]
 macro_rules! ascii_naive {
     ($name:ident,
@@ -492,7 +496,7 @@ macro_rules! ascii_to_ascii_simd_stride {
     #[inline(always)]
     pub unsafe fn $name(src: *const u8, dst: *mut u8) -> bool {
         let simd = $load(src);
-        if !is_ascii(simd) {
+        if !simd_is_ascii(simd) {
             return false;
         }
         $store(dst, simd);
@@ -508,7 +512,7 @@ macro_rules! ascii_to_basic_latin_simd_stride {
     #[inline(always)]
     pub unsafe fn $name(src: *const u8, dst: *mut u16) -> bool {
         let simd = $load(src);
-        if !is_ascii(simd) {
+        if !simd_is_ascii(simd) {
             return false;
         }
         let (first, second) = simd_unpack(simd);
@@ -541,7 +545,7 @@ macro_rules! basic_latin_to_ascii_simd_stride {
     pub unsafe fn $name(src: *const u16, dst: *mut u8) -> bool {
         let first = $load(src);
         let second = $load(src.offset(8));
-        if is_basic_latin(first | second) {
+        if simd_is_basic_latin(first | second) {
             $store(dst, simd_pack(first, second));
             true
         } else {
@@ -843,7 +847,7 @@ cfg_if! {
                 let len_minus_stride = len - STRIDE_SIZE;
                 loop {
                     let simd = unsafe { load16_unaligned(src.offset(offset as isize)) };
-                    if !is_ascii(simd) {
+                    if !simd_is_ascii(simd) {
                         break;
                     }
                     offset += STRIDE_SIZE;
@@ -914,10 +918,6 @@ cfg_if! {
             None
         }
     } else {
-        // `as` truncates, so works on 32-bit, too.
-        pub const ASCII_MASK: usize = 0x80808080_80808080u64 as usize;
-        pub const BASIC_LATIN_MASK: usize = 0xFF80FF80_FF80FF80u64 as usize;
-
         #[inline(always)]
         unsafe fn unpack_latin1_stride_alu(src: *const usize, dst: *mut usize) {
             let word = *src;

src/mem.rs

@@ -22,10 +22,21 @@ use super::EncoderResult;
 use utf_8::Utf8Decoder;
 use utf_8::Utf8Encoder;
 
+#[cfg(all(feature = "simd-accel", any(target_feature = "sse2", all(target_endian = "little", target_arch = "aarch64"))))]
+use simd_funcs::*;
+
+#[cfg(all(feature = "simd-accel", any(target_feature = "sse2", all(target_endian = "little", target_arch = "aarch64"))))]
+use simd::u8x16;
+
+#[cfg(all(feature = "simd-accel", any(target_feature = "sse2", all(target_endian = "little", target_arch = "aarch64"))))]
+use simd::u16x8;
+
 const SIMD_ALIGNMENT: usize = 16;
 
 const SIMD_ALIGNMENT_MASK: usize = 15;
 
+const SIMD_STRIDE_SIZE: usize = 16;
+
 const ALU_ALIGNMENT: usize = 8;
 
 const ALU_ALIGNMENT_MASK: usize = 7;
@@ -35,6 +46,7 @@ const ALU_STRIDE_SIZE: usize = 8;
 // `as` truncates, so works on 32-bit, too.
 const LATIN1_MASK: usize = 0xFF00FF00_FF00FF00u64 as usize;
 
+#[allow(unused_macros)]
 macro_rules! by_unit_check_alu {
     ($name:ident,
      $unit:ty,
@@ -71,9 +83,63 @@ macro_rules! by_unit_check_alu {
     })
 }
 
-by_unit_check_alu!(is_ascii_alu, u8, ASCII_MASK);
-by_unit_check_alu!(is_basic_latin_alu, u16, BASIC_LATIN_MASK);
-by_unit_check_alu!(is_utf16_latin1_alu, u16, LATIN1_MASK);
+#[allow(unused_macros)]
+macro_rules! by_unit_check_simd {
+    ($name:ident,
+     $unit:ty,
+     $splat:expr,
+     $simd_ty:ty,
+     $mask:ident,
+     $func:ident) => (
+    #[inline(always)]
+    fn $name(buffer: &[$unit]) -> bool {
+        let unit_size = ::std::mem::size_of::<$unit>();
+        let src = buffer.as_ptr();
+        let len = buffer.len();
+        let mut offset = 0usize;
+        let mut until_alignment = ((SIMD_ALIGNMENT - ((src as usize) & SIMD_ALIGNMENT_MASK)) &
+                                   SIMD_ALIGNMENT_MASK) / unit_size;
+        let mut accu = 0usize;
+        if until_alignment + SIMD_STRIDE_SIZE / unit_size <= len {
+            while until_alignment != 0 {
+                accu |= buffer[offset] as usize;
+                offset += 1;
+                until_alignment -= 1;
+            }
+            let len_minus_stride = len - SIMD_STRIDE_SIZE / unit_size;
+            let mut simd_accu = $splat;
+            loop {
+                simd_accu = simd_accu | unsafe { *(src.offset(offset as isize) as *const $simd_ty) };
+                offset += SIMD_STRIDE_SIZE / unit_size;
+                if offset > len_minus_stride {
+                    break;
+                }
+            }
+            // TODO: Could optimize away this branch on aarch64 and arm by
+            // ORing the horizontal max into accu.
+            if !$func(simd_accu) {
+                return false;
+            }
+        }
+        while offset < len {
+            accu |= buffer[offset] as usize;
+            offset += 1;
+        }
+        accu & $mask == 0
+    })
+}
+
+cfg_if!{
+    if #[cfg(all(feature = "simd-accel", any(target_feature = "sse2", all(target_endian = "little", target_arch = "aarch64"))))] {
+        by_unit_check_simd!(is_ascii_impl, u8, u8x16::splat(0), u8x16, ASCII_MASK, simd_is_ascii);
+        by_unit_check_simd!(is_basic_latin_impl, u16, u16x8::splat(0), u16x8, BASIC_LATIN_MASK, simd_is_basic_latin);
+        by_unit_check_simd!(is_utf16_latin1_impl, u16, u16x8::splat(0), u16x8, LATIN1_MASK, simd_is_latin1);
+    } else {
+        by_unit_check_alu!(is_ascii_impl, u8, ASCII_MASK);
+        by_unit_check_alu!(is_basic_latin_impl, u16, BASIC_LATIN_MASK);
+        by_unit_check_alu!(is_utf16_latin1_impl, u16, LATIN1_MASK);
+    }
+}
 
 #[inline(always)]
 fn utf16_valid_up_to_alu(buffer: &[u16]) -> usize {
@@ -115,7 +181,7 @@ fn utf16_valid_up_to_alu(buffer: &[u16]) -> usize {
 /// May read the entire buffer even if it isn't all-ASCII. (I.e. the function
 /// is not guaranteed to fail fast.)
 pub fn is_ascii(buffer: &[u8]) -> bool {
-    is_ascii_alu(buffer)
+    is_ascii_impl(buffer)
 }
 
 /// Checks whether the buffer is all-Basic Latin (i.e. UTF-16 representing
@@ -124,7 +190,7 @@ pub fn is_ascii(buffer: &[u8]) -> bool {
 /// May read the entire buffer even if it isn't all-ASCII. (I.e. the function
 /// is not guaranteed to fail fast.)
 pub fn is_basic_latin(buffer: &[u16]) -> bool {
-    is_basic_latin_alu(buffer)
+    is_basic_latin_impl(buffer)
 }
 
 /// Checks whether the buffer is valid UTF-8 representing only code points
@@ -179,7 +245,7 @@ pub fn is_str_latin1(buffer: &str) -> bool {
 /// May read the entire buffer even if it isn't all-Latin1. (I.e. the function
 /// is not guaranteed to fail fast.)
 pub fn is_utf16_latin1(buffer: &[u16]) -> bool {
-    is_utf16_latin1_alu(buffer)
+    is_utf16_latin1_impl(buffer)
 }
 
 /// Converts potentially-invalid UTF-8 to valid UTF-16 with errors replaced

src/simd_funcs.rs

@@ -89,7 +89,7 @@ cfg_if! {
 cfg_if! {
     if #[cfg(target_feature = "sse2")] {
         #[inline(always)]
-        pub fn is_ascii(s: u8x16) -> bool {
+        pub fn simd_is_ascii(s: u8x16) -> bool {
             unsafe {
                 let signed: i8x16 = ::std::mem::transmute_copy(&s);
                 x86_mm_movemask_epi8(signed) == 0
@@ -101,14 +101,14 @@ cfg_if! {
         }
 
         #[inline(always)]
-        pub fn is_ascii(s: u8x16) -> bool {
+        pub fn simd_is_ascii(s: u8x16) -> bool {
             unsafe {
                 aarch64_vmaxvq_u8(s) < 0x80
             }
         }
     } else {
         #[inline(always)]
-        pub fn is_ascii(s: u8x16) -> bool {
+        pub fn simd_is_ascii(s: u8x16) -> bool {
             let highest_ascii = u8x16::splat(0x7F);
             !s.gt(highest_ascii).any()
         }
@@ -122,17 +122,30 @@ cfg_if! {
         }
 
         #[inline(always)]
-        pub fn is_basic_latin(s: u16x8) -> bool {
+        pub fn simd_is_basic_latin(s: u16x8) -> bool {
             unsafe {
                 aarch64_vmaxvq_u16(s) < 0x80
             }
         }
+
+        #[inline(always)]
+        pub fn simd_is_latin1(s: u16x8) -> bool {
+            unsafe {
+                aarch64_vmaxvq_u16(s) < 0x100
+            }
+        }
     } else {
         #[inline(always)]
-        pub fn is_basic_latin(s: u16x8) -> bool {
+        pub fn simd_is_basic_latin(s: u16x8) -> bool {
             let highest_ascii = u16x8::splat(0x7F);
             !s.gt(highest_ascii).any()
         }
+
+        #[inline(always)]
+        pub fn simd_is_latin1(s: u16x8) -> bool {
+            let highest_latin1 = u16x8::splat(0xFF);
+            !s.gt(highest_latin1).any()
+        }
     }
 }
 
@@ -206,7 +219,7 @@ mod tests {
     }
 
     #[test]
-    fn test_is_basic_latin_success() {
+    fn test_simd_is_basic_latin_success() {
         let ascii: [u8; 16] = [0x61, 0x62, 0x63, 0x64, 0x65, 0x66, 0x67, 0x68, 0x69, 0x70, 0x71,
                                0x72, 0x73, 0x74, 0x75, 0x76];
         let basic_latin: [u16; 16] = [0x61, 0x62, 0x63, 0x64, 0x65, 0x66, 0x67, 0x68, 0x69, 0x70,
@@ -216,7 +229,7 @@ mod tests {
         let mut vec = Vec::with_capacity(16);
         vec.resize(16, 0u8);
         let ptr = vec.as_mut_ptr();
-        assert!(is_basic_latin(first | second));
+        assert!(simd_is_basic_latin(first | second));
         unsafe {
             store16_unaligned(ptr, simd_pack(first, second));
         }
@@ -224,46 +237,46 @@ mod tests {
     }
 
     #[test]
-    fn test_is_basic_latin_c0() {
+    fn test_simd_is_basic_latin_c0() {
         let input: [u16; 16] = [0x61, 0x62, 0x63, 0x81, 0x65, 0x66, 0x67, 0x68, 0x69, 0x70, 0x71,
                                 0x72, 0x73, 0x74, 0x75, 0x76];
         let first = unsafe { load8_unaligned(input.as_ptr()) };
         let second = unsafe { load8_unaligned(input.as_ptr().offset(8)) };
-        assert!(!is_basic_latin(first | second));
+        assert!(!simd_is_basic_latin(first | second));
     }
 
     #[test]
-    fn test_is_basic_latin_0fff() {
+    fn test_simd_is_basic_latin_0fff() {
         let input: [u16; 16] = [0x61, 0x62, 0x63, 0x0FFF, 0x65, 0x66, 0x67, 0x68, 0x69, 0x70,
                                 0x71, 0x72, 0x73, 0x74, 0x75, 0x76];
         let first = unsafe { load8_unaligned(input.as_ptr()) };
         let second = unsafe { load8_unaligned(input.as_ptr().offset(8)) };
-        assert!(!is_basic_latin(first | second));
+        assert!(!simd_is_basic_latin(first | second));
     }
 
     #[test]
-    fn test_is_basic_latin_ffff() {
+    fn test_simd_is_basic_latin_ffff() {
         let input: [u16; 16] = [0x61, 0x62, 0x63, 0xFFFF, 0x65, 0x66, 0x67, 0x68, 0x69, 0x70,
                                 0x71, 0x72, 0x73, 0x74, 0x75, 0x76];
         let first = unsafe { load8_unaligned(input.as_ptr()) };
         let second = unsafe { load8_unaligned(input.as_ptr().offset(8)) };
-        assert!(!is_basic_latin(first | second));
+        assert!(!simd_is_basic_latin(first | second));
     }
 
     #[test]
-    fn test_is_ascii_success() {
+    fn test_simd_is_ascii_success() {
         let ascii: [u8; 16] = [0x61, 0x62, 0x63, 0x64, 0x65, 0x66, 0x67, 0x68, 0x69, 0x70, 0x71,
                                0x72, 0x73, 0x74, 0x75, 0x76];
         let simd = unsafe { load16_unaligned(ascii.as_ptr()) };
-        assert!(is_ascii(simd));
+        assert!(simd_is_ascii(simd));
     }
 
     #[test]
-    fn test_is_ascii_failure() {
+    fn test_simd_is_ascii_failure() {
         let input: [u8; 16] = [0x61, 0x62, 0x63, 0x64, 0x81, 0x66, 0x67, 0x68, 0x69, 0x70, 0x71,
                                0x72, 0x73, 0x74, 0x75, 0x76];
         let simd = unsafe { load16_unaligned(input.as_ptr()) };
-        assert!(!is_ascii(simd));
+        assert!(!simd_is_ascii(simd));
     }
 
     #[cfg(target_feature = "sse2")]