Add NEON version for ambiguous-base detection

src/packed_seq.rs

@@ -1,4 +1,5 @@
 use traits::Seq;
+use wide::u16x8;
 
 use crate::{intrinsics::transpose, padded_it::ChunkIt};
 
@@ -195,39 +196,16 @@ pub fn char_is_ambiguous(base: u8) -> u8 {
     (table[pack_char_lossy(base) as usize] != (base & upper_mask)) as u8
 }
 
-/// Reverse the bits in the input.
-#[inline(always)]
-const fn rev_raw(word: u64) -> u64 {
-    #[cfg(any(target_arch = "arm", target_arch = "aarch64"))]
-    {
-        // ARM can reverse bits in a single instruction
-        word.reverse_bits()
-    }
-
-    #[cfg(not(any(target_arch = "arm", target_arch = "aarch64")))]
-    {
-        let mut res = word.swap_bytes();
-        res = ((res >> 4) & 0x0F0F_0F0F_0F0F_0F0F) | ((res & 0x0F0F_0F0F_0F0F_0F0F) << 4);
-        res = ((res >> 2) & 0x3333_3333_3333_3333) | ((res & 0x3333_3333_3333_3333) << 2);
-        res = ((res >> 1) & 0x5555_5555_5555_5555) | ((res & 0x5555_5555_5555_5555) << 1);
-        res ^ 0xAAAA_AAAA_AAAA_AAAA
-    }
-}
-
-/// Compute the reverse complement of a short sequence packed in a `u64`.
+/// Reverse `len` bits packed in a `u64`.
 #[inline(always)]
 pub const fn rev_u64(word: u64, len: usize) -> u64 {
-    rev_raw(word) >> (usize::BITS as usize - len)
+    word.reverse_bits() >> (usize::BITS as usize - len)
 }
 
+/// Reverse `len` bits packed in a `u128`.
 #[inline(always)]
 pub const fn rev_u128(word: u128, len: usize) -> u128 {
-    let low = word as u64;
-    let high = (word >> 64) as u64;
-    let rlow = rev_raw(low);
-    let rhigh = rev_raw(high);
-    let out = ((rlow as u128) << 64) | rhigh as u128;
-    out >> (u128::BITS as usize - len)
+    word.reverse_bits() >> (u128::BITS as usize - len)
 }
 
 // ======================================================================
@@ -268,6 +246,20 @@ pub(crate) fn read_slice(seq: &[u8], idx: usize) -> u32x8 {
     }
 }
 
+/// Read up to 16 bytes starting at idx.
+#[allow(unused)]
+#[inline(always)]
+pub(crate) fn read_slice_16(seq: &[u8], idx: usize) -> u16x8 {
+    // assert!(idx <= seq.len());
+    let mut result = [0u8; 16];
+    let num_bytes = 16.min(seq.len().saturating_sub(idx));
+    unsafe {
+        let src = seq.as_ptr().add(idx);
+        std::ptr::copy_nonoverlapping(src, result.as_mut_ptr(), num_bytes);
+        std::mem::transmute(result)
+    }
+}
+
 impl<'s, const B: usize> Seq<'s> for PackedSeqBase<'s, B>
 where
     Bits<B>: SupportedBits,
@@ -904,7 +896,6 @@ where
         #[allow(unused)]
         let mut last = unaligned;
 
-        // TODO: Vectorization for B=1?
         if B == 2 {
             #[cfg(all(target_arch = "x86_64", target_feature = "bmi2"))]
             {
@@ -950,7 +941,6 @@ where
             }
         }
         if B == 1 {
-            // FIXME: Add NEON version.
             #[cfg(all(target_arch = "x86_64", target_feature = "avx2"))]
             {
                 last = unaligned + len;
@@ -992,6 +982,41 @@ where
                     }
                 }
             }
+
+            #[cfg(target_feature = "neon")]
+            {
+                last = unaligned + len;
+                self.len = len;
+
+                for i in (unaligned..last).step_by(16) {
+                    use std::mem::transmute as t;
+
+                    use wide::CmpEq;
+                    type S = wide::i8x16;
+                    let chars: S = unsafe { t(read_slice_16(seq, i)) };
+                    let upper_mask = !(b'a' - b'A');
+                    // make everything upper case
+                    let chars = chars & S::splat(upper_mask as i8);
+                    let lossy_encoded = chars & S::splat(6);
+                    let table = unsafe { S::from(t::<_, S>(*b"AxCxTxGxxxxxxxxx")) };
+                    let lookup: S =
+                        unsafe { t(std::arch::aarch64::vqtbl1q_u8(t(table), t(lossy_encoded))) };
+                    let packed_bytes = !(chars.cmp_eq(lookup).move_mask() as u16);
+
+                    if i + 16 <= last {
+                        self.seq[idx + 0] = packed_bytes as u8;
+                        self.seq[idx + 1] = (packed_bytes >> 8) as u8;
+                        idx += 2;
+                    } else {
+                        let mut b = 0;
+                        while i + b < last {
+                            self.seq[idx] = (packed_bytes >> b) as u8;
+                            idx += 1;
+                            b += 8;
+                        }
+                    }
+                }
+            }
         }
 
         let mut packed_byte = 0;