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
114pub(crate) const PADDING: usize = 48;
116
117#[derive(Copy, Clone, Debug, MemSize, MemDbg)]
119pub struct PackedSeqBase<'s, const B: usize>
120where
121 Bits<B>: SupportedBits,
122{
123 seq: &'s [u8],
125 offset: usize,
127 len: usize,
129}
130
131#[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 pub(crate) seq: Vec<u8>,
140
141 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
150impl<'s, const B: usize> PackedSeqBase<'s, B>
153where
154 Bits<B>: SupportedBits,
155{
156 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 const C8: usize = 8 / B;
162 const C32: usize = 32 / B;
164 const C256: usize = 256 / B;
166 const K64: usize = (64 - 8) / B + 1;
168}
169
170impl<const B: usize> PackedSeqVecBase<B>
172where
173 Bits<B>: SupportedBits,
174{
175 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#[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#[inline(always)]
211pub fn pack_char_lossy(base: u8) -> u8 {
212 (base >> 1) & 3
213}
214#[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#[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#[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#[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#[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#[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#[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#[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#[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#[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#[inline(always)]
295pub const fn complement_base(base: u8) -> u8 {
296 base ^ 2
297}
298
299#[inline(always)]
301pub fn complement_base_simd(base: u32x8) -> u32x8 {
302 const TWO: u32x8 = u32x8::new([2; 8]);
303 base ^ TWO
304}
305
306#[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(); 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#[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#[inline(always)]
346pub fn char_is_ambiguous(base: u8) -> u8 {
347 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#[inline(always)]
355pub const fn rev_u64(word: u64, len: usize) -> u64 {
356 word.reverse_bits() >> (u64::BITS as usize - len)
357}
358
359#[inline(always)]
361pub const fn rev_u128(word: u128, len: usize) -> u128 {
362 word.reverse_bits() >> (u128::BITS as usize - len)
363}
364
365impl<'s, const B: usize> PackedSeqBase<'s, B>
368where
369 Bits<B>: SupportedBits,
370{
371 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 #[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 #[inline(always)]
394 pub fn unpack(&self) -> Vec<u8> {
395 self.iter_bp().map(unpack_base).collect()
396 }
397}
398
399#[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#[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#[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 #[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 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 #[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 #[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 let x = unsafe { (self.seq.as_ptr() as *const u128).read_unaligned() };
509 (x >> (B * self.offset)) & mask
510 }
511
512 #[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 .rev()
547 .copied()
548 .map(|mut res| {
549 match B {
550 2 => {
551 res = ((res >> 4) & 0x0F) | ((res & 0x0F) << 4);
554 res = ((res >> 2) & 0x33) | ((res & 0x33) << 2);
555 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 let new_offset = (Self::C8 - (self.offset + self.len) % Self::C8) % Self::C8;
567
568 if new_offset > 0 {
569 let shift = B * new_offset;
571 *seq.last_mut().unwrap() >>= shift;
572 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 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 (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 #[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 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 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 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 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 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 let chars = cur & Self::SIMD_CHAR_MASK;
742 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 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 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 *buf = vec![S::ZERO; buf_len];
819 } else {
820 buf.clear();
822 buf.resize(buf_len, S::ZERO);
823 }
824
825 let mut write_idx = 0;
826 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 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 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 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 let chars = upcoming & Self::SIMD_CHAR_MASK;
883 let chars_d = upcoming_d & Self::SIMD_CHAR_MASK;
884 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 #[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 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 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 *buf = vec![S::ZERO; buf_len];
969 } else {
970 buf.clear();
972 buf.resize(buf_len, S::ZERO);
973 }
974
975 let mut write_idx = 0;
976 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 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 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 if i == 0 {
1015 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 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 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 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 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 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 self.seq.resize(self.len.div_ceil(Self::C8), 0);
1154 self.seq.extend(seq.seq);
1156 self.len = end;
1157 start..end
1158 }
1159
1160 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 type S = wide::i8x32;
1264 let chars: S = unsafe { t(read_slice_32(seq, i)) };
1265 let upper_mask = !(b'a' - b'A');
1266 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 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 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 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 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 #[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 #[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 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}