Skip to main content

packed_seq/
packed_seq.rs

1use core::cell::RefCell;
2use std::ops::{Deref, DerefMut};
3use traits::Seq;
4use wide::u16x8;
5
6use crate::{intrinsics::transpose, padded_it::ChunkIt};
7
8use super::*;
9
10type SimdBuf = [S; 8];
11
12struct RecycledBox(Option<Box<SimdBuf>>);
13
14thread_local! {
15    static RECYCLED_BOX_CACHE: RefCell<Vec<Box<SimdBuf>>> = {
16        RefCell::new(vec![Box::new(SimdBuf::default())])
17    };
18}
19
20impl RecycledBox {
21    #[inline(always)]
22    fn take() -> Self {
23        let mut buf = RECYCLED_BOX_CACHE.with_borrow_mut(|v| RecycledBox(v.pop()));
24        buf.init_if_needed();
25        buf
26    }
27
28    #[inline(always)]
29    fn init_if_needed(&mut self) {
30        if self.0.is_none() {
31            self.0 = Some(Box::new(SimdBuf::default()));
32        }
33    }
34
35    #[inline(always)]
36    fn get(&self) -> &SimdBuf {
37        unsafe { self.0.as_ref().unwrap_unchecked() }
38    }
39
40    #[inline(always)]
41    fn get_mut(&mut self) -> &mut SimdBuf {
42        unsafe { self.0.as_mut().unwrap_unchecked() }
43    }
44}
45
46impl Deref for RecycledBox {
47    type Target = SimdBuf;
48
49    #[inline(always)]
50    fn deref(&self) -> &Self::Target {
51        self.get()
52    }
53}
54impl DerefMut for RecycledBox {
55    #[inline(always)]
56    fn deref_mut(&mut self) -> &mut SimdBuf {
57        self.get_mut()
58    }
59}
60
61impl Drop for RecycledBox {
62    #[inline(always)]
63    fn drop(&mut self) {
64        let mut x = None;
65        core::mem::swap(&mut x, &mut self.0);
66        RECYCLED_BOX_CACHE.with_borrow_mut(|v| v.push(unsafe { x.unwrap_unchecked() }));
67    }
68}
69
70struct RecycledVec(Vec<S>);
71
72thread_local! {
73    static RECYCLED_VEC_CACHE: RefCell<Vec<Vec<S>>> = {
74        RefCell::new(vec![])
75    };
76}
77
78impl RecycledVec {
79    #[inline(always)]
80    fn take() -> Self {
81        RecycledVec(RECYCLED_VEC_CACHE.with_borrow_mut(|v| v.pop().unwrap_or_default()))
82    }
83}
84
85impl Deref for RecycledVec {
86    type Target = Vec<S>;
87    #[inline(always)]
88    fn deref(&self) -> &Self::Target {
89        &self.0
90    }
91}
92impl DerefMut for RecycledVec {
93    #[inline(always)]
94    fn deref_mut(&mut self) -> &mut Self::Target {
95        &mut self.0
96    }
97}
98impl Drop for RecycledVec {
99    #[inline(always)]
100    fn drop(&mut self) {
101        RECYCLED_VEC_CACHE.with_borrow_mut(|v| v.push(std::mem::take(&mut self.0)));
102    }
103}
104
105#[doc(hidden)]
106pub struct Bits<const B: usize>;
107#[doc(hidden)]
108pub trait SupportedBits {}
109impl SupportedBits for Bits<1> {}
110impl SupportedBits for Bits<2> {}
111impl SupportedBits for Bits<4> {}
112impl SupportedBits for Bits<8> {}
113
114/// Number of padding bytes at the end of `PackedSeqVecBase::seq`.
115pub(crate) const PADDING: usize = 48;
116
117/// A variable-bit-width packed non-owned slice of DNA bases.
118#[derive(Copy, Clone, Debug, MemSize, MemDbg)]
119pub struct PackedSeqBase<'s, const B: usize>
120where
121    Bits<B>: SupportedBits,
122{
123    /// Packed data.
124    seq: &'s [u8],
125    /// Offset in bp from the start of the `seq`.
126    offset: usize,
127    /// Length of the sequence in bp, starting at `offset` from the start of `seq`.
128    len: usize,
129}
130
131/// A variable-bit-width packed owned sequence of DNA bases.
132#[derive(Clone, Debug, MemSize, MemDbg)]
133// #[cfg_attr(feature = "pyo3", pyo3::pyclass)] // Not supported: pyclass cannot have generic parameters
134#[cfg_attr(feature = "epserde", derive(epserde::Epserde))]
135pub struct PackedSeqVecBase<const B: usize> {
136    /// NOTE: We maintain the invariant that this has at least 48 bytes of padding
137    /// at the end after `len` finishes.
138    /// This ensures that `read_unaligned` in `as_64` works OK.
139    pub(crate) seq: Vec<u8>,
140
141    /// The length, in bp, of the underlying sequence. See `.len()`.
142    len: usize,
143}
144
145pub type PackedSeq<'s> = PackedSeqBase<'s, 2>;
146pub type PackedSeqVec = PackedSeqVecBase<2>;
147pub type BitSeq<'s> = PackedSeqBase<'s, 1>;
148pub type BitSeqVec = PackedSeqVecBase<1>;
149
150/// Convenience constants.
151/// B: bits per chat
152impl<'s, const B: usize> PackedSeqBase<'s, B>
153where
154    Bits<B>: SupportedBits,
155{
156    /// lowest B bits are 1.
157    const CHAR_MASK: u64 = (1 << B) - 1;
158    const SIMD_B: S = S::new([B as u32; 8]);
159    const SIMD_CHAR_MASK: S = S::new([(1 << B) - 1; 8]);
160    /// Chars per byte
161    const C8: usize = 8 / B;
162    /// Chars per u32
163    const C32: usize = 32 / B;
164    /// Chars per u256
165    const C256: usize = 256 / B;
166    /// Max length of a kmer that can be read as a single u64.
167    const K64: usize = (64 - 8) / B + 1;
168    /// Max length of a kmer that can be read as a single u128.
169    const K128: usize = (128 - 8) / B + 1;
170}
171
172/// Convenience constants.
173impl<const B: usize> PackedSeqVecBase<B>
174where
175    Bits<B>: SupportedBits,
176{
177    /// Chars per byte
178    const C8: usize = 8 / B;
179}
180
181impl<const B: usize> Default for PackedSeqVecBase<B>
182where
183    Bits<B>: SupportedBits,
184{
185    fn default() -> Self {
186        Self {
187            seq: vec![0; PADDING],
188            len: 0,
189        }
190    }
191}
192
193// ======================================================================
194// 2-BIT HELPER METHODS
195
196/// Pack an ASCII `ACTGactg` character into its 2-bit representation, and panic for anything else.
197#[inline(always)]
198pub fn pack_char(base: u8) -> u8 {
199    match base {
200        b'a' | b'A' => 0,
201        b'c' | b'C' => 1,
202        b'g' | b'G' => 3,
203        b't' | b'T' => 2,
204        _ => panic!(
205            "Unexpected character '{}' with ASCII value {base}. Expected one of ACTGactg.",
206            base as char
207        ),
208    }
209}
210
211/// Pack an ASCII `ACTGactg` character into its 2-bit representation, and silently convert other characters into 0..4 as well.
212#[inline(always)]
213pub fn pack_char_lossy(base: u8) -> u8 {
214    (base >> 1) & 3
215}
216/// Pack a slice of ASCII `ACTGactg` characters into its packed 2-bit kmer representation.
217/// Other characters are silently converted.
218#[inline(always)]
219pub fn pack_kmer_lossy(slice: &[u8]) -> u64 {
220    let mut kmer = 0;
221    for (i, &base) in slice.iter().enumerate() {
222        kmer |= (pack_char_lossy(base) as u64) << (2 * i);
223    }
224    kmer
225}
226/// Pack a slice of ASCII `ACTGactg` characters into its packed 2-bit kmer representation.
227#[inline(always)]
228pub fn pack_kmer_u128_lossy(slice: &[u8]) -> u128 {
229    let mut kmer = 0;
230    for (i, &base) in slice.iter().enumerate() {
231        kmer |= (pack_char_lossy(base) as u128) << (2 * i);
232    }
233    kmer
234}
235
236/// Unpack a 2-bit DNA base into the corresponding `ACTG` character.
237#[inline(always)]
238pub fn unpack_base(base: u8) -> u8 {
239    debug_assert!(base < 4, "Base {base} is not <4.");
240    b"ACTG"[base as usize]
241}
242
243/// Unpack a 2-bit encoded kmer into corresponding `ACTG` characters.
244/// Slice the returned array to the correct length.
245#[inline(always)]
246pub fn unpack_kmer(kmer: u64) -> [u8; 32] {
247    std::array::from_fn(|i| unpack_base(((kmer >> (2 * i)) & 3) as u8))
248}
249/// Unpack a 2-bit encoded kmer into corresponding `ACTG` character.
250#[inline(always)]
251pub fn unpack_kmer_into_vec(kmer: u64, k: usize, out: &mut Vec<u8>) {
252    out.clear();
253    out.extend((0..k).map(|i| unpack_base(((kmer >> (2 * i)) & 3) as u8)));
254}
255/// Unpack a 2-bit encoded kmer into corresponding `ACTG` character.
256#[inline(always)]
257pub fn unpack_kmer_to_vec(kmer: u64, k: usize) -> Vec<u8> {
258    let mut out = vec![];
259    unpack_kmer_into_vec(kmer, k, &mut out);
260    out
261}
262
263/// Unpack a 2-bit encoded kmer into corresponding `ACTG` characters.
264/// Slice the returned array to the correct length.
265#[inline(always)]
266pub fn unpack_kmer_u128(kmer: u128) -> [u8; 64] {
267    std::array::from_fn(|i| unpack_base(((kmer >> (2 * i)) & 3) as u8))
268}
269/// Unpack a 2-bit encoded kmer into corresponding `ACTG` character.
270#[inline(always)]
271pub fn unpack_kmer_u128_into_vec(kmer: u128, k: usize, out: &mut Vec<u8>) {
272    out.clear();
273    out.extend((0..k).map(|i| unpack_base(((kmer >> (2 * i)) & 3) as u8)));
274}
275/// Unpack a 2-bit encoded kmer into corresponding `ACTG` character.
276#[inline(always)]
277pub fn unpack_kmer_u128_to_vec(kmer: u128, k: usize) -> Vec<u8> {
278    let mut out = vec![];
279    unpack_kmer_u128_into_vec(kmer, k, &mut out);
280    out
281}
282
283/// Complement an ASCII character: `A<>T` and `C<>G`.
284#[inline(always)]
285pub const fn complement_char(base: u8) -> u8 {
286    match base {
287        b'A' => b'T',
288        b'C' => b'G',
289        b'G' => b'C',
290        b'T' => b'A',
291        _ => panic!("Unexpected character. Expected one of ACTGactg.",),
292    }
293}
294
295/// Complement a 2-bit base: `0<>2` and `1<>3`.
296#[inline(always)]
297pub const fn complement_base(base: u8) -> u8 {
298    base ^ 2
299}
300
301/// Complement 8 lanes of 2-bit bases: `0<>2` and `1<>3`.
302#[inline(always)]
303pub fn complement_base_simd(base: u32x8) -> u32x8 {
304    const TWO: u32x8 = u32x8::new([2; 8]);
305    base ^ TWO
306}
307
308/// Reverse complement the 2-bit pairs in the input.
309#[inline(always)]
310const fn revcomp_raw(word: u64) -> u64 {
311    #[cfg(any(target_arch = "arm", target_arch = "aarch64"))]
312    {
313        let mut res = word.reverse_bits(); // ARM can reverse bits in a single instruction
314        res = ((res >> 1) & 0x5555_5555_5555_5555) | ((res & 0x5555_5555_5555_5555) << 1);
315        res ^ 0xAAAA_AAAA_AAAA_AAAA
316    }
317
318    #[cfg(not(any(target_arch = "arm", target_arch = "aarch64")))]
319    {
320        let mut res = word.swap_bytes();
321        res = ((res >> 4) & 0x0F0F_0F0F_0F0F_0F0F) | ((res & 0x0F0F_0F0F_0F0F_0F0F) << 4);
322        res = ((res >> 2) & 0x3333_3333_3333_3333) | ((res & 0x3333_3333_3333_3333) << 2);
323        res ^ 0xAAAA_AAAA_AAAA_AAAA
324    }
325}
326
327/// Compute the reverse complement of a short sequence packed in a `u64`.
328#[inline(always)]
329pub const fn revcomp_u64(word: u64, len: usize) -> u64 {
330    revcomp_raw(word) >> (u64::BITS as usize - 2 * len)
331}
332
333#[inline(always)]
334pub const fn revcomp_u128(word: u128, len: usize) -> u128 {
335    let low = word as u64;
336    let high = (word >> 64) as u64;
337    let rlow = revcomp_raw(low);
338    let rhigh = revcomp_raw(high);
339    let out = ((rlow as u128) << 64) | rhigh as u128;
340    out >> (u128::BITS as usize - 2 * len)
341}
342
343// ======================================================================
344// 1-BIT HELPER METHODS
345
346/// 1 when a char is ambiguous.
347#[inline(always)]
348pub fn char_is_ambiguous(base: u8) -> u8 {
349    // (!matches!(base, b'A' | b'C'  | b'G'  | b'T' | b'a' | b'c'  | b'g'  | b't')) as u8
350    let table = b"ACTG";
351    let upper_mask = !(b'a' - b'A');
352    (table[pack_char_lossy(base) as usize] != (base & upper_mask)) as u8
353}
354
355/// Reverse `len` bits packed in a `u64`.
356#[inline(always)]
357pub const fn rev_u64(word: u64, len: usize) -> u64 {
358    word.reverse_bits() >> (u64::BITS as usize - len)
359}
360
361/// Reverse `len` bits packed in a `u128`.
362#[inline(always)]
363pub const fn rev_u128(word: u128, len: usize) -> u128 {
364    word.reverse_bits() >> (u128::BITS as usize - len)
365}
366
367// ======================================================================
368
369impl<'s, const B: usize> PackedSeqBase<'s, B>
370where
371    Bits<B>: SupportedBits,
372{
373    /// Creates a `Seq` from a slice of packed bytes, an offset in bp and a length in bp.
374    ///
375    /// The slice should have at least 48 bytes of padding after `offset + len`.
376    /// Otherwise, the function will panic.
377    pub fn from_raw_parts(seq: &'s [u8], offset: usize, len: usize) -> Self {
378        assert!(offset + len + PADDING * Self::C8 <= seq.len() * Self::C8);
379        Self { seq, offset, len }
380    }
381
382    /// Shrink `seq` to only just cover the data.
383    #[inline(always)]
384    pub fn normalize(&self) -> Self {
385        let start_byte = self.offset / Self::C8;
386        let end_byte = (self.offset + self.len).div_ceil(Self::C8);
387        Self {
388            seq: &self.seq[start_byte..end_byte + PADDING],
389            offset: self.offset % Self::C8,
390            len: self.len,
391        }
392    }
393
394    /// Return a `Vec<u8>` of ASCII `ACTG` characters.
395    #[inline(always)]
396    pub fn unpack(&self) -> Vec<u8> {
397        self.iter_bp().map(unpack_base).collect()
398    }
399}
400
401impl<'s> PackedSeqBase<'s, 2> {
402    /// Iterate over the reverse complement as packed 2-bit bases.
403    #[inline(always)]
404    pub fn iter_bp_rc(self) -> impl ExactSizeIterator<Item = u8> + 's {
405        (0..self.len)
406            .rev()
407            .map(move |i| complement_base(self.get(i)))
408    }
409
410    /// Append this packed slice as ASCII `ACTG` bases to `out`.
411    pub fn unpack_into(&self, out: &mut Vec<u8>) {
412        out.extend((0..self.len).map(|i| unpack_base(self.get(i))));
413    }
414
415    /// Append this the reverse complement packed slice as ASCII `ACTG` bases to `out`.
416    pub fn unpack_rc_into(&self, out: &mut Vec<u8>) {
417        out.extend(
418            (0..self.len)
419                .rev()
420                .map(|i| unpack_base(complement_base(self.get(i)))),
421        );
422    }
423}
424
425/// Read up to 32 bytes starting at idx.
426#[inline(always)]
427pub(crate) unsafe fn read_slice_32_unchecked(seq: &[u8], idx: usize) -> u32x8 {
428    unsafe {
429        let src = seq.as_ptr().add(idx);
430        debug_assert!(idx + 32 <= seq.len());
431        std::mem::transmute::<_, *const u32x8>(src).read_unaligned()
432    }
433}
434
435/// Read up to 32 bytes starting at idx.
436#[inline(always)]
437pub(crate) fn read_slice_32(seq: &[u8], idx: usize) -> u32x8 {
438    unsafe {
439        let src = seq.as_ptr().add(idx);
440        if idx + 32 <= seq.len() {
441            std::mem::transmute::<_, *const u32x8>(src).read_unaligned()
442        } else {
443            let num_bytes = seq.len().saturating_sub(idx);
444            let mut result = [0u8; 32];
445            std::ptr::copy_nonoverlapping(src, result.as_mut_ptr(), num_bytes);
446            std::mem::transmute(result)
447        }
448    }
449}
450
451/// Read up to 16 bytes starting at idx.
452#[allow(unused)]
453#[inline(always)]
454pub(crate) fn read_slice_16(seq: &[u8], idx: usize) -> u16x8 {
455    unsafe {
456        let src = seq.as_ptr().add(idx);
457        if idx + 16 <= seq.len() {
458            std::mem::transmute::<_, *const u16x8>(src).read_unaligned()
459        } else {
460            let num_bytes = seq.len().saturating_sub(idx);
461            let mut result = [0u8; 16];
462            std::ptr::copy_nonoverlapping(src, result.as_mut_ptr(), num_bytes);
463            std::mem::transmute(result)
464        }
465    }
466}
467
468impl<'s, const B: usize> Seq<'s> for PackedSeqBase<'s, B>
469where
470    Bits<B>: SupportedBits,
471{
472    const BITS_PER_CHAR: usize = B;
473    const BASES_PER_BYTE: usize = Self::C8;
474    type SeqVec = PackedSeqVecBase<B>;
475
476    #[inline(always)]
477    fn len(&self) -> usize {
478        self.len
479    }
480
481    #[inline(always)]
482    fn is_empty(&self) -> bool {
483        self.len == 0
484    }
485
486    #[inline(always)]
487    fn get_ascii(&self, index: usize) -> u8 {
488        unpack_base(self.get(index))
489    }
490
491    /// Convert a short sequence (kmer) to a packed representation as `u64`.
492    /// Panics if `self` is longer than 32 characters.
493    #[inline(always)]
494    fn as_u64(&self) -> u64 {
495        assert!(self.len() <= 64 / B);
496
497        let mask = u64::MAX >> (64 - B * self.len());
498
499        // The unaligned read is OK, because we ensure that the underlying `PackedSeqVecBase::seq` always
500        // has at least 48 bytes of padding at the end.
501        if self.len() <= Self::K64 {
502            let x = unsafe { (self.seq.as_ptr() as *const u64).read_unaligned() };
503            (x >> (B * self.offset)) & mask
504        } else {
505            let x = unsafe { (self.seq.as_ptr() as *const u128).read_unaligned() };
506            (x >> (B * self.offset)) as u64 & mask
507        }
508    }
509
510    /// Convert a short sequence (kmer) to a packed representation of its reverse complement as `usize`.
511    /// Panics if `self` is longer than 32 characters.
512    #[inline(always)]
513    fn revcomp_as_u64(&self) -> u64 {
514        match B {
515            1 => rev_u64(self.as_u64(), self.len()),
516            2 => revcomp_u64(self.as_u64(), self.len()),
517            _ => panic!("Rev(comp) is only supported for 1-bit and 2-bit alphabets."),
518        }
519    }
520
521    /// Convert a short sequence (kmer) to a packed representation as `u128`.
522    /// Panics if `self` is longer than 61 characters (for B=2) and the read spans more than 16 bytes.
523    #[inline(always)]
524    fn as_u128(&self) -> u128 {
525        let mask = u128::MAX >> (128 - B * self.len());
526
527        // Sequences >61 long cannot be read with a single unaligned u128 read.
528        if self.len() > Self::K128 {
529            assert!(self.len() <= 128 / B);
530            if self.offset + self.len() > 128 / B {
531                let x = unsafe { (self.seq.as_ptr() as *const u128).read_unaligned() };
532                let y = unsafe { self.seq.as_ptr().offset(16).read() as u128 };
533                return ((x >> (B * self.offset)) | (y << (128 - B * self.offset))) & mask;
534            }
535        }
536
537        // The unaligned read is OK, because we ensure that the underlying `PackedSeqVecBase::seq` always
538        // has at least 48 bytes of padding at the end.
539        let x = unsafe { (self.seq.as_ptr() as *const u128).read_unaligned() };
540        (x >> (B * self.offset)) & mask
541    }
542
543    /// Convert a short sequence (kmer) to a packed representation of its reverse complement as `usize`.
544    /// Panics if `self` is longer than 64 characters.
545    #[inline(always)]
546    fn revcomp_as_u128(&self) -> u128 {
547        match B {
548            1 => rev_u128(self.as_u128(), self.len()),
549            2 => revcomp_u128(self.as_u128(), self.len()),
550            _ => panic!("Rev(comp) is only supported for 1-bit and 2-bit alphabets."),
551        }
552    }
553
554    #[inline(always)]
555    fn to_vec(&self) -> PackedSeqVecBase<B> {
556        assert_eq!(self.offset, 0);
557        PackedSeqVecBase {
558            seq: self
559                .seq
560                .iter()
561                .copied()
562                .chain(std::iter::repeat_n(0u8, PADDING))
563                .collect(),
564            len: self.len,
565        }
566    }
567
568    fn to_revcomp(&self) -> PackedSeqVecBase<B> {
569        match B {
570            1 | 2 => {}
571            _ => panic!("Can only reverse (&complement) 1-bit and 2-bit packed sequences.",),
572        }
573
574        let mut seq = self.seq[..(self.offset + self.len).div_ceil(Self::C8)]
575            .iter()
576            // 1. reverse the bytes
577            .rev()
578            .copied()
579            .map(|mut res| {
580                match B {
581                    2 => {
582                        // 2. swap the bases in the byte
583                        // This is auto-vectorized.
584                        res = ((res >> 4) & 0x0F) | ((res & 0x0F) << 4);
585                        res = ((res >> 2) & 0x33) | ((res & 0x33) << 2);
586                        // Complement the bases.
587                        res ^ 0xAA
588                    }
589                    1 => res.reverse_bits(),
590                    _ => unreachable!(),
591                }
592            })
593            .chain(std::iter::repeat_n(0u8, PADDING))
594            .collect::<Vec<u8>>();
595
596        // 3. Shift away the offset.
597        let new_offset = (Self::C8 - (self.offset + self.len) % Self::C8) % Self::C8;
598
599        if new_offset > 0 {
600            // Shift everything left by `2*new_offset` bits.
601            let shift = B * new_offset;
602            *seq.last_mut().unwrap() >>= shift;
603            // This loop is also auto-vectorized.
604            for i in 0..seq.len() - 1 {
605                seq[i] = (seq[i] >> shift) | (seq[i + 1] << (8 - shift));
606            }
607        }
608
609        PackedSeqVecBase { seq, len: self.len }
610    }
611
612    #[inline(always)]
613    fn slice(&self, range: Range<usize>) -> Self {
614        debug_assert!(
615            range.end <= self.len,
616            "Slice index out of bounds: {} > {}",
617            range.end,
618            self.len
619        );
620        PackedSeqBase {
621            seq: self.seq,
622            offset: self.offset + range.start,
623            len: range.end - range.start,
624        }
625        .normalize()
626    }
627
628    #[inline(always)]
629    fn iter_bp(self) -> impl ExactSizeIterator<Item = u8> {
630        assert!(self.len <= self.seq.len() * Self::C8);
631
632        let this = self.normalize();
633
634        // read u64 at a time?
635        let mut byte = 0;
636        (0..this.len + this.offset)
637            .map(
638                #[inline(always)]
639                move |i| {
640                    if i % Self::C8 == 0 {
641                        byte = this.seq[i / Self::C8];
642                    }
643                    // Shift byte instead of i?
644                    (byte >> (B * (i % Self::C8))) & Self::CHAR_MASK as u8
645                },
646            )
647            .advance(this.offset)
648    }
649
650    #[inline(always)]
651    fn par_iter_bp(self, context: usize) -> PaddedIt<impl ChunkIt<S>> {
652        self.par_iter_bp_with_buf(context, RecycledBox::take())
653    }
654
655    #[inline(always)]
656    fn par_iter_bp_delayed(self, context: usize, delay: Delay) -> PaddedIt<impl ChunkIt<(S, S)>> {
657        self.par_iter_bp_delayed_with_factor(context, delay, 1)
658    }
659
660    /// NOTE: When `self` starts does not start at a byte boundary, the
661    /// 'delayed' character is not guaranteed to be `0`.
662    #[inline(always)]
663    fn par_iter_bp_delayed_2(
664        self,
665        context: usize,
666        delay1: Delay,
667        delay2: Delay,
668    ) -> PaddedIt<impl ChunkIt<(S, S, S)>> {
669        self.par_iter_bp_delayed_2_with_factor_and_buf(
670            context,
671            delay1,
672            delay2,
673            1,
674            RecycledVec::take(),
675        )
676    }
677
678    /// Compares 29 characters at a time.
679    fn cmp_lcp(&self, other: &Self) -> (std::cmp::Ordering, usize) {
680        let mut lcp = 0;
681        let min_len = self.len.min(other.len);
682        for i in (0..min_len).step_by(Self::K64) {
683            let len = (min_len - i).min(Self::K64);
684            let this = self.slice(i..i + len);
685            let other = other.slice(i..i + len);
686            let this_word = this.as_u64();
687            let other_word = other.as_u64();
688            if this_word != other_word {
689                // Unfortunately, bases are packed in little endian order, so the default order is reversed.
690                let eq = this_word ^ other_word;
691                let t = eq.trailing_zeros() as usize / B * B;
692                lcp += t / B;
693                let mask = (Self::CHAR_MASK) << t;
694                return ((this_word & mask).cmp(&(other_word & mask)), lcp);
695            }
696            lcp += len;
697        }
698        (self.len.cmp(&other.len), lcp)
699    }
700
701    #[inline(always)]
702    fn get(&self, index: usize) -> u8 {
703        let offset = self.offset + index;
704        let idx = offset / Self::C8;
705        let offset = offset % Self::C8;
706        (self.seq[idx] >> (B * offset)) & Self::CHAR_MASK as u8
707    }
708}
709
710impl<'s, const B: usize> PackedSeqBase<'s, B>
711where
712    Bits<B>: SupportedBits,
713{
714    #[inline(always)]
715    pub fn par_iter_bp_with_buf<BUF: DerefMut<Target = [S; 8]>>(
716        self,
717        context: usize,
718        mut buf: BUF,
719    ) -> PaddedIt<impl ChunkIt<S> + use<'s, B, BUF>> {
720        #[cfg(target_endian = "big")]
721        panic!("Big endian architectures are not supported.");
722
723        let this = self.normalize();
724        let o = this.offset;
725        assert!(o < Self::C8);
726
727        let num_kmers = if this.len == 0 {
728            0
729        } else {
730            (this.len + o).saturating_sub(context - 1)
731        };
732        // without +o, since we don't need them in the stride.
733        let num_kmers_stride = this.len.saturating_sub(context - 1);
734        let n = num_kmers_stride.div_ceil(L).next_multiple_of(Self::C8);
735        let bytes_per_chunk = n / Self::C8;
736        let padding = Self::C8 * L * bytes_per_chunk - num_kmers_stride;
737
738        let offsets: [usize; 8] = from_fn(|l| l * bytes_per_chunk);
739        let mut cur = S::ZERO;
740
741        let par_len = if num_kmers == 0 {
742            0
743        } else {
744            n + context + o - 1
745        };
746
747        let last_i = par_len.saturating_sub(1) / Self::C32 * Self::C32;
748        // Safety check for the `read_slice_32_unchecked`:
749        assert!(offsets[7] + (last_i / Self::C8) + 32 <= this.seq.len());
750
751        let it = (0..par_len)
752            .map(
753                #[inline(always)]
754                move |i| {
755                    if i % Self::C32 == 0 {
756                        if i % Self::C256 == 0 {
757                            // Read a u256 for each lane containing the next 128 characters.
758                            let data: [u32x8; 8] = from_fn(
759                                #[inline(always)]
760                                |lane| unsafe {
761                                    read_slice_32_unchecked(
762                                        this.seq,
763                                        offsets[lane] + (i / Self::C8),
764                                    )
765                                },
766                            );
767                            *buf = transpose(data);
768                        }
769                        cur = buf[(i % Self::C256) / Self::C32];
770                    }
771                    // Extract the last 2 bits of each character.
772                    let chars = cur & Self::SIMD_CHAR_MASK;
773                    // Shift remaining characters to the right.
774                    cur = cur >> Self::SIMD_B;
775                    chars
776                },
777            )
778            .advance(o);
779
780        PaddedIt { it, padding }
781    }
782
783    #[inline(always)]
784    pub fn par_iter_bp_delayed_with_factor(
785        self,
786        context: usize,
787        delay: Delay,
788        factor: usize,
789    ) -> PaddedIt<impl ChunkIt<(S, S)> + use<'s, B>> {
790        self.par_iter_bp_delayed_with_factor_and_buf(context, delay, factor, RecycledVec::take())
791    }
792
793    #[inline(always)]
794    pub fn par_iter_bp_delayed_with_buf<BUF: DerefMut<Target = Vec<S>>>(
795        self,
796        context: usize,
797        delay: Delay,
798        buf: BUF,
799    ) -> PaddedIt<impl ChunkIt<(S, S)> + use<'s, B, BUF>> {
800        self.par_iter_bp_delayed_with_factor_and_buf(context, delay, 1, buf)
801    }
802
803    #[inline(always)]
804    pub fn par_iter_bp_delayed_with_factor_and_buf<BUF: DerefMut<Target = Vec<S>>>(
805        self,
806        context: usize,
807        Delay(delay): Delay,
808        factor: usize,
809        mut buf: BUF,
810    ) -> PaddedIt<impl ChunkIt<(S, S)> + use<'s, B, BUF>> {
811        #[cfg(target_endian = "big")]
812        panic!("Big endian architectures are not supported.");
813
814        assert!(
815            delay < usize::MAX / 2,
816            "Delay={} should be >=0.",
817            delay as isize
818        );
819
820        let this = self.normalize();
821        let o = this.offset;
822        assert!(o < Self::C8);
823
824        let num_kmers = if this.len == 0 {
825            0
826        } else {
827            (this.len + o).saturating_sub(context - 1)
828        };
829        // without +o, since we don't need them in the stride.
830        let num_kmers_stride = this.len.saturating_sub(context - 1);
831        let n = num_kmers_stride
832            .div_ceil(L)
833            .next_multiple_of(factor * Self::C8);
834        let bytes_per_chunk = n / Self::C8;
835        let padding = Self::C8 * L * bytes_per_chunk - num_kmers_stride;
836
837        let offsets: [usize; 8] = from_fn(|l| l * bytes_per_chunk);
838        let mut upcoming = S::ZERO;
839        let mut upcoming_d = S::ZERO;
840
841        // Even buf_len is nice to only have the write==buf_len check once.
842        // We also make it the next power of 2, for faster modulo operations.
843        // delay/16: number of bp in a u32.
844        // +8: some 'random' padding
845        let buf_len = (delay / Self::C32 + 8).next_power_of_two();
846        let buf_mask = buf_len - 1;
847        if buf.capacity() < buf_len {
848            // This has better codegen than `vec.clear(); vec.resize()`, since the inner `do_reserve_and_handle` of resize is not inlined.
849            *buf = vec![S::ZERO; buf_len];
850        } else {
851            // NOTE: Buf needs to be filled with zeros to guarantee returning 0 values for out-of-bounds characters.
852            buf.clear();
853            buf.resize(buf_len, S::ZERO);
854        }
855
856        let mut write_idx = 0;
857        // We compensate for the first delay/16 triggers of the check below that
858        // happen before the delay is actually reached.
859        let mut read_idx = (buf_len - delay / Self::C32) % buf_len;
860
861        let par_len = if num_kmers == 0 {
862            0
863        } else {
864            n + context + o - 1
865        };
866
867        let last_i = par_len.saturating_sub(1) / Self::C32 * Self::C32;
868        // Safety check for the `read_slice_32_unchecked`:
869        assert!(offsets[7] + (last_i / Self::C8) + 32 <= this.seq.len());
870
871        let it = (0..par_len)
872            .map(
873                #[inline(always)]
874                move |i| {
875                    if i % Self::C32 == 0 {
876                        if i % Self::C256 == 0 {
877                            // Read a u256 for each lane containing the next 128 characters.
878                            let data: [u32x8; 8] = from_fn(
879                                #[inline(always)]
880                                |lane| unsafe {
881                                    read_slice_32_unchecked(
882                                        this.seq,
883                                        offsets[lane] + (i / Self::C8),
884                                    )
885                                },
886                            );
887                            unsafe {
888                                *TryInto::<&mut [u32x8; 8]>::try_into(
889                                    buf.get_unchecked_mut(write_idx..write_idx + 8),
890                                )
891                                .unwrap_unchecked() = transpose(data);
892                            }
893                            if i == 0 {
894                                // Mask out chars before the offset.
895                                let elem = !((1u32 << (B * o)) - 1);
896                                let mask = S::splat(elem);
897                                unsafe { assert_unchecked(write_idx < buf.len()) };
898                                buf[write_idx] &= mask;
899                            }
900                        }
901                        unsafe { assert_unchecked(write_idx < buf.len()) };
902                        upcoming = buf[write_idx];
903                        write_idx += 1;
904                        write_idx &= buf_mask;
905                    }
906                    if i % Self::C32 == delay % Self::C32 {
907                        unsafe { assert_unchecked(read_idx < buf.len()) };
908                        upcoming_d = buf[read_idx];
909                        read_idx += 1;
910                        read_idx &= buf_mask;
911                    }
912                    // Extract the last 2 bits of each character.
913                    let chars = upcoming & Self::SIMD_CHAR_MASK;
914                    let chars_d = upcoming_d & Self::SIMD_CHAR_MASK;
915                    // Shift remaining characters to the right.
916                    upcoming = upcoming >> Self::SIMD_B;
917                    upcoming_d = upcoming_d >> Self::SIMD_B;
918                    (chars, chars_d)
919                },
920            )
921            .advance(o);
922
923        PaddedIt { it, padding }
924    }
925
926    #[inline(always)]
927    pub fn par_iter_bp_delayed_2_with_factor(
928        self,
929        context: usize,
930        delay1: Delay,
931        delay2: Delay,
932        factor: usize,
933    ) -> PaddedIt<impl ChunkIt<(S, S, S)> + use<'s, B>> {
934        self.par_iter_bp_delayed_2_with_factor_and_buf(
935            context,
936            delay1,
937            delay2,
938            factor,
939            RecycledVec::take(),
940        )
941    }
942
943    #[inline(always)]
944    pub fn par_iter_bp_delayed_2_with_buf<BUF: DerefMut<Target = Vec<S>>>(
945        self,
946        context: usize,
947        delay1: Delay,
948        delay2: Delay,
949        buf: BUF,
950    ) -> PaddedIt<impl ChunkIt<(S, S, S)> + use<'s, B, BUF>> {
951        self.par_iter_bp_delayed_2_with_factor_and_buf(context, delay1, delay2, 1, buf)
952    }
953
954    /// When iterating over 2-bit and 1-bit encoded data in parallel,
955    /// one must ensure that they have the same stride.
956    /// On the larger type, set `factor` as the ratio to the smaller one,
957    /// so that the stride in bytes is a multiple of `factor`,
958    /// so that the smaller type also has a byte-aligned stride.
959    #[inline(always)]
960    pub fn par_iter_bp_delayed_2_with_factor_and_buf<BUF: DerefMut<Target = Vec<S>>>(
961        self,
962        context: usize,
963        Delay(delay1): Delay,
964        Delay(delay2): Delay,
965        factor: usize,
966        mut buf: BUF,
967    ) -> PaddedIt<impl ChunkIt<(S, S, S)> + use<'s, B, BUF>> {
968        #[cfg(target_endian = "big")]
969        panic!("Big endian architectures are not supported.");
970
971        let this = self.normalize();
972        let o = this.offset;
973        assert!(o < Self::C8);
974        assert!(delay1 <= delay2, "Delay1 must be at most delay2.");
975
976        let num_kmers = if this.len == 0 {
977            0
978        } else {
979            (this.len + o).saturating_sub(context - 1)
980        };
981        // without +o, since we don't need them in the stride.
982        let num_kmers_stride = this.len.saturating_sub(context - 1);
983        let n = num_kmers_stride
984            .div_ceil(L)
985            .next_multiple_of(factor * Self::C8);
986        let bytes_per_chunk = n / Self::C8;
987        let padding = Self::C8 * L * bytes_per_chunk - num_kmers_stride;
988
989        let offsets: [usize; 8] = from_fn(|l| l * bytes_per_chunk);
990        let mut upcoming = S::ZERO;
991        let mut upcoming_d1 = S::ZERO;
992        let mut upcoming_d2 = S::ZERO;
993
994        // Even buf_len is nice to only have the write==buf_len check once.
995        let buf_len = (delay2 / Self::C32 + 8).next_power_of_two();
996        let buf_mask = buf_len - 1;
997        if buf.capacity() < buf_len {
998            // This has better codegen than `vec.clear(); vec.resize()`, since the inner `do_reserve_and_handle` of resize is not inlined.
999            *buf = vec![S::ZERO; buf_len];
1000        } else {
1001            // NOTE: Buf needs to be filled with zeros to guarantee returning 0 values for out-of-bounds characters.
1002            buf.clear();
1003            buf.resize(buf_len, S::ZERO);
1004        }
1005
1006        let mut write_idx = 0;
1007        // We compensate for the first delay/16 triggers of the check below that
1008        // happen before the delay is actually reached.
1009        let mut read_idx1 = (buf_len - delay1 / Self::C32) % buf_len;
1010        let mut read_idx2 = (buf_len - delay2 / Self::C32) % buf_len;
1011
1012        let par_len = if num_kmers == 0 {
1013            0
1014        } else {
1015            n + context + o - 1
1016        };
1017
1018        let last_i = par_len.saturating_sub(1) / Self::C32 * Self::C32;
1019        // Safety check for the `read_slice_32_unchecked`:
1020        assert!(offsets[7] + (last_i / Self::C8) + 32 <= this.seq.len());
1021
1022        let it = (0..par_len)
1023            .map(
1024                #[inline(always)]
1025                move |i| {
1026                    if i % Self::C32 == 0 {
1027                        if i % Self::C256 == 0 {
1028                            // Read a u256 for each lane containing the next 128 characters.
1029                            let data: [u32x8; 8] = from_fn(
1030                                #[inline(always)]
1031                                |lane| unsafe {
1032                                    read_slice_32_unchecked(
1033                                        this.seq,
1034                                        offsets[lane] + (i / Self::C8),
1035                                    )
1036                                },
1037                            );
1038                            unsafe {
1039                                *TryInto::<&mut [u32x8; 8]>::try_into(
1040                                    buf.get_unchecked_mut(write_idx..write_idx + 8),
1041                                )
1042                                .unwrap_unchecked() = transpose(data);
1043                            }
1044                            // FIXME DROP THIS?
1045                            if i == 0 {
1046                                // Mask out chars before the offset.
1047                                let elem = !((1u32 << (B * o)) - 1);
1048                                let mask = S::splat(elem);
1049                                buf[write_idx] &= mask;
1050                            }
1051                        }
1052                        upcoming = buf[write_idx];
1053                        write_idx += 1;
1054                        write_idx &= buf_mask;
1055                    }
1056                    if i % Self::C32 == delay1 % Self::C32 {
1057                        unsafe { assert_unchecked(read_idx1 < buf.len()) };
1058                        upcoming_d1 = buf[read_idx1];
1059                        read_idx1 += 1;
1060                        read_idx1 &= buf_mask;
1061                    }
1062                    if i % Self::C32 == delay2 % Self::C32 {
1063                        unsafe { assert_unchecked(read_idx2 < buf.len()) };
1064                        upcoming_d2 = buf[read_idx2];
1065                        read_idx2 += 1;
1066                        read_idx2 &= buf_mask;
1067                    }
1068                    // Extract the last 2 bits of each character.
1069                    let chars = upcoming & Self::SIMD_CHAR_MASK;
1070                    let chars_d1 = upcoming_d1 & Self::SIMD_CHAR_MASK;
1071                    let chars_d2 = upcoming_d2 & Self::SIMD_CHAR_MASK;
1072                    // Shift remaining characters to the right.
1073                    upcoming = upcoming >> Self::SIMD_B;
1074                    upcoming_d1 = upcoming_d1 >> Self::SIMD_B;
1075                    upcoming_d2 = upcoming_d2 >> Self::SIMD_B;
1076                    (chars, chars_d1, chars_d2)
1077                },
1078            )
1079            .advance(o);
1080
1081        PaddedIt { it, padding }
1082    }
1083}
1084
1085impl<const B: usize> PartialEq for PackedSeqBase<'_, B>
1086where
1087    Bits<B>: SupportedBits,
1088{
1089    /// Compares 29 characters at a time.
1090    fn eq(&self, other: &Self) -> bool {
1091        if self.len != other.len {
1092            return false;
1093        }
1094        for i in (0..self.len).step_by(Self::K64) {
1095            let len = (self.len - i).min(Self::K64);
1096            let this = self.slice(i..i + len);
1097            let that = other.slice(i..i + len);
1098            if this.as_u64() != that.as_u64() {
1099                return false;
1100            }
1101        }
1102        true
1103    }
1104}
1105
1106impl<const B: usize> Eq for PackedSeqBase<'_, B> where Bits<B>: SupportedBits {}
1107
1108impl<const B: usize> PartialOrd for PackedSeqBase<'_, B>
1109where
1110    Bits<B>: SupportedBits,
1111{
1112    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
1113        Some(self.cmp(other))
1114    }
1115}
1116
1117impl<const B: usize> Ord for PackedSeqBase<'_, B>
1118where
1119    Bits<B>: SupportedBits,
1120{
1121    /// Compares 29 characters at a time.
1122    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
1123        let min_len = self.len.min(other.len);
1124        for i in (0..min_len).step_by(Self::K64) {
1125            let len = (min_len - i).min(Self::K64);
1126            let this = self.slice(i..i + len);
1127            let other = other.slice(i..i + len);
1128            let this_word = this.as_u64();
1129            let other_word = other.as_u64();
1130            if this_word != other_word {
1131                // Unfortunately, bases are packed in little endian order, so the default order is reversed.
1132                let eq = this_word ^ other_word;
1133                let t = eq.trailing_zeros() as usize / B * B;
1134                let mask = (Self::CHAR_MASK) << t;
1135                return (this_word & mask).cmp(&(other_word & mask));
1136            }
1137        }
1138        self.len.cmp(&other.len)
1139    }
1140}
1141
1142impl<const B: usize> SeqVec for PackedSeqVecBase<B>
1143where
1144    Bits<B>: SupportedBits,
1145{
1146    type Seq<'s> = PackedSeqBase<'s, B>;
1147
1148    #[inline(always)]
1149    fn into_raw(mut self) -> Vec<u8> {
1150        self.seq.resize(self.len.div_ceil(Self::C8), 0);
1151        self.seq
1152    }
1153
1154    #[inline(always)]
1155    fn as_slice(&self) -> Self::Seq<'_> {
1156        PackedSeqBase {
1157            seq: &self.seq[..self.len.div_ceil(Self::C8) + PADDING],
1158            offset: 0,
1159            len: self.len,
1160        }
1161    }
1162
1163    #[inline(always)]
1164    fn len(&self) -> usize {
1165        self.len
1166    }
1167
1168    #[inline(always)]
1169    fn is_empty(&self) -> bool {
1170        self.len == 0
1171    }
1172
1173    #[inline(always)]
1174    fn clear(&mut self) {
1175        self.seq.clear();
1176        self.len = 0;
1177    }
1178
1179    fn push_seq<'a>(&mut self, seq: PackedSeqBase<'_, B>) -> Range<usize> {
1180        let start = self.len.next_multiple_of(Self::C8) + seq.offset;
1181        let end = start + seq.len();
1182
1183        // Shrink away the padding.
1184        self.seq.resize(self.len.div_ceil(Self::C8), 0);
1185        // Extend.
1186        self.seq.extend(seq.seq);
1187
1188        // For the ambiguous N-seq, mark the entire gap as ambiguous.
1189        if B == 1 {
1190            if self.len % 8 != 0 {
1191                self.seq[self.len / 8] |= u8::MAX << (self.len % 8);
1192            }
1193            self.seq[self.len.div_ceil(8)] |= (1u8 << seq.offset) - 1;
1194        }
1195        self.len = end;
1196        start..end
1197    }
1198
1199    /// For `PackedSeqVec` (2-bit encoding): map ASCII `ACGT` (and `acgt`) to DNA `0132` in that order.
1200    /// Other characters are silently mapped into `0..4`.
1201    ///
1202    /// Returns the range of the newly added sequence in the packed representation.
1203    ///
1204    /// Uses the BMI2 `pext` instruction when available, based on the
1205    /// `n_to_bits_pext` method described at
1206    /// <https://github.com/Daniel-Liu-c0deb0t/cute-nucleotides>.
1207    ///
1208    /// For `BitSeqVec` (1-bit encoding): map `ACGTacgt` to `0`, and everything else to `1`.
1209    fn push_ascii(&mut self, seq: &[u8]) -> Range<usize> {
1210        match B {
1211            1 | 2 => {}
1212            _ => panic!(
1213                "Can only use ASCII input for 2-bit DNA packing, or 1-bit ambiguous indicators."
1214            ),
1215        }
1216
1217        if seq.is_empty() {
1218            return self.len..self.len;
1219        }
1220
1221        self.seq
1222            .resize((self.len + seq.len()).div_ceil(Self::C8) + PADDING, 0);
1223
1224        let start_aligned = self.len.next_multiple_of(Self::C8);
1225        let start = self.len;
1226        let len = seq.len();
1227        let mut idx = self.len / Self::C8;
1228
1229        let parse_base = |base| match B {
1230            1 => char_is_ambiguous(base),
1231            2 => pack_char_lossy(base),
1232            _ => unreachable!(),
1233        };
1234
1235        let unaligned = core::cmp::min(start_aligned - start, len);
1236        if unaligned > 0 {
1237            let mut packed_byte = self.seq[idx];
1238            for &base in &seq[..unaligned] {
1239                packed_byte |= parse_base(base) << ((self.len % Self::C8) * B);
1240                self.len += 1;
1241            }
1242            self.seq[idx] = packed_byte;
1243            idx += 1;
1244        }
1245
1246        #[allow(unused)]
1247        let mut last = unaligned;
1248
1249        if B == 2 {
1250            #[cfg(all(target_arch = "x86_64", target_feature = "bmi2"))]
1251            {
1252                last = unaligned + (len - unaligned) / 8 * 8;
1253
1254                for i in (unaligned..last).step_by(8) {
1255                    let chunk =
1256                        unsafe { seq.get_unchecked(i..i + 8).try_into().unwrap_unchecked() };
1257                    let ascii = u64::from_le_bytes(chunk);
1258                    let packed_bytes =
1259                        unsafe { std::arch::x86_64::_pext_u64(ascii, 0x0606060606060606) } as u16;
1260                    unsafe {
1261                        self.seq
1262                            .get_unchecked_mut(idx..(idx + 2))
1263                            .copy_from_slice(&packed_bytes.to_le_bytes())
1264                    };
1265                    idx += 2;
1266                    self.len += 8;
1267                }
1268            }
1269
1270            #[cfg(target_feature = "neon")]
1271            {
1272                use core::arch::aarch64::{
1273                    vandq_u8, vdup_n_u8, vld1q_u8, vpadd_u8, vshlq_u8, vst1_u8,
1274                };
1275                use core::mem::transmute;
1276
1277                last = unaligned + (len - unaligned) / 16 * 16;
1278
1279                for i in (unaligned..last).step_by(16) {
1280                    unsafe {
1281                        let ascii = vld1q_u8(seq.as_ptr().add(i));
1282                        let masked_bits = vandq_u8(ascii, transmute([6i8; 16]));
1283                        let (bits_0, bits_1) = transmute(vshlq_u8(
1284                            masked_bits,
1285                            transmute([-1i8, 1, 3, 5, -1, 1, 3, 5, -1, 1, 3, 5, -1, 1, 3, 5]),
1286                        ));
1287                        let half_packed = vpadd_u8(bits_0, bits_1);
1288                        let packed = vpadd_u8(half_packed, vdup_n_u8(0));
1289                        vst1_u8(self.seq.as_mut_ptr().add(idx), packed);
1290                        idx += Self::C8;
1291                        self.len += 16;
1292                    }
1293                }
1294            }
1295
1296            #[cfg(all(target_arch = "wasm32", target_feature = "simd128"))]
1297            {
1298                use core::mem::transmute as t;
1299                use wide::{u8x16, u32x4};
1300
1301                last = unaligned + (len - unaligned) / 16 * 16;
1302
1303                let pick_low_bytes = u8x16::new([
1304                    0, 4, 8, 12, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255,
1305                ]);
1306
1307                for i in (unaligned..last).step_by(16) {
1308                    unsafe {
1309                        let ascii = seq.as_ptr().add(i).cast::<u8x16>().read_unaligned();
1310                        let bits: u32x4 = t((ascii >> 1) & u8x16::splat(3));
1311                        let bits = bits | (bits >> 6);
1312                        let bits = bits | (bits >> 12);
1313                        let packed: [u32; 4] =
1314                            t(t::<_, u8x16>(bits).shuffle_zeroing(pick_low_bytes));
1315                        self.seq
1316                            .get_unchecked_mut(idx..(idx + 4))
1317                            .copy_from_slice(&packed[0].to_le_bytes());
1318                        idx += 4;
1319                        self.len += 16;
1320                    }
1321                }
1322            }
1323        }
1324
1325        if B == 1 {
1326            #[cfg(all(target_arch = "x86_64", target_feature = "avx2"))]
1327            {
1328                last = len;
1329                self.len += len - unaligned;
1330
1331                let mut last_i = unaligned;
1332
1333                for i in (unaligned..last).step_by(32) {
1334                    use std::mem::transmute as t;
1335
1336                    // Wide doesn't have u8x32, so this is messy here...
1337                    type S = wide::i8x32;
1338                    let chars: S = unsafe { t(read_slice_32(seq, i)) };
1339                    let upper_mask = !(b'a' - b'A');
1340                    // make everything upper case
1341                    let chars = chars & S::splat(upper_mask as i8);
1342                    let lossy_encoded = chars & S::splat(6);
1343                    let table = unsafe { S::from(t::<_, S>(*b"AxCxTxGxxxxxxxxxAxCxTxGxxxxxxxxx")) };
1344                    let lookup: S = unsafe {
1345                        t(std::arch::x86_64::_mm256_shuffle_epi8(
1346                            t(table),
1347                            t(lossy_encoded),
1348                        ))
1349                    };
1350                    let packed_bytes = !(chars.simd_eq(lookup).to_bitmask() as u32);
1351
1352                    last_i = i;
1353                    unsafe {
1354                        self.seq
1355                            .get_unchecked_mut(idx..(idx + 4))
1356                            .copy_from_slice(&packed_bytes.to_le_bytes())
1357                    };
1358                    idx += 4;
1359                }
1360
1361                // Fix up trailing bytes.
1362                if unaligned < last {
1363                    idx -= 4;
1364                    let mut val = unsafe {
1365                        u32::from_le_bytes(
1366                            self.seq
1367                                .get_unchecked(idx..(idx + 4))
1368                                .try_into()
1369                                .unwrap_unchecked(),
1370                        )
1371                    };
1372                    // keep only the `last - last_i` low bits.
1373                    let keep = last - last_i;
1374                    val <<= 32 - keep;
1375                    val >>= 32 - keep;
1376                    unsafe {
1377                        self.seq
1378                            .get_unchecked_mut(idx..(idx + 4))
1379                            .copy_from_slice(&val.to_le_bytes())
1380                    };
1381                    idx += keep.div_ceil(8);
1382                }
1383            }
1384
1385            #[cfg(target_feature = "neon")]
1386            {
1387                use core::arch::aarch64::*;
1388                use core::mem::transmute;
1389
1390                last = unaligned + (len - unaligned) / 64 * 64;
1391
1392                for i in (unaligned..last).step_by(64) {
1393                    unsafe {
1394                        let ptr = seq.as_ptr().add(i);
1395                        let chars = vld4q_u8(ptr);
1396
1397                        let upper_mask = vdupq_n_u8(!(b'a' - b'A'));
1398                        let chars = neon::map_8x16x4(chars, |v| vandq_u8(v, upper_mask));
1399
1400                        let two_bits_mask = vdupq_n_u8(6);
1401                        let lossy_encoded = neon::map_8x16x4(chars, |v| vandq_u8(v, two_bits_mask));
1402
1403                        let table = transmute(*b"AxCxTxGxxxxxxxxx");
1404                        let lookup = neon::map_8x16x4(lossy_encoded, |v| vqtbl1q_u8(table, v));
1405
1406                        let mask = neon::map_two_8x16x4(chars, lookup, |v1, v2| vceqq_u8(v1, v2));
1407                        let packed_bytes = !neon::movemask_64(mask);
1408
1409                        self.seq[idx..(idx + 8)].copy_from_slice(&packed_bytes.to_le_bytes());
1410                        idx += 8;
1411                        self.len += 64;
1412                    }
1413                }
1414            }
1415
1416            #[cfg(all(target_arch = "wasm32", target_feature = "simd128"))]
1417            {
1418                use wide::u8x16;
1419
1420                last = unaligned + (len - unaligned) / 16 * 16;
1421
1422                let upper_mask = u8x16::splat(!(b'a' - b'A'));
1423                let two_bits_mask = u8x16::splat(6);
1424                let table = u8x16::new(*b"AxCxTxGxxxxxxxxx");
1425
1426                for i in (unaligned..last).step_by(16) {
1427                    unsafe {
1428                        let chars = seq.as_ptr().add(i).cast::<u8x16>().read_unaligned();
1429                        let chars = chars & upper_mask;
1430                        let lossy_encoded = chars & two_bits_mask;
1431                        let lookup = table.shuffle_zeroing(lossy_encoded);
1432                        let packed_bytes = !(chars.simd_eq(lookup).to_bitmask() as u16);
1433
1434                        self.seq
1435                            .get_unchecked_mut(idx..(idx + 2))
1436                            .copy_from_slice(&packed_bytes.to_le_bytes());
1437                        idx += 2;
1438                        self.len += 16;
1439                    }
1440                }
1441            }
1442        }
1443
1444        let mut packed_byte = 0;
1445        for &base in &seq[last..] {
1446            packed_byte |= parse_base(base) << ((self.len % Self::C8) * B);
1447            self.len += 1;
1448            if self.len % Self::C8 == 0 {
1449                self.seq[idx] = packed_byte;
1450                idx += 1;
1451                packed_byte = 0;
1452            }
1453        }
1454        if self.len % Self::C8 != 0 && last < len {
1455            self.seq[idx] = packed_byte;
1456            idx += 1;
1457        }
1458        assert_eq!(idx + PADDING, self.seq.len());
1459        start..start + len
1460    }
1461
1462    #[cfg(feature = "rand")]
1463    fn random(n: usize) -> Self {
1464        use rand::Rng;
1465
1466        let byte_len = n.div_ceil(Self::C8);
1467        let mut seq = vec![0; byte_len + PADDING];
1468        rand::make_rng::<rand::rngs::SmallRng>().fill_bytes(&mut seq[..byte_len]);
1469        // Ensure that the last byte is padded with zeros.
1470        if n % Self::C8 != 0 {
1471            seq[byte_len - 1] &= (1 << (B * (n % Self::C8))) - 1;
1472        }
1473
1474        Self { seq, len: n }
1475    }
1476}
1477
1478impl PackedSeqVecBase<2> {
1479    /// Creates a `SeqVec` from a vector of packed bytes and a length in bp.
1480    ///
1481    /// The vector should have at least 48 bytes of padding after `len`.
1482    /// Otherwise, the vector will be resized to be padded with zeros.
1483    pub fn from_raw_parts(mut seq: Vec<u8>, len: usize) -> Self {
1484        assert!(len <= seq.len() * Self::C8);
1485        seq.resize(len.div_ceil(Self::C8) + PADDING, 0);
1486        Self { seq, len }
1487    }
1488
1489    /// Read a fasta/fastq file at the given path into a `PackedSeq`.
1490    /// Returns the sequence as well as the slices of all records.
1491    ///
1492    /// Currently, all slices are bit-packed, without any padding.
1493    #[cfg(feature = "file_io")]
1494    pub fn from_fastx(path: &std::path::Path) -> (Self, Vec<Range<usize>>) {
1495        use helicase::{HelicaseParser, input::FromFile};
1496
1497        let mut ranges = vec![];
1498
1499        const CONFIG: helicase::Config = helicase::ParserOptions::default()
1500            .ignore_headers()
1501            .ignore_quality()
1502            .dna_packed()
1503            .reject_non_actg()
1504            .config()
1505            | helicase::config::advanced::MERGE_RECORDS;
1506        let mut parser = helicase::FastxParser::<CONFIG>::from_file(&path).unwrap();
1507
1508        let mut start = 0;
1509        while let Some(event) = parser.next() {
1510            event.unwrap();
1511            let end = parser.get_dna_packed().len();
1512            ranges.push(start..end);
1513            start = end;
1514        }
1515        let (len, seq) = parser.get_dna_packed_owned().into_packed_bytes();
1516
1517        let seq = PackedSeqVec::from_raw_parts(seq, len);
1518        (seq, ranges)
1519    }
1520
1521    /// Read a plain FASTA/FASTQ file using `num_threads` workers.
1522    /// NOTE: This may return data out-of-order!
1523    ///
1524    /// Returns the sequence and record ranges in chunk completion order.
1525    /// There may be padding between chunks; use the returned ranges to access records.
1526    /// A thread count of zero disables multi-threading.
1527    /// Compression formats without parallel decoding support are read sequentially.
1528    #[cfg(feature = "file_io")]
1529    pub fn from_fastx_parallel(
1530        path: &std::path::Path,
1531        num_threads: usize,
1532    ) -> (Self, Vec<Range<usize>>) {
1533        use helicase::HelicaseParser;
1534
1535        const CONFIG: helicase::Config = helicase::ParserOptions::default()
1536            .ignore_headers()
1537            .ignore_quality()
1538            .dna_packed()
1539            .reject_non_actg()
1540            .config()
1541            | helicase::config::advanced::MERGE_RECORDS;
1542        let chunks = std::sync::Mutex::new(Vec::new());
1543        // Spawn threads that parse the file in chunks and save the packed data.
1544        helicase::ParallelFastxParser::<CONFIG>::from_file(path)
1545            .unwrap()
1546            .threads(num_threads)
1547            .for_each(|mut parser| {
1548                let mut ranges = Vec::new();
1549                let mut start = 0;
1550                while let Some(event) = parser.next() {
1551                    event?;
1552                    let end = parser.get_dna_packed().len();
1553                    ranges.push(start..end);
1554                    start = end;
1555                }
1556                let (len, mut seq) = parser.get_dna_packed_owned().into_packed_bytes();
1557                seq.truncate(len.div_ceil(Self::C8));
1558                chunks.lock().unwrap().push((seq, ranges));
1559                Ok(())
1560            })
1561            .unwrap();
1562
1563        // Accumulate per-thread lengths.
1564        let chunks = chunks.into_inner().unwrap();
1565        let total_bytes = chunks.iter().map(|(seq, _)| seq.len()).sum::<usize>();
1566        let total_records = chunks.iter().map(|(_, ranges)| ranges.len()).sum::<usize>();
1567
1568        // Allocate space for the final sequence and the final ranges.
1569        let mut seq = vec![0; total_bytes + PADDING];
1570        let mut ranges = Vec::<Range<usize>>::with_capacity(total_records);
1571
1572        // Spawn threads. Each thread takes a chunk of chunks and writes them to the global storage.
1573        std::thread::scope(|scope| {
1574            let mut seq_out = &mut seq[..total_bytes];
1575            let mut ranges_out = &mut ranges.spare_capacity_mut()[..total_records];
1576            let mut offset = 0;
1577            let threads = num_threads.max(1);
1578            let groups = chunks.chunks(chunks.len().div_ceil(threads));
1579            for group in groups {
1580                let thread_bytes = group.iter().map(|(seq, _)| seq.len()).sum();
1581                let thread_records = group.iter().map(|(_, ranges)| ranges.len()).sum();
1582
1583                // claim space for thread
1584                let (mut dest, rest) = seq_out.split_at_mut(thread_bytes);
1585                seq_out = rest;
1586                let (mut range_dest, rest) = ranges_out.split_at_mut(thread_records);
1587                ranges_out = rest;
1588
1589                let mut base_offset = offset * Self::C8;
1590                offset += thread_bytes;
1591
1592                scope.spawn(move || {
1593                    for (chunk, chunk_ranges) in group {
1594                        // claim space for chunk
1595                        let (chunk_dest, rest) = dest.split_at_mut(chunk.len());
1596                        dest = rest;
1597                        let (chunk_range_dest, rest) = range_dest.split_at_mut(chunk_ranges.len());
1598                        range_dest = rest;
1599
1600                        // copy chunk
1601                        chunk_dest.copy_from_slice(&chunk);
1602
1603                        // copy ranges
1604                        let chunk_offset = base_offset;
1605                        base_offset += chunk.len() * Self::C8;
1606                        for (dest, range) in chunk_range_dest.iter_mut().zip(chunk_ranges) {
1607                            dest.write(chunk_offset + range.start..chunk_offset + range.end);
1608                        }
1609                    }
1610                });
1611            }
1612        });
1613        // SAFETY: The disjoint slices cover all `range_len` slots, and each slot
1614        // was initialized above. The scope joins every worker before returning.
1615        unsafe { ranges.set_len(total_records) };
1616        (Self::from_raw_parts(seq, total_bytes * Self::C8), ranges)
1617    }
1618}
1619
1620impl PackedSeqVecBase<1> {
1621    pub fn with_len(n: usize) -> Self {
1622        Self {
1623            seq: vec![0; n.div_ceil(Self::C8) + PADDING],
1624            len: n,
1625        }
1626    }
1627
1628    #[cfg(feature = "rand")]
1629    pub fn random(len: usize, n_frac: f32) -> Self {
1630        let byte_len = len.div_ceil(Self::C8);
1631        let mut seq = vec![0; byte_len + PADDING];
1632
1633        assert!(
1634            (0.0..=0.3).contains(&n_frac),
1635            "n_frac={} should be in [0, 0.3]",
1636            n_frac
1637        );
1638
1639        for _ in 0..(len as f32 * n_frac) as usize {
1640            let idx = rand::random::<u64>() as usize % len;
1641            let byte = idx / Self::C8;
1642            let offset = idx % Self::C8;
1643            seq[byte] |= 1 << offset;
1644        }
1645
1646        Self { seq, len }
1647    }
1648}
1649
1650impl<'s> PackedSeqBase<'s, 1> {
1651    /// An iterator indicating for each kmer whether it contains ambiguous bases.
1652    ///
1653    /// Returns n-(k-1) elements.
1654    #[inline(always)]
1655    pub fn iter_kmer_ambiguity(self, k: usize) -> impl ExactSizeIterator<Item = bool> + use<'s> {
1656        let this = self.normalize();
1657        assert!(k > 0);
1658        assert!(k <= Self::K64);
1659        (this.offset..this.offset + this.len.saturating_sub(k - 1))
1660            .map(move |i| self.read_kmer(k, i) != 0)
1661    }
1662
1663    /// A parallel iterator indicating for each kmer whether it contains ambiguous bases.
1664    ///
1665    /// First element is the 'kmer' consisting only of the first character of each chunk.
1666    ///
1667    /// `k`: length of windows to check
1668    /// `context`: number of overlapping iterations +1. To determine stride of each lane.
1669    /// `skip`: Set to `context-1` to skip the iterations added by the context.
1670    #[inline(always)]
1671    pub fn par_iter_kmer_ambiguity(
1672        self,
1673        k: usize,
1674        context: usize,
1675        skip: usize,
1676    ) -> PaddedIt<impl ChunkIt<S> + use<'s>> {
1677        #[cfg(target_endian = "big")]
1678        panic!("Big endian architectures are not supported.");
1679
1680        assert!(k > 0, "par_iter_kmers requires k>0, but k={k}");
1681        assert!(k <= 96, "par_iter_kmers requires k<=96, but k={k}");
1682
1683        let this = self.normalize();
1684        let o = this.offset;
1685        assert!(o < Self::C8);
1686
1687        let delay = k - 1;
1688
1689        let it = self.par_iter_bp_delayed(context, Delay(delay));
1690
1691        let mut cnt = u32x8::ZERO;
1692
1693        it.map(
1694            #[inline(always)]
1695            move |(a, r)| {
1696                cnt += a;
1697                let out = cnt.simd_gt(S::ZERO);
1698                cnt -= r;
1699                out
1700            },
1701        )
1702        .advance(skip)
1703    }
1704
1705    #[inline(always)]
1706    pub fn par_iter_kmer_ambiguity_with_buf(
1707        self,
1708        k: usize,
1709        context: usize,
1710        skip: usize,
1711        buf: &'s mut Vec<S>,
1712    ) -> PaddedIt<impl ChunkIt<S> + use<'s>> {
1713        #[cfg(target_endian = "big")]
1714        panic!("Big endian architectures are not supported.");
1715
1716        assert!(k > 0, "par_iter_kmers requires k>0, but k={k}");
1717        assert!(k <= 96, "par_iter_kmers requires k<=96, but k={k}");
1718
1719        let this = self.normalize();
1720        let o = this.offset;
1721        assert!(o < Self::C8);
1722
1723        let delay = k - 1;
1724
1725        let it = self.par_iter_bp_delayed_with_buf(context, Delay(delay), buf);
1726
1727        let mut cnt = u32x8::ZERO;
1728
1729        it.map(
1730            #[inline(always)]
1731            move |(a, r)| {
1732                cnt += a;
1733                let out = cnt.simd_gt(S::ZERO);
1734                cnt -= r;
1735                out
1736            },
1737        )
1738        .advance(skip)
1739    }
1740}
1741
1742#[cfg(target_feature = "neon")]
1743mod neon {
1744    use core::arch::aarch64::*;
1745
1746    #[inline(always)]
1747    pub fn movemask_64(v: uint8x16x4_t) -> u64 {
1748        // https://stackoverflow.com/questions/74722950/convert-vector-compare-mask-into-bit-mask-in-aarch64-simd-or-arm-neon/74748402#74748402
1749        unsafe {
1750            let acc = vsriq_n_u8(vsriq_n_u8(v.3, v.2, 1), vsriq_n_u8(v.1, v.0, 1), 2);
1751            vget_lane_u64(
1752                vreinterpret_u64_u8(vshrn_n_u16(
1753                    vreinterpretq_u16_u8(vsriq_n_u8(acc, acc, 4)),
1754                    4,
1755                )),
1756                0,
1757            )
1758        }
1759    }
1760
1761    #[inline(always)]
1762    pub fn map_8x16x4<F>(v: uint8x16x4_t, mut f: F) -> uint8x16x4_t
1763    where
1764        F: FnMut(uint8x16_t) -> uint8x16_t,
1765    {
1766        uint8x16x4_t(f(v.0), f(v.1), f(v.2), f(v.3))
1767    }
1768
1769    #[inline(always)]
1770    pub fn map_two_8x16x4<F>(v1: uint8x16x4_t, v2: uint8x16x4_t, mut f: F) -> uint8x16x4_t
1771    where
1772        F: FnMut(uint8x16_t, uint8x16_t) -> uint8x16_t,
1773    {
1774        uint8x16x4_t(f(v1.0, v2.0), f(v1.1, v2.1), f(v1.2, v2.2), f(v1.3, v2.3))
1775    }
1776}