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)]
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}
169
170/// Convenience constants.
171impl<const B: usize> PackedSeqVecBase<B>
172where
173    Bits<B>: SupportedBits,
174{
175    /// Chars per byte
176    const C8: usize = 8 / B;
177}
178
179impl<const B: usize> Default for PackedSeqVecBase<B>
180where
181    Bits<B>: SupportedBits,
182{
183    fn default() -> Self {
184        Self {
185            seq: vec![0; PADDING],
186            len: 0,
187        }
188    }
189}
190
191// ======================================================================
192// 2-BIT HELPER METHODS
193
194/// Pack an ASCII `ACTGactg` character into its 2-bit representation, and panic for anything else.
195#[inline(always)]
196pub fn pack_char(base: u8) -> u8 {
197    match base {
198        b'a' | b'A' => 0,
199        b'c' | b'C' => 1,
200        b'g' | b'G' => 3,
201        b't' | b'T' => 2,
202        _ => panic!(
203            "Unexpected character '{}' with ASCII value {base}. Expected one of ACTGactg.",
204            base as char
205        ),
206    }
207}
208
209/// Pack an ASCII `ACTGactg` character into its 2-bit representation, and silently convert other characters into 0..4 as well.
210#[inline(always)]
211pub fn pack_char_lossy(base: u8) -> u8 {
212    (base >> 1) & 3
213}
214/// Pack a slice of ASCII `ACTGactg` characters into its packed 2-bit kmer representation.
215/// Other characters are silently converted.
216#[inline(always)]
217pub fn pack_kmer_lossy(slice: &[u8]) -> u64 {
218    let mut kmer = 0;
219    for (i, &base) in slice.iter().enumerate() {
220        kmer |= (pack_char_lossy(base) as u64) << (2 * i);
221    }
222    kmer
223}
224/// Pack a slice of ASCII `ACTGactg` characters into its packed 2-bit kmer representation.
225#[inline(always)]
226pub fn pack_kmer_u128_lossy(slice: &[u8]) -> u128 {
227    let mut kmer = 0;
228    for (i, &base) in slice.iter().enumerate() {
229        kmer |= (pack_char_lossy(base) as u128) << (2 * i);
230    }
231    kmer
232}
233
234/// Unpack a 2-bit DNA base into the corresponding `ACTG` character.
235#[inline(always)]
236pub fn unpack_base(base: u8) -> u8 {
237    debug_assert!(base < 4, "Base {base} is not <4.");
238    b"ACTG"[base as usize]
239}
240
241/// Unpack a 2-bit encoded kmer into corresponding `ACTG` characters.
242/// Slice the returned array to the correct length.
243#[inline(always)]
244pub fn unpack_kmer(kmer: u64) -> [u8; 32] {
245    std::array::from_fn(|i| unpack_base(((kmer >> (2 * i)) & 3) as u8))
246}
247/// Unpack a 2-bit encoded kmer into corresponding `ACTG` character.
248#[inline(always)]
249pub fn unpack_kmer_into_vec(kmer: u64, k: usize, out: &mut Vec<u8>) {
250    out.clear();
251    out.extend((0..k).map(|i| unpack_base(((kmer >> (2 * i)) & 3) as u8)));
252}
253/// Unpack a 2-bit encoded kmer into corresponding `ACTG` character.
254#[inline(always)]
255pub fn unpack_kmer_to_vec(kmer: u64, k: usize) -> Vec<u8> {
256    let mut out = vec![];
257    unpack_kmer_into_vec(kmer, k, &mut out);
258    out
259}
260
261/// Unpack a 2-bit encoded kmer into corresponding `ACTG` characters.
262/// Slice the returned array to the correct length.
263#[inline(always)]
264pub fn unpack_kmer_u128(kmer: u128) -> [u8; 64] {
265    std::array::from_fn(|i| unpack_base(((kmer >> (2 * i)) & 3) as u8))
266}
267/// Unpack a 2-bit encoded kmer into corresponding `ACTG` character.
268#[inline(always)]
269pub fn unpack_kmer_u128_into_vec(kmer: u128, k: usize, out: &mut Vec<u8>) {
270    out.clear();
271    out.extend((0..k).map(|i| unpack_base(((kmer >> (2 * i)) & 3) as u8)));
272}
273/// Unpack a 2-bit encoded kmer into corresponding `ACTG` character.
274#[inline(always)]
275pub fn unpack_kmer_u128_to_vec(kmer: u128, k: usize) -> Vec<u8> {
276    let mut out = vec![];
277    unpack_kmer_u128_into_vec(kmer, k, &mut out);
278    out
279}
280
281/// Complement an ASCII character: `A<>T` and `C<>G`.
282#[inline(always)]
283pub const fn complement_char(base: u8) -> u8 {
284    match base {
285        b'A' => b'T',
286        b'C' => b'G',
287        b'G' => b'C',
288        b'T' => b'A',
289        _ => panic!("Unexpected character. Expected one of ACTGactg.",),
290    }
291}
292
293/// Complement a 2-bit base: `0<>2` and `1<>3`.
294#[inline(always)]
295pub const fn complement_base(base: u8) -> u8 {
296    base ^ 2
297}
298
299/// Complement 8 lanes of 2-bit bases: `0<>2` and `1<>3`.
300#[inline(always)]
301pub fn complement_base_simd(base: u32x8) -> u32x8 {
302    const TWO: u32x8 = u32x8::new([2; 8]);
303    base ^ TWO
304}
305
306/// Reverse complement the 2-bit pairs in the input.
307#[inline(always)]
308const fn revcomp_raw(word: u64) -> u64 {
309    #[cfg(any(target_arch = "arm", target_arch = "aarch64"))]
310    {
311        let mut res = word.reverse_bits(); // ARM can reverse bits in a single instruction
312        res = ((res >> 1) & 0x5555_5555_5555_5555) | ((res & 0x5555_5555_5555_5555) << 1);
313        res ^ 0xAAAA_AAAA_AAAA_AAAA
314    }
315
316    #[cfg(not(any(target_arch = "arm", target_arch = "aarch64")))]
317    {
318        let mut res = word.swap_bytes();
319        res = ((res >> 4) & 0x0F0F_0F0F_0F0F_0F0F) | ((res & 0x0F0F_0F0F_0F0F_0F0F) << 4);
320        res = ((res >> 2) & 0x3333_3333_3333_3333) | ((res & 0x3333_3333_3333_3333) << 2);
321        res ^ 0xAAAA_AAAA_AAAA_AAAA
322    }
323}
324
325/// Compute the reverse complement of a short sequence packed in a `u64`.
326#[inline(always)]
327pub const fn revcomp_u64(word: u64, len: usize) -> u64 {
328    revcomp_raw(word) >> (u64::BITS as usize - 2 * len)
329}
330
331#[inline(always)]
332pub const fn revcomp_u128(word: u128, len: usize) -> u128 {
333    let low = word as u64;
334    let high = (word >> 64) as u64;
335    let rlow = revcomp_raw(low);
336    let rhigh = revcomp_raw(high);
337    let out = ((rlow as u128) << 64) | rhigh as u128;
338    out >> (u128::BITS as usize - 2 * len)
339}
340
341// ======================================================================
342// 1-BIT HELPER METHODS
343
344/// 1 when a char is ambiguous.
345#[inline(always)]
346pub fn char_is_ambiguous(base: u8) -> u8 {
347    // (!matches!(base, b'A' | b'C'  | b'G'  | b'T' | b'a' | b'c'  | b'g'  | b't')) as u8
348    let table = b"ACTG";
349    let upper_mask = !(b'a' - b'A');
350    (table[pack_char_lossy(base) as usize] != (base & upper_mask)) as u8
351}
352
353/// Reverse `len` bits packed in a `u64`.
354#[inline(always)]
355pub const fn rev_u64(word: u64, len: usize) -> u64 {
356    word.reverse_bits() >> (u64::BITS as usize - len)
357}
358
359/// Reverse `len` bits packed in a `u128`.
360#[inline(always)]
361pub const fn rev_u128(word: u128, len: usize) -> u128 {
362    word.reverse_bits() >> (u128::BITS as usize - len)
363}
364
365// ======================================================================
366
367impl<'s, const B: usize> PackedSeqBase<'s, B>
368where
369    Bits<B>: SupportedBits,
370{
371    /// Creates a `Seq` from a slice of packed bytes, an offset in bp and a length in bp.
372    ///
373    /// The slice should have at least 48 bytes of padding after `offset + len`.
374    /// Otherwise, the function will panic.
375    pub fn from_raw_parts(seq: &'s [u8], offset: usize, len: usize) -> Self {
376        assert!(offset + len + PADDING * Self::C8 <= seq.len() * Self::C8);
377        Self { seq, offset, len }
378    }
379
380    /// Shrink `seq` to only just cover the data.
381    #[inline(always)]
382    pub fn normalize(&self) -> Self {
383        let start_byte = self.offset / Self::C8;
384        let end_byte = (self.offset + self.len).div_ceil(Self::C8);
385        Self {
386            seq: &self.seq[start_byte..end_byte + PADDING],
387            offset: self.offset % Self::C8,
388            len: self.len,
389        }
390    }
391
392    /// Return a `Vec<u8>` of ASCII `ACTG` characters.
393    #[inline(always)]
394    pub fn unpack(&self) -> Vec<u8> {
395        self.iter_bp().map(unpack_base).collect()
396    }
397}
398
399/// Read up to 32 bytes starting at idx.
400#[inline(always)]
401pub(crate) unsafe fn read_slice_32_unchecked(seq: &[u8], idx: usize) -> u32x8 {
402    unsafe {
403        let src = seq.as_ptr().add(idx);
404        debug_assert!(idx + 32 <= seq.len());
405        std::mem::transmute::<_, *const u32x8>(src).read_unaligned()
406    }
407}
408
409/// Read up to 32 bytes starting at idx.
410#[inline(always)]
411pub(crate) fn read_slice_32(seq: &[u8], idx: usize) -> u32x8 {
412    unsafe {
413        let src = seq.as_ptr().add(idx);
414        if idx + 32 <= seq.len() {
415            std::mem::transmute::<_, *const u32x8>(src).read_unaligned()
416        } else {
417            let num_bytes = seq.len().saturating_sub(idx);
418            let mut result = [0u8; 32];
419            std::ptr::copy_nonoverlapping(src, result.as_mut_ptr(), num_bytes);
420            std::mem::transmute(result)
421        }
422    }
423}
424
425/// Read up to 16 bytes starting at idx.
426#[allow(unused)]
427#[inline(always)]
428pub(crate) fn read_slice_16(seq: &[u8], idx: usize) -> u16x8 {
429    unsafe {
430        let src = seq.as_ptr().add(idx);
431        if idx + 16 <= seq.len() {
432            std::mem::transmute::<_, *const u16x8>(src).read_unaligned()
433        } else {
434            let num_bytes = seq.len().saturating_sub(idx);
435            let mut result = [0u8; 16];
436            std::ptr::copy_nonoverlapping(src, result.as_mut_ptr(), num_bytes);
437            std::mem::transmute(result)
438        }
439    }
440}
441
442impl<'s, const B: usize> Seq<'s> for PackedSeqBase<'s, B>
443where
444    Bits<B>: SupportedBits,
445{
446    const BITS_PER_CHAR: usize = B;
447    const BASES_PER_BYTE: usize = Self::C8;
448    type SeqVec = PackedSeqVecBase<B>;
449
450    #[inline(always)]
451    fn len(&self) -> usize {
452        self.len
453    }
454
455    #[inline(always)]
456    fn is_empty(&self) -> bool {
457        self.len == 0
458    }
459
460    #[inline(always)]
461    fn get_ascii(&self, index: usize) -> u8 {
462        unpack_base(self.get(index))
463    }
464
465    /// Convert a short sequence (kmer) to a packed representation as `u64`.
466    /// Panics if `self` is longer than 32 characters.
467    #[inline(always)]
468    fn as_u64(&self) -> u64 {
469        assert!(self.len() <= 64 / B);
470
471        let mask = u64::MAX >> (64 - B * self.len());
472
473        // The unaligned read is OK, because we ensure that the underlying `PackedSeqVecBase::seq` always
474        // has at least 48 bytes of padding at the end.
475        if self.len() <= Self::K64 {
476            let x = unsafe { (self.seq.as_ptr() as *const u64).read_unaligned() };
477            (x >> (B * self.offset)) & mask
478        } else {
479            let x = unsafe { (self.seq.as_ptr() as *const u128).read_unaligned() };
480            (x >> (B * self.offset)) as u64 & mask
481        }
482    }
483
484    /// Convert a short sequence (kmer) to a packed representation of its reverse complement as `usize`.
485    /// Panics if `self` is longer than 32 characters.
486    #[inline(always)]
487    fn revcomp_as_u64(&self) -> u64 {
488        match B {
489            1 => rev_u64(self.as_u64(), self.len()),
490            2 => revcomp_u64(self.as_u64(), self.len()),
491            _ => panic!("Rev(comp) is only supported for 1-bit and 2-bit alphabets."),
492        }
493    }
494
495    /// Convert a short sequence (kmer) to a packed representation as `u128`.
496    /// Panics if `self` is longer than 64 characters.
497    #[inline(always)]
498    fn as_u128(&self) -> u128 {
499        assert!(
500            self.len() <= (128 - 8) / B + 1,
501            "Sequences >61 long cannot be read with a single unaligned u128 read."
502        );
503
504        let mask = u128::MAX >> (128 - B * self.len());
505
506        // The unaligned read is OK, because we ensure that the underlying `PackedSeqVecBase::seq` always
507        // has at least 48 bytes of padding at the end.
508        let x = unsafe { (self.seq.as_ptr() as *const u128).read_unaligned() };
509        (x >> (B * self.offset)) & mask
510    }
511
512    /// Convert a short sequence (kmer) to a packed representation of its reverse complement as `usize`.
513    /// Panics if `self` is longer than 64 characters.
514    #[inline(always)]
515    fn revcomp_as_u128(&self) -> u128 {
516        match B {
517            1 => rev_u128(self.as_u128(), self.len()),
518            2 => revcomp_u128(self.as_u128(), self.len()),
519            _ => panic!("Rev(comp) is only supported for 1-bit and 2-bit alphabets."),
520        }
521    }
522
523    #[inline(always)]
524    fn to_vec(&self) -> PackedSeqVecBase<B> {
525        assert_eq!(self.offset, 0);
526        PackedSeqVecBase {
527            seq: self
528                .seq
529                .iter()
530                .copied()
531                .chain(std::iter::repeat_n(0u8, PADDING))
532                .collect(),
533            len: self.len,
534        }
535    }
536
537    fn to_revcomp(&self) -> PackedSeqVecBase<B> {
538        match B {
539            1 | 2 => {}
540            _ => panic!("Can only reverse (&complement) 1-bit and 2-bit packed sequences.",),
541        }
542
543        let mut seq = self.seq[..(self.offset + self.len).div_ceil(Self::C8)]
544            .iter()
545            // 1. reverse the bytes
546            .rev()
547            .copied()
548            .map(|mut res| {
549                match B {
550                    2 => {
551                        // 2. swap the bases in the byte
552                        // This is auto-vectorized.
553                        res = ((res >> 4) & 0x0F) | ((res & 0x0F) << 4);
554                        res = ((res >> 2) & 0x33) | ((res & 0x33) << 2);
555                        // Complement the bases.
556                        res ^ 0xAA
557                    }
558                    1 => res.reverse_bits(),
559                    _ => unreachable!(),
560                }
561            })
562            .chain(std::iter::repeat_n(0u8, PADDING))
563            .collect::<Vec<u8>>();
564
565        // 3. Shift away the offset.
566        let new_offset = (Self::C8 - (self.offset + self.len) % Self::C8) % Self::C8;
567
568        if new_offset > 0 {
569            // Shift everything left by `2*new_offset` bits.
570            let shift = B * new_offset;
571            *seq.last_mut().unwrap() >>= shift;
572            // This loop is also auto-vectorized.
573            for i in 0..seq.len() - 1 {
574                seq[i] = (seq[i] >> shift) | (seq[i + 1] << (8 - shift));
575            }
576        }
577
578        PackedSeqVecBase { seq, len: self.len }
579    }
580
581    #[inline(always)]
582    fn slice(&self, range: Range<usize>) -> Self {
583        debug_assert!(
584            range.end <= self.len,
585            "Slice index out of bounds: {} > {}",
586            range.end,
587            self.len
588        );
589        PackedSeqBase {
590            seq: self.seq,
591            offset: self.offset + range.start,
592            len: range.end - range.start,
593        }
594        .normalize()
595    }
596
597    #[inline(always)]
598    fn iter_bp(self) -> impl ExactSizeIterator<Item = u8> {
599        assert!(self.len <= self.seq.len() * Self::C8);
600
601        let this = self.normalize();
602
603        // read u64 at a time?
604        let mut byte = 0;
605        (0..this.len + this.offset)
606            .map(
607                #[inline(always)]
608                move |i| {
609                    if i % Self::C8 == 0 {
610                        byte = this.seq[i / Self::C8];
611                    }
612                    // Shift byte instead of i?
613                    (byte >> (B * (i % Self::C8))) & Self::CHAR_MASK as u8
614                },
615            )
616            .advance(this.offset)
617    }
618
619    #[inline(always)]
620    fn par_iter_bp(self, context: usize) -> PaddedIt<impl ChunkIt<S>> {
621        self.par_iter_bp_with_buf(context, RecycledBox::take())
622    }
623
624    #[inline(always)]
625    fn par_iter_bp_delayed(self, context: usize, delay: Delay) -> PaddedIt<impl ChunkIt<(S, S)>> {
626        self.par_iter_bp_delayed_with_factor(context, delay, 1)
627    }
628
629    /// NOTE: When `self` starts does not start at a byte boundary, the
630    /// 'delayed' character is not guaranteed to be `0`.
631    #[inline(always)]
632    fn par_iter_bp_delayed_2(
633        self,
634        context: usize,
635        delay1: Delay,
636        delay2: Delay,
637    ) -> PaddedIt<impl ChunkIt<(S, S, S)>> {
638        self.par_iter_bp_delayed_2_with_factor_and_buf(
639            context,
640            delay1,
641            delay2,
642            1,
643            RecycledVec::take(),
644        )
645    }
646
647    /// Compares 29 characters at a time.
648    fn cmp_lcp(&self, other: &Self) -> (std::cmp::Ordering, usize) {
649        let mut lcp = 0;
650        let min_len = self.len.min(other.len);
651        for i in (0..min_len).step_by(Self::K64) {
652            let len = (min_len - i).min(Self::K64);
653            let this = self.slice(i..i + len);
654            let other = other.slice(i..i + len);
655            let this_word = this.as_u64();
656            let other_word = other.as_u64();
657            if this_word != other_word {
658                // Unfortunately, bases are packed in little endian order, so the default order is reversed.
659                let eq = this_word ^ other_word;
660                let t = eq.trailing_zeros() as usize / B * B;
661                lcp += t / B;
662                let mask = (Self::CHAR_MASK) << t;
663                return ((this_word & mask).cmp(&(other_word & mask)), lcp);
664            }
665            lcp += len;
666        }
667        (self.len.cmp(&other.len), lcp)
668    }
669
670    #[inline(always)]
671    fn get(&self, index: usize) -> u8 {
672        let offset = self.offset + index;
673        let idx = offset / Self::C8;
674        let offset = offset % Self::C8;
675        (self.seq[idx] >> (B * offset)) & Self::CHAR_MASK as u8
676    }
677}
678
679impl<'s, const B: usize> PackedSeqBase<'s, B>
680where
681    Bits<B>: SupportedBits,
682{
683    #[inline(always)]
684    pub fn par_iter_bp_with_buf<BUF: DerefMut<Target = [S; 8]>>(
685        self,
686        context: usize,
687        mut buf: BUF,
688    ) -> PaddedIt<impl ChunkIt<S> + use<'s, B, BUF>> {
689        #[cfg(target_endian = "big")]
690        panic!("Big endian architectures are not supported.");
691
692        let this = self.normalize();
693        let o = this.offset;
694        assert!(o < Self::C8);
695
696        let num_kmers = if this.len == 0 {
697            0
698        } else {
699            (this.len + o).saturating_sub(context - 1)
700        };
701        // without +o, since we don't need them in the stride.
702        let num_kmers_stride = this.len.saturating_sub(context - 1);
703        let n = num_kmers_stride.div_ceil(L).next_multiple_of(Self::C8);
704        let bytes_per_chunk = n / Self::C8;
705        let padding = Self::C8 * L * bytes_per_chunk - num_kmers_stride;
706
707        let offsets: [usize; 8] = from_fn(|l| l * bytes_per_chunk);
708        let mut cur = S::ZERO;
709
710        let par_len = if num_kmers == 0 {
711            0
712        } else {
713            n + context + o - 1
714        };
715
716        let last_i = par_len.saturating_sub(1) / Self::C32 * Self::C32;
717        // Safety check for the `read_slice_32_unchecked`:
718        assert!(offsets[7] + (last_i / Self::C8) + 32 <= this.seq.len());
719
720        let it = (0..par_len)
721            .map(
722                #[inline(always)]
723                move |i| {
724                    if i % Self::C32 == 0 {
725                        if i % Self::C256 == 0 {
726                            // Read a u256 for each lane containing the next 128 characters.
727                            let data: [u32x8; 8] = from_fn(
728                                #[inline(always)]
729                                |lane| unsafe {
730                                    read_slice_32_unchecked(
731                                        this.seq,
732                                        offsets[lane] + (i / Self::C8),
733                                    )
734                                },
735                            );
736                            *buf = transpose(data);
737                        }
738                        cur = buf[(i % Self::C256) / Self::C32];
739                    }
740                    // Extract the last 2 bits of each character.
741                    let chars = cur & Self::SIMD_CHAR_MASK;
742                    // Shift remaining characters to the right.
743                    cur = cur >> Self::SIMD_B;
744                    chars
745                },
746            )
747            .advance(o);
748
749        PaddedIt { it, padding }
750    }
751
752    #[inline(always)]
753    pub fn par_iter_bp_delayed_with_factor(
754        self,
755        context: usize,
756        delay: Delay,
757        factor: usize,
758    ) -> PaddedIt<impl ChunkIt<(S, S)> + use<'s, B>> {
759        self.par_iter_bp_delayed_with_factor_and_buf(context, delay, factor, RecycledVec::take())
760    }
761
762    #[inline(always)]
763    pub fn par_iter_bp_delayed_with_buf<BUF: DerefMut<Target = Vec<S>>>(
764        self,
765        context: usize,
766        delay: Delay,
767        buf: BUF,
768    ) -> PaddedIt<impl ChunkIt<(S, S)> + use<'s, B, BUF>> {
769        self.par_iter_bp_delayed_with_factor_and_buf(context, delay, 1, buf)
770    }
771
772    #[inline(always)]
773    pub fn par_iter_bp_delayed_with_factor_and_buf<BUF: DerefMut<Target = Vec<S>>>(
774        self,
775        context: usize,
776        Delay(delay): Delay,
777        factor: usize,
778        mut buf: BUF,
779    ) -> PaddedIt<impl ChunkIt<(S, S)> + use<'s, B, BUF>> {
780        #[cfg(target_endian = "big")]
781        panic!("Big endian architectures are not supported.");
782
783        assert!(
784            delay < usize::MAX / 2,
785            "Delay={} should be >=0.",
786            delay as isize
787        );
788
789        let this = self.normalize();
790        let o = this.offset;
791        assert!(o < Self::C8);
792
793        let num_kmers = if this.len == 0 {
794            0
795        } else {
796            (this.len + o).saturating_sub(context - 1)
797        };
798        // without +o, since we don't need them in the stride.
799        let num_kmers_stride = this.len.saturating_sub(context - 1);
800        let n = num_kmers_stride
801            .div_ceil(L)
802            .next_multiple_of(factor * Self::C8);
803        let bytes_per_chunk = n / Self::C8;
804        let padding = Self::C8 * L * bytes_per_chunk - num_kmers_stride;
805
806        let offsets: [usize; 8] = from_fn(|l| l * bytes_per_chunk);
807        let mut upcoming = S::ZERO;
808        let mut upcoming_d = S::ZERO;
809
810        // Even buf_len is nice to only have the write==buf_len check once.
811        // We also make it the next power of 2, for faster modulo operations.
812        // delay/16: number of bp in a u32.
813        // +8: some 'random' padding
814        let buf_len = (delay / Self::C32 + 8).next_power_of_two();
815        let buf_mask = buf_len - 1;
816        if buf.capacity() < buf_len {
817            // This has better codegen than `vec.clear(); vec.resize()`, since the inner `do_reserve_and_handle` of resize is not inlined.
818            *buf = vec![S::ZERO; buf_len];
819        } else {
820            // NOTE: Buf needs to be filled with zeros to guarantee returning 0 values for out-of-bounds characters.
821            buf.clear();
822            buf.resize(buf_len, S::ZERO);
823        }
824
825        let mut write_idx = 0;
826        // We compensate for the first delay/16 triggers of the check below that
827        // happen before the delay is actually reached.
828        let mut read_idx = (buf_len - delay / Self::C32) % buf_len;
829
830        let par_len = if num_kmers == 0 {
831            0
832        } else {
833            n + context + o - 1
834        };
835
836        let last_i = par_len.saturating_sub(1) / Self::C32 * Self::C32;
837        // Safety check for the `read_slice_32_unchecked`:
838        assert!(offsets[7] + (last_i / Self::C8) + 32 <= this.seq.len());
839
840        let it = (0..par_len)
841            .map(
842                #[inline(always)]
843                move |i| {
844                    if i % Self::C32 == 0 {
845                        if i % Self::C256 == 0 {
846                            // Read a u256 for each lane containing the next 128 characters.
847                            let data: [u32x8; 8] = from_fn(
848                                #[inline(always)]
849                                |lane| unsafe {
850                                    read_slice_32_unchecked(
851                                        this.seq,
852                                        offsets[lane] + (i / Self::C8),
853                                    )
854                                },
855                            );
856                            unsafe {
857                                *TryInto::<&mut [u32x8; 8]>::try_into(
858                                    buf.get_unchecked_mut(write_idx..write_idx + 8),
859                                )
860                                .unwrap_unchecked() = transpose(data);
861                            }
862                            if i == 0 {
863                                // Mask out chars before the offset.
864                                let elem = !((1u32 << (B * o)) - 1);
865                                let mask = S::splat(elem);
866                                unsafe { assert_unchecked(write_idx < buf.len()) };
867                                buf[write_idx] &= mask;
868                            }
869                        }
870                        unsafe { assert_unchecked(write_idx < buf.len()) };
871                        upcoming = buf[write_idx];
872                        write_idx += 1;
873                        write_idx &= buf_mask;
874                    }
875                    if i % Self::C32 == delay % Self::C32 {
876                        unsafe { assert_unchecked(read_idx < buf.len()) };
877                        upcoming_d = buf[read_idx];
878                        read_idx += 1;
879                        read_idx &= buf_mask;
880                    }
881                    // Extract the last 2 bits of each character.
882                    let chars = upcoming & Self::SIMD_CHAR_MASK;
883                    let chars_d = upcoming_d & Self::SIMD_CHAR_MASK;
884                    // Shift remaining characters to the right.
885                    upcoming = upcoming >> Self::SIMD_B;
886                    upcoming_d = upcoming_d >> Self::SIMD_B;
887                    (chars, chars_d)
888                },
889            )
890            .advance(o);
891
892        PaddedIt { it, padding }
893    }
894
895    #[inline(always)]
896    pub fn par_iter_bp_delayed_2_with_factor(
897        self,
898        context: usize,
899        delay1: Delay,
900        delay2: Delay,
901        factor: usize,
902    ) -> PaddedIt<impl ChunkIt<(S, S, S)> + use<'s, B>> {
903        self.par_iter_bp_delayed_2_with_factor_and_buf(
904            context,
905            delay1,
906            delay2,
907            factor,
908            RecycledVec::take(),
909        )
910    }
911
912    #[inline(always)]
913    pub fn par_iter_bp_delayed_2_with_buf<BUF: DerefMut<Target = Vec<S>>>(
914        self,
915        context: usize,
916        delay1: Delay,
917        delay2: Delay,
918        buf: BUF,
919    ) -> PaddedIt<impl ChunkIt<(S, S, S)> + use<'s, B, BUF>> {
920        self.par_iter_bp_delayed_2_with_factor_and_buf(context, delay1, delay2, 1, buf)
921    }
922
923    /// When iterating over 2-bit and 1-bit encoded data in parallel,
924    /// one must ensure that they have the same stride.
925    /// On the larger type, set `factor` as the ratio to the smaller one,
926    /// so that the stride in bytes is a multiple of `factor`,
927    /// so that the smaller type also has a byte-aligned stride.
928    #[inline(always)]
929    pub fn par_iter_bp_delayed_2_with_factor_and_buf<BUF: DerefMut<Target = Vec<S>>>(
930        self,
931        context: usize,
932        Delay(delay1): Delay,
933        Delay(delay2): Delay,
934        factor: usize,
935        mut buf: BUF,
936    ) -> PaddedIt<impl ChunkIt<(S, S, S)> + use<'s, B, BUF>> {
937        #[cfg(target_endian = "big")]
938        panic!("Big endian architectures are not supported.");
939
940        let this = self.normalize();
941        let o = this.offset;
942        assert!(o < Self::C8);
943        assert!(delay1 <= delay2, "Delay1 must be at most delay2.");
944
945        let num_kmers = if this.len == 0 {
946            0
947        } else {
948            (this.len + o).saturating_sub(context - 1)
949        };
950        // without +o, since we don't need them in the stride.
951        let num_kmers_stride = this.len.saturating_sub(context - 1);
952        let n = num_kmers_stride
953            .div_ceil(L)
954            .next_multiple_of(factor * Self::C8);
955        let bytes_per_chunk = n / Self::C8;
956        let padding = Self::C8 * L * bytes_per_chunk - num_kmers_stride;
957
958        let offsets: [usize; 8] = from_fn(|l| l * bytes_per_chunk);
959        let mut upcoming = S::ZERO;
960        let mut upcoming_d1 = S::ZERO;
961        let mut upcoming_d2 = S::ZERO;
962
963        // Even buf_len is nice to only have the write==buf_len check once.
964        let buf_len = (delay2 / Self::C32 + 8).next_power_of_two();
965        let buf_mask = buf_len - 1;
966        if buf.capacity() < buf_len {
967            // This has better codegen than `vec.clear(); vec.resize()`, since the inner `do_reserve_and_handle` of resize is not inlined.
968            *buf = vec![S::ZERO; buf_len];
969        } else {
970            // NOTE: Buf needs to be filled with zeros to guarantee returning 0 values for out-of-bounds characters.
971            buf.clear();
972            buf.resize(buf_len, S::ZERO);
973        }
974
975        let mut write_idx = 0;
976        // We compensate for the first delay/16 triggers of the check below that
977        // happen before the delay is actually reached.
978        let mut read_idx1 = (buf_len - delay1 / Self::C32) % buf_len;
979        let mut read_idx2 = (buf_len - delay2 / Self::C32) % buf_len;
980
981        let par_len = if num_kmers == 0 {
982            0
983        } else {
984            n + context + o - 1
985        };
986
987        let last_i = par_len.saturating_sub(1) / Self::C32 * Self::C32;
988        // Safety check for the `read_slice_32_unchecked`:
989        assert!(offsets[7] + (last_i / Self::C8) + 32 <= this.seq.len());
990
991        let it = (0..par_len)
992            .map(
993                #[inline(always)]
994                move |i| {
995                    if i % Self::C32 == 0 {
996                        if i % Self::C256 == 0 {
997                            // Read a u256 for each lane containing the next 128 characters.
998                            let data: [u32x8; 8] = from_fn(
999                                #[inline(always)]
1000                                |lane| unsafe {
1001                                    read_slice_32_unchecked(
1002                                        this.seq,
1003                                        offsets[lane] + (i / Self::C8),
1004                                    )
1005                                },
1006                            );
1007                            unsafe {
1008                                *TryInto::<&mut [u32x8; 8]>::try_into(
1009                                    buf.get_unchecked_mut(write_idx..write_idx + 8),
1010                                )
1011                                .unwrap_unchecked() = transpose(data);
1012                            }
1013                            // FIXME DROP THIS?
1014                            if i == 0 {
1015                                // Mask out chars before the offset.
1016                                let elem = !((1u32 << (B * o)) - 1);
1017                                let mask = S::splat(elem);
1018                                buf[write_idx] &= mask;
1019                            }
1020                        }
1021                        upcoming = buf[write_idx];
1022                        write_idx += 1;
1023                        write_idx &= buf_mask;
1024                    }
1025                    if i % Self::C32 == delay1 % Self::C32 {
1026                        unsafe { assert_unchecked(read_idx1 < buf.len()) };
1027                        upcoming_d1 = buf[read_idx1];
1028                        read_idx1 += 1;
1029                        read_idx1 &= buf_mask;
1030                    }
1031                    if i % Self::C32 == delay2 % Self::C32 {
1032                        unsafe { assert_unchecked(read_idx2 < buf.len()) };
1033                        upcoming_d2 = buf[read_idx2];
1034                        read_idx2 += 1;
1035                        read_idx2 &= buf_mask;
1036                    }
1037                    // Extract the last 2 bits of each character.
1038                    let chars = upcoming & Self::SIMD_CHAR_MASK;
1039                    let chars_d1 = upcoming_d1 & Self::SIMD_CHAR_MASK;
1040                    let chars_d2 = upcoming_d2 & Self::SIMD_CHAR_MASK;
1041                    // Shift remaining characters to the right.
1042                    upcoming = upcoming >> Self::SIMD_B;
1043                    upcoming_d1 = upcoming_d1 >> Self::SIMD_B;
1044                    upcoming_d2 = upcoming_d2 >> Self::SIMD_B;
1045                    (chars, chars_d1, chars_d2)
1046                },
1047            )
1048            .advance(o);
1049
1050        PaddedIt { it, padding }
1051    }
1052}
1053
1054impl<const B: usize> PartialEq for PackedSeqBase<'_, B>
1055where
1056    Bits<B>: SupportedBits,
1057{
1058    /// Compares 29 characters at a time.
1059    fn eq(&self, other: &Self) -> bool {
1060        if self.len != other.len {
1061            return false;
1062        }
1063        for i in (0..self.len).step_by(Self::K64) {
1064            let len = (self.len - i).min(Self::K64);
1065            let this = self.slice(i..i + len);
1066            let that = other.slice(i..i + len);
1067            if this.as_u64() != that.as_u64() {
1068                return false;
1069            }
1070        }
1071        true
1072    }
1073}
1074
1075impl<const B: usize> Eq for PackedSeqBase<'_, B> where Bits<B>: SupportedBits {}
1076
1077impl<const B: usize> PartialOrd for PackedSeqBase<'_, B>
1078where
1079    Bits<B>: SupportedBits,
1080{
1081    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
1082        Some(self.cmp(other))
1083    }
1084}
1085
1086impl<const B: usize> Ord for PackedSeqBase<'_, B>
1087where
1088    Bits<B>: SupportedBits,
1089{
1090    /// Compares 29 characters at a time.
1091    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
1092        let min_len = self.len.min(other.len);
1093        for i in (0..min_len).step_by(Self::K64) {
1094            let len = (min_len - i).min(Self::K64);
1095            let this = self.slice(i..i + len);
1096            let other = other.slice(i..i + len);
1097            let this_word = this.as_u64();
1098            let other_word = other.as_u64();
1099            if this_word != other_word {
1100                // Unfortunately, bases are packed in little endian order, so the default order is reversed.
1101                let eq = this_word ^ other_word;
1102                let t = eq.trailing_zeros() as usize / B * B;
1103                let mask = (Self::CHAR_MASK) << t;
1104                return (this_word & mask).cmp(&(other_word & mask));
1105            }
1106        }
1107        self.len.cmp(&other.len)
1108    }
1109}
1110
1111impl<const B: usize> SeqVec for PackedSeqVecBase<B>
1112where
1113    Bits<B>: SupportedBits,
1114{
1115    type Seq<'s> = PackedSeqBase<'s, B>;
1116
1117    #[inline(always)]
1118    fn into_raw(mut self) -> Vec<u8> {
1119        self.seq.resize(self.len.div_ceil(Self::C8), 0);
1120        self.seq
1121    }
1122
1123    #[inline(always)]
1124    fn as_slice(&self) -> Self::Seq<'_> {
1125        PackedSeqBase {
1126            seq: &self.seq[..self.len.div_ceil(Self::C8) + PADDING],
1127            offset: 0,
1128            len: self.len,
1129        }
1130    }
1131
1132    #[inline(always)]
1133    fn len(&self) -> usize {
1134        self.len
1135    }
1136
1137    #[inline(always)]
1138    fn is_empty(&self) -> bool {
1139        self.len == 0
1140    }
1141
1142    #[inline(always)]
1143    fn clear(&mut self) {
1144        self.seq.clear();
1145        self.len = 0;
1146    }
1147
1148    fn push_seq<'a>(&mut self, seq: PackedSeqBase<'_, B>) -> Range<usize> {
1149        let start = self.len.next_multiple_of(Self::C8) + seq.offset;
1150        let end = start + seq.len();
1151
1152        // Shrink away the padding.
1153        self.seq.resize(self.len.div_ceil(Self::C8), 0);
1154        // Extend.
1155        self.seq.extend(seq.seq);
1156        self.len = end;
1157        start..end
1158    }
1159
1160    /// For `PackedSeqVec` (2-bit encoding): map ASCII `ACGT` (and `acgt`) to DNA `0132` in that order.
1161    /// Other characters are silently mapped into `0..4`.
1162    ///
1163    /// Uses the BMI2 `pext` instruction when available, based on the
1164    /// `n_to_bits_pext` method described at
1165    /// <https://github.com/Daniel-Liu-c0deb0t/cute-nucleotides>.
1166    ///
1167    /// For `BitSeqVec` (1-bit encoding): map `ACGTacgt` to `0`, and everything else to `1`.
1168    fn push_ascii(&mut self, seq: &[u8]) -> Range<usize> {
1169        match B {
1170            1 | 2 => {}
1171            _ => panic!(
1172                "Can only use ASCII input for 2-bit DNA packing, or 1-bit ambiguous indicators."
1173            ),
1174        }
1175
1176        self.seq
1177            .resize((self.len + seq.len()).div_ceil(Self::C8) + PADDING, 0);
1178        let start_aligned = self.len.next_multiple_of(Self::C8);
1179        let start = self.len;
1180        let len = seq.len();
1181        let mut idx = self.len / Self::C8;
1182
1183        let parse_base = |base| match B {
1184            1 => char_is_ambiguous(base),
1185            2 => pack_char_lossy(base),
1186            _ => unreachable!(),
1187        };
1188
1189        let unaligned = core::cmp::min(start_aligned - start, len);
1190        if unaligned > 0 {
1191            let mut packed_byte = self.seq[idx];
1192            for &base in &seq[..unaligned] {
1193                packed_byte |= parse_base(base) << ((self.len % Self::C8) * B);
1194                self.len += 1;
1195            }
1196            self.seq[idx] = packed_byte;
1197            idx += 1;
1198        }
1199
1200        #[allow(unused)]
1201        let mut last = unaligned;
1202
1203        if B == 2 {
1204            #[cfg(all(target_arch = "x86_64", target_feature = "bmi2"))]
1205            {
1206                last = unaligned + (len - unaligned) / 8 * 8;
1207
1208                for i in (unaligned..last).step_by(8) {
1209                    let chunk =
1210                        unsafe { seq.get_unchecked(i..i + 8).try_into().unwrap_unchecked() };
1211                    let ascii = u64::from_le_bytes(chunk);
1212                    let packed_bytes =
1213                        unsafe { std::arch::x86_64::_pext_u64(ascii, 0x0606060606060606) } as u16;
1214                    unsafe {
1215                        self.seq
1216                            .get_unchecked_mut(idx..(idx + 2))
1217                            .copy_from_slice(&packed_bytes.to_le_bytes())
1218                    };
1219                    idx += 2;
1220                    self.len += 8;
1221                }
1222            }
1223
1224            #[cfg(target_feature = "neon")]
1225            {
1226                use core::arch::aarch64::{
1227                    vandq_u8, vdup_n_u8, vld1q_u8, vpadd_u8, vshlq_u8, vst1_u8,
1228                };
1229                use core::mem::transmute;
1230
1231                last = unaligned + (len - unaligned) / 16 * 16;
1232
1233                for i in (unaligned..last).step_by(16) {
1234                    unsafe {
1235                        let ascii = vld1q_u8(seq.as_ptr().add(i));
1236                        let masked_bits = vandq_u8(ascii, transmute([6i8; 16]));
1237                        let (bits_0, bits_1) = transmute(vshlq_u8(
1238                            masked_bits,
1239                            transmute([-1i8, 1, 3, 5, -1, 1, 3, 5, -1, 1, 3, 5, -1, 1, 3, 5]),
1240                        ));
1241                        let half_packed = vpadd_u8(bits_0, bits_1);
1242                        let packed = vpadd_u8(half_packed, vdup_n_u8(0));
1243                        vst1_u8(self.seq.as_mut_ptr().add(idx), packed);
1244                        idx += Self::C8;
1245                        self.len += 16;
1246                    }
1247                }
1248            }
1249        }
1250
1251        if B == 1 {
1252            #[cfg(all(target_arch = "x86_64", target_feature = "avx2"))]
1253            {
1254                last = len;
1255                self.len += len - unaligned;
1256
1257                let mut last_i = unaligned;
1258
1259                for i in (unaligned..last).step_by(32) {
1260                    use std::mem::transmute as t;
1261
1262                    // Wide doesn't have u8x32, so this is messy here...
1263                    type S = wide::i8x32;
1264                    let chars: S = unsafe { t(read_slice_32(seq, i)) };
1265                    let upper_mask = !(b'a' - b'A');
1266                    // make everything upper case
1267                    let chars = chars & S::splat(upper_mask as i8);
1268                    let lossy_encoded = chars & S::splat(6);
1269                    let table = unsafe { S::from(t::<_, S>(*b"AxCxTxGxxxxxxxxxAxCxTxGxxxxxxxxx")) };
1270                    let lookup: S = unsafe {
1271                        t(std::arch::x86_64::_mm256_shuffle_epi8(
1272                            t(table),
1273                            t(lossy_encoded),
1274                        ))
1275                    };
1276                    let packed_bytes = !(chars.simd_eq(lookup).to_bitmask() as u32);
1277
1278                    last_i = i;
1279                    unsafe {
1280                        self.seq
1281                            .get_unchecked_mut(idx..(idx + 4))
1282                            .copy_from_slice(&packed_bytes.to_le_bytes())
1283                    };
1284                    idx += 4;
1285                }
1286
1287                // Fix up trailing bytes.
1288                if unaligned < last {
1289                    idx -= 4;
1290                    let mut val = unsafe {
1291                        u32::from_le_bytes(
1292                            self.seq
1293                                .get_unchecked(idx..(idx + 4))
1294                                .try_into()
1295                                .unwrap_unchecked(),
1296                        )
1297                    };
1298                    // keep only the `last - last_i` low bits.
1299                    let keep = last - last_i;
1300                    val <<= 32 - keep;
1301                    val >>= 32 - keep;
1302                    unsafe {
1303                        self.seq
1304                            .get_unchecked_mut(idx..(idx + 4))
1305                            .copy_from_slice(&val.to_le_bytes())
1306                    };
1307                    idx += keep.div_ceil(8);
1308                }
1309            }
1310
1311            #[cfg(target_feature = "neon")]
1312            {
1313                use core::arch::aarch64::*;
1314                use core::mem::transmute;
1315
1316                last = unaligned + (len - unaligned) / 64 * 64;
1317
1318                for i in (unaligned..last).step_by(64) {
1319                    unsafe {
1320                        let ptr = seq.as_ptr().add(i);
1321                        let chars = vld4q_u8(ptr);
1322
1323                        let upper_mask = vdupq_n_u8(!(b'a' - b'A'));
1324                        let chars = neon::map_8x16x4(chars, |v| vandq_u8(v, upper_mask));
1325
1326                        let two_bits_mask = vdupq_n_u8(6);
1327                        let lossy_encoded = neon::map_8x16x4(chars, |v| vandq_u8(v, two_bits_mask));
1328
1329                        let table = transmute(*b"AxCxTxGxxxxxxxxx");
1330                        let lookup = neon::map_8x16x4(lossy_encoded, |v| vqtbl1q_u8(table, v));
1331
1332                        let mask = neon::map_two_8x16x4(chars, lookup, |v1, v2| vceqq_u8(v1, v2));
1333                        let packed_bytes = !neon::movemask_64(mask);
1334
1335                        self.seq[idx..(idx + 8)].copy_from_slice(&packed_bytes.to_le_bytes());
1336                        idx += 8;
1337                        self.len += 64;
1338                    }
1339                }
1340            }
1341        }
1342
1343        let mut packed_byte = 0;
1344        for &base in &seq[last..] {
1345            packed_byte |= parse_base(base) << ((self.len % Self::C8) * B);
1346            self.len += 1;
1347            if self.len % Self::C8 == 0 {
1348                self.seq[idx] = packed_byte;
1349                idx += 1;
1350                packed_byte = 0;
1351            }
1352        }
1353        if self.len % Self::C8 != 0 && last < len {
1354            self.seq[idx] = packed_byte;
1355            idx += 1;
1356        }
1357        assert_eq!(idx + PADDING, self.seq.len());
1358        start..start + len
1359    }
1360
1361    #[cfg(feature = "rand")]
1362    fn random(n: usize) -> Self {
1363        use rand::Rng;
1364
1365        let byte_len = n.div_ceil(Self::C8);
1366        let mut seq = vec![0; byte_len + PADDING];
1367        rand::make_rng::<rand::rngs::SmallRng>().fill_bytes(&mut seq[..byte_len]);
1368        // Ensure that the last byte is padded with zeros.
1369        if n % Self::C8 != 0 {
1370            seq[byte_len - 1] &= (1 << (B * (n % Self::C8))) - 1;
1371        }
1372
1373        Self { seq, len: n }
1374    }
1375}
1376
1377impl<const B: usize> PackedSeqVecBase<B>
1378where
1379    Bits<B>: SupportedBits,
1380{
1381    /// Creates a `SeqVec` from a vector of packed bytes and a length in bp.
1382    ///
1383    /// The vector should have at least 48 bytes of padding after `len`.
1384    /// Otherwise, the vector will be resized to be padded with zeros.
1385    pub fn from_raw_parts(mut seq: Vec<u8>, len: usize) -> Self {
1386        assert!(len <= seq.len() * Self::C8);
1387        seq.resize(len.div_ceil(Self::C8) + PADDING, 0);
1388        Self { seq, len }
1389    }
1390}
1391
1392impl PackedSeqVecBase<1> {
1393    pub fn with_len(n: usize) -> Self {
1394        Self {
1395            seq: vec![0; n.div_ceil(Self::C8) + PADDING],
1396            len: n,
1397        }
1398    }
1399
1400    #[cfg(feature = "rand")]
1401    pub fn random(len: usize, n_frac: f32) -> Self {
1402        let byte_len = len.div_ceil(Self::C8);
1403        let mut seq = vec![0; byte_len + PADDING];
1404
1405        assert!(
1406            (0.0..=0.3).contains(&n_frac),
1407            "n_frac={} should be in [0, 0.3]",
1408            n_frac
1409        );
1410
1411        for _ in 0..(len as f32 * n_frac) as usize {
1412            let idx = rand::random::<u64>() as usize % len;
1413            let byte = idx / Self::C8;
1414            let offset = idx % Self::C8;
1415            seq[byte] |= 1 << offset;
1416        }
1417
1418        Self { seq, len }
1419    }
1420}
1421
1422impl<'s> PackedSeqBase<'s, 1> {
1423    /// An iterator indicating for each kmer whether it contains ambiguous bases.
1424    ///
1425    /// Returns n-(k-1) elements.
1426    #[inline(always)]
1427    pub fn iter_kmer_ambiguity(self, k: usize) -> impl ExactSizeIterator<Item = bool> + use<'s> {
1428        let this = self.normalize();
1429        assert!(k > 0);
1430        assert!(k <= Self::K64);
1431        (this.offset..this.offset + this.len.saturating_sub(k - 1))
1432            .map(move |i| self.read_kmer(k, i) != 0)
1433    }
1434
1435    /// A parallel iterator indicating for each kmer whether it contains ambiguous bases.
1436    ///
1437    /// First element is the 'kmer' consisting only of the first character of each chunk.
1438    ///
1439    /// `k`: length of windows to check
1440    /// `context`: number of overlapping iterations +1. To determine stride of each lane.
1441    /// `skip`: Set to `context-1` to skip the iterations added by the context.
1442    #[inline(always)]
1443    pub fn par_iter_kmer_ambiguity(
1444        self,
1445        k: usize,
1446        context: usize,
1447        skip: usize,
1448    ) -> PaddedIt<impl ChunkIt<S> + use<'s>> {
1449        #[cfg(target_endian = "big")]
1450        panic!("Big endian architectures are not supported.");
1451
1452        assert!(k > 0, "par_iter_kmers requires k>0, but k={k}");
1453        assert!(k <= 96, "par_iter_kmers requires k<=96, but k={k}");
1454
1455        let this = self.normalize();
1456        let o = this.offset;
1457        assert!(o < Self::C8);
1458
1459        let delay = k - 1;
1460
1461        let it = self.par_iter_bp_delayed(context, Delay(delay));
1462
1463        let mut cnt = u32x8::ZERO;
1464
1465        it.map(
1466            #[inline(always)]
1467            move |(a, r)| {
1468                cnt += a;
1469                let out = cnt.simd_gt(S::ZERO);
1470                cnt -= r;
1471                out
1472            },
1473        )
1474        .advance(skip)
1475    }
1476
1477    #[inline(always)]
1478    pub fn par_iter_kmer_ambiguity_with_buf(
1479        self,
1480        k: usize,
1481        context: usize,
1482        skip: usize,
1483        buf: &'s mut Vec<S>,
1484    ) -> PaddedIt<impl ChunkIt<S> + use<'s>> {
1485        #[cfg(target_endian = "big")]
1486        panic!("Big endian architectures are not supported.");
1487
1488        assert!(k > 0, "par_iter_kmers requires k>0, but k={k}");
1489        assert!(k <= 96, "par_iter_kmers requires k<=96, but k={k}");
1490
1491        let this = self.normalize();
1492        let o = this.offset;
1493        assert!(o < Self::C8);
1494
1495        let delay = k - 1;
1496
1497        let it = self.par_iter_bp_delayed_with_buf(context, Delay(delay), buf);
1498
1499        let mut cnt = u32x8::ZERO;
1500
1501        it.map(
1502            #[inline(always)]
1503            move |(a, r)| {
1504                cnt += a;
1505                let out = cnt.simd_gt(S::ZERO);
1506                cnt -= r;
1507                out
1508            },
1509        )
1510        .advance(skip)
1511    }
1512}
1513
1514#[cfg(target_feature = "neon")]
1515mod neon {
1516    use core::arch::aarch64::*;
1517
1518    #[inline(always)]
1519    pub fn movemask_64(v: uint8x16x4_t) -> u64 {
1520        // https://stackoverflow.com/questions/74722950/convert-vector-compare-mask-into-bit-mask-in-aarch64-simd-or-arm-neon/74748402#74748402
1521        unsafe {
1522            let acc = vsriq_n_u8(vsriq_n_u8(v.3, v.2, 1), vsriq_n_u8(v.1, v.0, 1), 2);
1523            vget_lane_u64(
1524                vreinterpret_u64_u8(vshrn_n_u16(
1525                    vreinterpretq_u16_u8(vsriq_n_u8(acc, acc, 4)),
1526                    4,
1527                )),
1528                0,
1529            )
1530        }
1531    }
1532
1533    #[inline(always)]
1534    pub fn map_8x16x4<F>(v: uint8x16x4_t, mut f: F) -> uint8x16x4_t
1535    where
1536        F: FnMut(uint8x16_t) -> uint8x16_t,
1537    {
1538        uint8x16x4_t(f(v.0), f(v.1), f(v.2), f(v.3))
1539    }
1540
1541    #[inline(always)]
1542    pub fn map_two_8x16x4<F>(v1: uint8x16x4_t, v2: uint8x16x4_t, mut f: F) -> uint8x16x4_t
1543    where
1544        F: FnMut(uint8x16_t, uint8x16_t) -> uint8x16_t,
1545    {
1546        uint8x16x4_t(f(v1.0, v2.0), f(v1.1, v2.1), f(v1.2, v2.2), f(v1.3, v2.3))
1547    }
1548}