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 = "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 const K128: usize = (128 - 8) / B + 1;
170}
171
172impl<const B: usize> PackedSeqVecBase<B>
174where
175 Bits<B>: SupportedBits,
176{
177 const C8: usize = 8 / B;
179}
180
181impl<const B: usize> Default for PackedSeqVecBase<B>
182where
183 Bits<B>: SupportedBits,
184{
185 fn default() -> Self {
186 Self {
187 seq: vec![0; PADDING],
188 len: 0,
189 }
190 }
191}
192
193#[inline(always)]
198pub fn pack_char(base: u8) -> u8 {
199 match base {
200 b'a' | b'A' => 0,
201 b'c' | b'C' => 1,
202 b'g' | b'G' => 3,
203 b't' | b'T' => 2,
204 _ => panic!(
205 "Unexpected character '{}' with ASCII value {base}. Expected one of ACTGactg.",
206 base as char
207 ),
208 }
209}
210
211#[inline(always)]
213pub fn pack_char_lossy(base: u8) -> u8 {
214 (base >> 1) & 3
215}
216#[inline(always)]
219pub fn pack_kmer_lossy(slice: &[u8]) -> u64 {
220 let mut kmer = 0;
221 for (i, &base) in slice.iter().enumerate() {
222 kmer |= (pack_char_lossy(base) as u64) << (2 * i);
223 }
224 kmer
225}
226#[inline(always)]
228pub fn pack_kmer_u128_lossy(slice: &[u8]) -> u128 {
229 let mut kmer = 0;
230 for (i, &base) in slice.iter().enumerate() {
231 kmer |= (pack_char_lossy(base) as u128) << (2 * i);
232 }
233 kmer
234}
235
236#[inline(always)]
238pub fn unpack_base(base: u8) -> u8 {
239 debug_assert!(base < 4, "Base {base} is not <4.");
240 b"ACTG"[base as usize]
241}
242
243#[inline(always)]
246pub fn unpack_kmer(kmer: u64) -> [u8; 32] {
247 std::array::from_fn(|i| unpack_base(((kmer >> (2 * i)) & 3) as u8))
248}
249#[inline(always)]
251pub fn unpack_kmer_into_vec(kmer: u64, k: usize, out: &mut Vec<u8>) {
252 out.clear();
253 out.extend((0..k).map(|i| unpack_base(((kmer >> (2 * i)) & 3) as u8)));
254}
255#[inline(always)]
257pub fn unpack_kmer_to_vec(kmer: u64, k: usize) -> Vec<u8> {
258 let mut out = vec![];
259 unpack_kmer_into_vec(kmer, k, &mut out);
260 out
261}
262
263#[inline(always)]
266pub fn unpack_kmer_u128(kmer: u128) -> [u8; 64] {
267 std::array::from_fn(|i| unpack_base(((kmer >> (2 * i)) & 3) as u8))
268}
269#[inline(always)]
271pub fn unpack_kmer_u128_into_vec(kmer: u128, k: usize, out: &mut Vec<u8>) {
272 out.clear();
273 out.extend((0..k).map(|i| unpack_base(((kmer >> (2 * i)) & 3) as u8)));
274}
275#[inline(always)]
277pub fn unpack_kmer_u128_to_vec(kmer: u128, k: usize) -> Vec<u8> {
278 let mut out = vec![];
279 unpack_kmer_u128_into_vec(kmer, k, &mut out);
280 out
281}
282
283#[inline(always)]
285pub const fn complement_char(base: u8) -> u8 {
286 match base {
287 b'A' => b'T',
288 b'C' => b'G',
289 b'G' => b'C',
290 b'T' => b'A',
291 _ => panic!("Unexpected character. Expected one of ACTGactg.",),
292 }
293}
294
295#[inline(always)]
297pub const fn complement_base(base: u8) -> u8 {
298 base ^ 2
299}
300
301#[inline(always)]
303pub fn complement_base_simd(base: u32x8) -> u32x8 {
304 const TWO: u32x8 = u32x8::new([2; 8]);
305 base ^ TWO
306}
307
308#[inline(always)]
310const fn revcomp_raw(word: u64) -> u64 {
311 #[cfg(any(target_arch = "arm", target_arch = "aarch64"))]
312 {
313 let mut res = word.reverse_bits(); res = ((res >> 1) & 0x5555_5555_5555_5555) | ((res & 0x5555_5555_5555_5555) << 1);
315 res ^ 0xAAAA_AAAA_AAAA_AAAA
316 }
317
318 #[cfg(not(any(target_arch = "arm", target_arch = "aarch64")))]
319 {
320 let mut res = word.swap_bytes();
321 res = ((res >> 4) & 0x0F0F_0F0F_0F0F_0F0F) | ((res & 0x0F0F_0F0F_0F0F_0F0F) << 4);
322 res = ((res >> 2) & 0x3333_3333_3333_3333) | ((res & 0x3333_3333_3333_3333) << 2);
323 res ^ 0xAAAA_AAAA_AAAA_AAAA
324 }
325}
326
327#[inline(always)]
329pub const fn revcomp_u64(word: u64, len: usize) -> u64 {
330 revcomp_raw(word) >> (u64::BITS as usize - 2 * len)
331}
332
333#[inline(always)]
334pub const fn revcomp_u128(word: u128, len: usize) -> u128 {
335 let low = word as u64;
336 let high = (word >> 64) as u64;
337 let rlow = revcomp_raw(low);
338 let rhigh = revcomp_raw(high);
339 let out = ((rlow as u128) << 64) | rhigh as u128;
340 out >> (u128::BITS as usize - 2 * len)
341}
342
343#[inline(always)]
348pub fn char_is_ambiguous(base: u8) -> u8 {
349 let table = b"ACTG";
351 let upper_mask = !(b'a' - b'A');
352 (table[pack_char_lossy(base) as usize] != (base & upper_mask)) as u8
353}
354
355#[inline(always)]
357pub const fn rev_u64(word: u64, len: usize) -> u64 {
358 word.reverse_bits() >> (u64::BITS as usize - len)
359}
360
361#[inline(always)]
363pub const fn rev_u128(word: u128, len: usize) -> u128 {
364 word.reverse_bits() >> (u128::BITS as usize - len)
365}
366
367impl<'s, const B: usize> PackedSeqBase<'s, B>
370where
371 Bits<B>: SupportedBits,
372{
373 pub fn from_raw_parts(seq: &'s [u8], offset: usize, len: usize) -> Self {
378 assert!(offset + len + PADDING * Self::C8 <= seq.len() * Self::C8);
379 Self { seq, offset, len }
380 }
381
382 #[inline(always)]
384 pub fn normalize(&self) -> Self {
385 let start_byte = self.offset / Self::C8;
386 let end_byte = (self.offset + self.len).div_ceil(Self::C8);
387 Self {
388 seq: &self.seq[start_byte..end_byte + PADDING],
389 offset: self.offset % Self::C8,
390 len: self.len,
391 }
392 }
393
394 #[inline(always)]
396 pub fn unpack(&self) -> Vec<u8> {
397 self.iter_bp().map(unpack_base).collect()
398 }
399}
400
401impl<'s> PackedSeqBase<'s, 2> {
402 #[inline(always)]
404 pub fn iter_bp_rc(self) -> impl ExactSizeIterator<Item = u8> + 's {
405 (0..self.len)
406 .rev()
407 .map(move |i| complement_base(self.get(i)))
408 }
409
410 pub fn unpack_into(&self, out: &mut Vec<u8>) {
412 out.extend((0..self.len).map(|i| unpack_base(self.get(i))));
413 }
414
415 pub fn unpack_rc_into(&self, out: &mut Vec<u8>) {
417 out.extend(
418 (0..self.len)
419 .rev()
420 .map(|i| unpack_base(complement_base(self.get(i)))),
421 );
422 }
423}
424
425#[inline(always)]
427pub(crate) unsafe fn read_slice_32_unchecked(seq: &[u8], idx: usize) -> u32x8 {
428 unsafe {
429 let src = seq.as_ptr().add(idx);
430 debug_assert!(idx + 32 <= seq.len());
431 std::mem::transmute::<_, *const u32x8>(src).read_unaligned()
432 }
433}
434
435#[inline(always)]
437pub(crate) fn read_slice_32(seq: &[u8], idx: usize) -> u32x8 {
438 unsafe {
439 let src = seq.as_ptr().add(idx);
440 if idx + 32 <= seq.len() {
441 std::mem::transmute::<_, *const u32x8>(src).read_unaligned()
442 } else {
443 let num_bytes = seq.len().saturating_sub(idx);
444 let mut result = [0u8; 32];
445 std::ptr::copy_nonoverlapping(src, result.as_mut_ptr(), num_bytes);
446 std::mem::transmute(result)
447 }
448 }
449}
450
451#[allow(unused)]
453#[inline(always)]
454pub(crate) fn read_slice_16(seq: &[u8], idx: usize) -> u16x8 {
455 unsafe {
456 let src = seq.as_ptr().add(idx);
457 if idx + 16 <= seq.len() {
458 std::mem::transmute::<_, *const u16x8>(src).read_unaligned()
459 } else {
460 let num_bytes = seq.len().saturating_sub(idx);
461 let mut result = [0u8; 16];
462 std::ptr::copy_nonoverlapping(src, result.as_mut_ptr(), num_bytes);
463 std::mem::transmute(result)
464 }
465 }
466}
467
468impl<'s, const B: usize> Seq<'s> for PackedSeqBase<'s, B>
469where
470 Bits<B>: SupportedBits,
471{
472 const BITS_PER_CHAR: usize = B;
473 const BASES_PER_BYTE: usize = Self::C8;
474 type SeqVec = PackedSeqVecBase<B>;
475
476 #[inline(always)]
477 fn len(&self) -> usize {
478 self.len
479 }
480
481 #[inline(always)]
482 fn is_empty(&self) -> bool {
483 self.len == 0
484 }
485
486 #[inline(always)]
487 fn get_ascii(&self, index: usize) -> u8 {
488 unpack_base(self.get(index))
489 }
490
491 #[inline(always)]
494 fn as_u64(&self) -> u64 {
495 assert!(self.len() <= 64 / B);
496
497 let mask = u64::MAX >> (64 - B * self.len());
498
499 if self.len() <= Self::K64 {
502 let x = unsafe { (self.seq.as_ptr() as *const u64).read_unaligned() };
503 (x >> (B * self.offset)) & mask
504 } else {
505 let x = unsafe { (self.seq.as_ptr() as *const u128).read_unaligned() };
506 (x >> (B * self.offset)) as u64 & mask
507 }
508 }
509
510 #[inline(always)]
513 fn revcomp_as_u64(&self) -> u64 {
514 match B {
515 1 => rev_u64(self.as_u64(), self.len()),
516 2 => revcomp_u64(self.as_u64(), self.len()),
517 _ => panic!("Rev(comp) is only supported for 1-bit and 2-bit alphabets."),
518 }
519 }
520
521 #[inline(always)]
524 fn as_u128(&self) -> u128 {
525 let mask = u128::MAX >> (128 - B * self.len());
526
527 if self.len() > Self::K128 {
529 assert!(self.len() <= 128 / B);
530 if self.offset + self.len() > 128 / B {
531 let x = unsafe { (self.seq.as_ptr() as *const u128).read_unaligned() };
532 let y = unsafe { self.seq.as_ptr().offset(16).read() as u128 };
533 return ((x >> (B * self.offset)) | (y << (128 - B * self.offset))) & mask;
534 }
535 }
536
537 let x = unsafe { (self.seq.as_ptr() as *const u128).read_unaligned() };
540 (x >> (B * self.offset)) & mask
541 }
542
543 #[inline(always)]
546 fn revcomp_as_u128(&self) -> u128 {
547 match B {
548 1 => rev_u128(self.as_u128(), self.len()),
549 2 => revcomp_u128(self.as_u128(), self.len()),
550 _ => panic!("Rev(comp) is only supported for 1-bit and 2-bit alphabets."),
551 }
552 }
553
554 #[inline(always)]
555 fn to_vec(&self) -> PackedSeqVecBase<B> {
556 assert_eq!(self.offset, 0);
557 PackedSeqVecBase {
558 seq: self
559 .seq
560 .iter()
561 .copied()
562 .chain(std::iter::repeat_n(0u8, PADDING))
563 .collect(),
564 len: self.len,
565 }
566 }
567
568 fn to_revcomp(&self) -> PackedSeqVecBase<B> {
569 match B {
570 1 | 2 => {}
571 _ => panic!("Can only reverse (&complement) 1-bit and 2-bit packed sequences.",),
572 }
573
574 let mut seq = self.seq[..(self.offset + self.len).div_ceil(Self::C8)]
575 .iter()
576 .rev()
578 .copied()
579 .map(|mut res| {
580 match B {
581 2 => {
582 res = ((res >> 4) & 0x0F) | ((res & 0x0F) << 4);
585 res = ((res >> 2) & 0x33) | ((res & 0x33) << 2);
586 res ^ 0xAA
588 }
589 1 => res.reverse_bits(),
590 _ => unreachable!(),
591 }
592 })
593 .chain(std::iter::repeat_n(0u8, PADDING))
594 .collect::<Vec<u8>>();
595
596 let new_offset = (Self::C8 - (self.offset + self.len) % Self::C8) % Self::C8;
598
599 if new_offset > 0 {
600 let shift = B * new_offset;
602 *seq.last_mut().unwrap() >>= shift;
603 for i in 0..seq.len() - 1 {
605 seq[i] = (seq[i] >> shift) | (seq[i + 1] << (8 - shift));
606 }
607 }
608
609 PackedSeqVecBase { seq, len: self.len }
610 }
611
612 #[inline(always)]
613 fn slice(&self, range: Range<usize>) -> Self {
614 debug_assert!(
615 range.end <= self.len,
616 "Slice index out of bounds: {} > {}",
617 range.end,
618 self.len
619 );
620 PackedSeqBase {
621 seq: self.seq,
622 offset: self.offset + range.start,
623 len: range.end - range.start,
624 }
625 .normalize()
626 }
627
628 #[inline(always)]
629 fn iter_bp(self) -> impl ExactSizeIterator<Item = u8> {
630 assert!(self.len <= self.seq.len() * Self::C8);
631
632 let this = self.normalize();
633
634 let mut byte = 0;
636 (0..this.len + this.offset)
637 .map(
638 #[inline(always)]
639 move |i| {
640 if i % Self::C8 == 0 {
641 byte = this.seq[i / Self::C8];
642 }
643 (byte >> (B * (i % Self::C8))) & Self::CHAR_MASK as u8
645 },
646 )
647 .advance(this.offset)
648 }
649
650 #[inline(always)]
651 fn par_iter_bp(self, context: usize) -> PaddedIt<impl ChunkIt<S>> {
652 self.par_iter_bp_with_buf(context, RecycledBox::take())
653 }
654
655 #[inline(always)]
656 fn par_iter_bp_delayed(self, context: usize, delay: Delay) -> PaddedIt<impl ChunkIt<(S, S)>> {
657 self.par_iter_bp_delayed_with_factor(context, delay, 1)
658 }
659
660 #[inline(always)]
663 fn par_iter_bp_delayed_2(
664 self,
665 context: usize,
666 delay1: Delay,
667 delay2: Delay,
668 ) -> PaddedIt<impl ChunkIt<(S, S, S)>> {
669 self.par_iter_bp_delayed_2_with_factor_and_buf(
670 context,
671 delay1,
672 delay2,
673 1,
674 RecycledVec::take(),
675 )
676 }
677
678 fn cmp_lcp(&self, other: &Self) -> (std::cmp::Ordering, usize) {
680 let mut lcp = 0;
681 let min_len = self.len.min(other.len);
682 for i in (0..min_len).step_by(Self::K64) {
683 let len = (min_len - i).min(Self::K64);
684 let this = self.slice(i..i + len);
685 let other = other.slice(i..i + len);
686 let this_word = this.as_u64();
687 let other_word = other.as_u64();
688 if this_word != other_word {
689 let eq = this_word ^ other_word;
691 let t = eq.trailing_zeros() as usize / B * B;
692 lcp += t / B;
693 let mask = (Self::CHAR_MASK) << t;
694 return ((this_word & mask).cmp(&(other_word & mask)), lcp);
695 }
696 lcp += len;
697 }
698 (self.len.cmp(&other.len), lcp)
699 }
700
701 #[inline(always)]
702 fn get(&self, index: usize) -> u8 {
703 let offset = self.offset + index;
704 let idx = offset / Self::C8;
705 let offset = offset % Self::C8;
706 (self.seq[idx] >> (B * offset)) & Self::CHAR_MASK as u8
707 }
708}
709
710impl<'s, const B: usize> PackedSeqBase<'s, B>
711where
712 Bits<B>: SupportedBits,
713{
714 #[inline(always)]
715 pub fn par_iter_bp_with_buf<BUF: DerefMut<Target = [S; 8]>>(
716 self,
717 context: usize,
718 mut buf: BUF,
719 ) -> PaddedIt<impl ChunkIt<S> + use<'s, B, BUF>> {
720 #[cfg(target_endian = "big")]
721 panic!("Big endian architectures are not supported.");
722
723 let this = self.normalize();
724 let o = this.offset;
725 assert!(o < Self::C8);
726
727 let num_kmers = if this.len == 0 {
728 0
729 } else {
730 (this.len + o).saturating_sub(context - 1)
731 };
732 let num_kmers_stride = this.len.saturating_sub(context - 1);
734 let n = num_kmers_stride.div_ceil(L).next_multiple_of(Self::C8);
735 let bytes_per_chunk = n / Self::C8;
736 let padding = Self::C8 * L * bytes_per_chunk - num_kmers_stride;
737
738 let offsets: [usize; 8] = from_fn(|l| l * bytes_per_chunk);
739 let mut cur = S::ZERO;
740
741 let par_len = if num_kmers == 0 {
742 0
743 } else {
744 n + context + o - 1
745 };
746
747 let last_i = par_len.saturating_sub(1) / Self::C32 * Self::C32;
748 assert!(offsets[7] + (last_i / Self::C8) + 32 <= this.seq.len());
750
751 let it = (0..par_len)
752 .map(
753 #[inline(always)]
754 move |i| {
755 if i % Self::C32 == 0 {
756 if i % Self::C256 == 0 {
757 let data: [u32x8; 8] = from_fn(
759 #[inline(always)]
760 |lane| unsafe {
761 read_slice_32_unchecked(
762 this.seq,
763 offsets[lane] + (i / Self::C8),
764 )
765 },
766 );
767 *buf = transpose(data);
768 }
769 cur = buf[(i % Self::C256) / Self::C32];
770 }
771 let chars = cur & Self::SIMD_CHAR_MASK;
773 cur = cur >> Self::SIMD_B;
775 chars
776 },
777 )
778 .advance(o);
779
780 PaddedIt { it, padding }
781 }
782
783 #[inline(always)]
784 pub fn par_iter_bp_delayed_with_factor(
785 self,
786 context: usize,
787 delay: Delay,
788 factor: usize,
789 ) -> PaddedIt<impl ChunkIt<(S, S)> + use<'s, B>> {
790 self.par_iter_bp_delayed_with_factor_and_buf(context, delay, factor, RecycledVec::take())
791 }
792
793 #[inline(always)]
794 pub fn par_iter_bp_delayed_with_buf<BUF: DerefMut<Target = Vec<S>>>(
795 self,
796 context: usize,
797 delay: Delay,
798 buf: BUF,
799 ) -> PaddedIt<impl ChunkIt<(S, S)> + use<'s, B, BUF>> {
800 self.par_iter_bp_delayed_with_factor_and_buf(context, delay, 1, buf)
801 }
802
803 #[inline(always)]
804 pub fn par_iter_bp_delayed_with_factor_and_buf<BUF: DerefMut<Target = Vec<S>>>(
805 self,
806 context: usize,
807 Delay(delay): Delay,
808 factor: usize,
809 mut buf: BUF,
810 ) -> PaddedIt<impl ChunkIt<(S, S)> + use<'s, B, BUF>> {
811 #[cfg(target_endian = "big")]
812 panic!("Big endian architectures are not supported.");
813
814 assert!(
815 delay < usize::MAX / 2,
816 "Delay={} should be >=0.",
817 delay as isize
818 );
819
820 let this = self.normalize();
821 let o = this.offset;
822 assert!(o < Self::C8);
823
824 let num_kmers = if this.len == 0 {
825 0
826 } else {
827 (this.len + o).saturating_sub(context - 1)
828 };
829 let num_kmers_stride = this.len.saturating_sub(context - 1);
831 let n = num_kmers_stride
832 .div_ceil(L)
833 .next_multiple_of(factor * Self::C8);
834 let bytes_per_chunk = n / Self::C8;
835 let padding = Self::C8 * L * bytes_per_chunk - num_kmers_stride;
836
837 let offsets: [usize; 8] = from_fn(|l| l * bytes_per_chunk);
838 let mut upcoming = S::ZERO;
839 let mut upcoming_d = S::ZERO;
840
841 let buf_len = (delay / Self::C32 + 8).next_power_of_two();
846 let buf_mask = buf_len - 1;
847 if buf.capacity() < buf_len {
848 *buf = vec![S::ZERO; buf_len];
850 } else {
851 buf.clear();
853 buf.resize(buf_len, S::ZERO);
854 }
855
856 let mut write_idx = 0;
857 let mut read_idx = (buf_len - delay / Self::C32) % buf_len;
860
861 let par_len = if num_kmers == 0 {
862 0
863 } else {
864 n + context + o - 1
865 };
866
867 let last_i = par_len.saturating_sub(1) / Self::C32 * Self::C32;
868 assert!(offsets[7] + (last_i / Self::C8) + 32 <= this.seq.len());
870
871 let it = (0..par_len)
872 .map(
873 #[inline(always)]
874 move |i| {
875 if i % Self::C32 == 0 {
876 if i % Self::C256 == 0 {
877 let data: [u32x8; 8] = from_fn(
879 #[inline(always)]
880 |lane| unsafe {
881 read_slice_32_unchecked(
882 this.seq,
883 offsets[lane] + (i / Self::C8),
884 )
885 },
886 );
887 unsafe {
888 *TryInto::<&mut [u32x8; 8]>::try_into(
889 buf.get_unchecked_mut(write_idx..write_idx + 8),
890 )
891 .unwrap_unchecked() = transpose(data);
892 }
893 if i == 0 {
894 let elem = !((1u32 << (B * o)) - 1);
896 let mask = S::splat(elem);
897 unsafe { assert_unchecked(write_idx < buf.len()) };
898 buf[write_idx] &= mask;
899 }
900 }
901 unsafe { assert_unchecked(write_idx < buf.len()) };
902 upcoming = buf[write_idx];
903 write_idx += 1;
904 write_idx &= buf_mask;
905 }
906 if i % Self::C32 == delay % Self::C32 {
907 unsafe { assert_unchecked(read_idx < buf.len()) };
908 upcoming_d = buf[read_idx];
909 read_idx += 1;
910 read_idx &= buf_mask;
911 }
912 let chars = upcoming & Self::SIMD_CHAR_MASK;
914 let chars_d = upcoming_d & Self::SIMD_CHAR_MASK;
915 upcoming = upcoming >> Self::SIMD_B;
917 upcoming_d = upcoming_d >> Self::SIMD_B;
918 (chars, chars_d)
919 },
920 )
921 .advance(o);
922
923 PaddedIt { it, padding }
924 }
925
926 #[inline(always)]
927 pub fn par_iter_bp_delayed_2_with_factor(
928 self,
929 context: usize,
930 delay1: Delay,
931 delay2: Delay,
932 factor: usize,
933 ) -> PaddedIt<impl ChunkIt<(S, S, S)> + use<'s, B>> {
934 self.par_iter_bp_delayed_2_with_factor_and_buf(
935 context,
936 delay1,
937 delay2,
938 factor,
939 RecycledVec::take(),
940 )
941 }
942
943 #[inline(always)]
944 pub fn par_iter_bp_delayed_2_with_buf<BUF: DerefMut<Target = Vec<S>>>(
945 self,
946 context: usize,
947 delay1: Delay,
948 delay2: Delay,
949 buf: BUF,
950 ) -> PaddedIt<impl ChunkIt<(S, S, S)> + use<'s, B, BUF>> {
951 self.par_iter_bp_delayed_2_with_factor_and_buf(context, delay1, delay2, 1, buf)
952 }
953
954 #[inline(always)]
960 pub fn par_iter_bp_delayed_2_with_factor_and_buf<BUF: DerefMut<Target = Vec<S>>>(
961 self,
962 context: usize,
963 Delay(delay1): Delay,
964 Delay(delay2): Delay,
965 factor: usize,
966 mut buf: BUF,
967 ) -> PaddedIt<impl ChunkIt<(S, S, S)> + use<'s, B, BUF>> {
968 #[cfg(target_endian = "big")]
969 panic!("Big endian architectures are not supported.");
970
971 let this = self.normalize();
972 let o = this.offset;
973 assert!(o < Self::C8);
974 assert!(delay1 <= delay2, "Delay1 must be at most delay2.");
975
976 let num_kmers = if this.len == 0 {
977 0
978 } else {
979 (this.len + o).saturating_sub(context - 1)
980 };
981 let num_kmers_stride = this.len.saturating_sub(context - 1);
983 let n = num_kmers_stride
984 .div_ceil(L)
985 .next_multiple_of(factor * Self::C8);
986 let bytes_per_chunk = n / Self::C8;
987 let padding = Self::C8 * L * bytes_per_chunk - num_kmers_stride;
988
989 let offsets: [usize; 8] = from_fn(|l| l * bytes_per_chunk);
990 let mut upcoming = S::ZERO;
991 let mut upcoming_d1 = S::ZERO;
992 let mut upcoming_d2 = S::ZERO;
993
994 let buf_len = (delay2 / Self::C32 + 8).next_power_of_two();
996 let buf_mask = buf_len - 1;
997 if buf.capacity() < buf_len {
998 *buf = vec![S::ZERO; buf_len];
1000 } else {
1001 buf.clear();
1003 buf.resize(buf_len, S::ZERO);
1004 }
1005
1006 let mut write_idx = 0;
1007 let mut read_idx1 = (buf_len - delay1 / Self::C32) % buf_len;
1010 let mut read_idx2 = (buf_len - delay2 / Self::C32) % buf_len;
1011
1012 let par_len = if num_kmers == 0 {
1013 0
1014 } else {
1015 n + context + o - 1
1016 };
1017
1018 let last_i = par_len.saturating_sub(1) / Self::C32 * Self::C32;
1019 assert!(offsets[7] + (last_i / Self::C8) + 32 <= this.seq.len());
1021
1022 let it = (0..par_len)
1023 .map(
1024 #[inline(always)]
1025 move |i| {
1026 if i % Self::C32 == 0 {
1027 if i % Self::C256 == 0 {
1028 let data: [u32x8; 8] = from_fn(
1030 #[inline(always)]
1031 |lane| unsafe {
1032 read_slice_32_unchecked(
1033 this.seq,
1034 offsets[lane] + (i / Self::C8),
1035 )
1036 },
1037 );
1038 unsafe {
1039 *TryInto::<&mut [u32x8; 8]>::try_into(
1040 buf.get_unchecked_mut(write_idx..write_idx + 8),
1041 )
1042 .unwrap_unchecked() = transpose(data);
1043 }
1044 if i == 0 {
1046 let elem = !((1u32 << (B * o)) - 1);
1048 let mask = S::splat(elem);
1049 buf[write_idx] &= mask;
1050 }
1051 }
1052 upcoming = buf[write_idx];
1053 write_idx += 1;
1054 write_idx &= buf_mask;
1055 }
1056 if i % Self::C32 == delay1 % Self::C32 {
1057 unsafe { assert_unchecked(read_idx1 < buf.len()) };
1058 upcoming_d1 = buf[read_idx1];
1059 read_idx1 += 1;
1060 read_idx1 &= buf_mask;
1061 }
1062 if i % Self::C32 == delay2 % Self::C32 {
1063 unsafe { assert_unchecked(read_idx2 < buf.len()) };
1064 upcoming_d2 = buf[read_idx2];
1065 read_idx2 += 1;
1066 read_idx2 &= buf_mask;
1067 }
1068 let chars = upcoming & Self::SIMD_CHAR_MASK;
1070 let chars_d1 = upcoming_d1 & Self::SIMD_CHAR_MASK;
1071 let chars_d2 = upcoming_d2 & Self::SIMD_CHAR_MASK;
1072 upcoming = upcoming >> Self::SIMD_B;
1074 upcoming_d1 = upcoming_d1 >> Self::SIMD_B;
1075 upcoming_d2 = upcoming_d2 >> Self::SIMD_B;
1076 (chars, chars_d1, chars_d2)
1077 },
1078 )
1079 .advance(o);
1080
1081 PaddedIt { it, padding }
1082 }
1083}
1084
1085impl<const B: usize> PartialEq for PackedSeqBase<'_, B>
1086where
1087 Bits<B>: SupportedBits,
1088{
1089 fn eq(&self, other: &Self) -> bool {
1091 if self.len != other.len {
1092 return false;
1093 }
1094 for i in (0..self.len).step_by(Self::K64) {
1095 let len = (self.len - i).min(Self::K64);
1096 let this = self.slice(i..i + len);
1097 let that = other.slice(i..i + len);
1098 if this.as_u64() != that.as_u64() {
1099 return false;
1100 }
1101 }
1102 true
1103 }
1104}
1105
1106impl<const B: usize> Eq for PackedSeqBase<'_, B> where Bits<B>: SupportedBits {}
1107
1108impl<const B: usize> PartialOrd for PackedSeqBase<'_, B>
1109where
1110 Bits<B>: SupportedBits,
1111{
1112 fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
1113 Some(self.cmp(other))
1114 }
1115}
1116
1117impl<const B: usize> Ord for PackedSeqBase<'_, B>
1118where
1119 Bits<B>: SupportedBits,
1120{
1121 fn cmp(&self, other: &Self) -> std::cmp::Ordering {
1123 let min_len = self.len.min(other.len);
1124 for i in (0..min_len).step_by(Self::K64) {
1125 let len = (min_len - i).min(Self::K64);
1126 let this = self.slice(i..i + len);
1127 let other = other.slice(i..i + len);
1128 let this_word = this.as_u64();
1129 let other_word = other.as_u64();
1130 if this_word != other_word {
1131 let eq = this_word ^ other_word;
1133 let t = eq.trailing_zeros() as usize / B * B;
1134 let mask = (Self::CHAR_MASK) << t;
1135 return (this_word & mask).cmp(&(other_word & mask));
1136 }
1137 }
1138 self.len.cmp(&other.len)
1139 }
1140}
1141
1142impl<const B: usize> SeqVec for PackedSeqVecBase<B>
1143where
1144 Bits<B>: SupportedBits,
1145{
1146 type Seq<'s> = PackedSeqBase<'s, B>;
1147
1148 #[inline(always)]
1149 fn into_raw(mut self) -> Vec<u8> {
1150 self.seq.resize(self.len.div_ceil(Self::C8), 0);
1151 self.seq
1152 }
1153
1154 #[inline(always)]
1155 fn as_slice(&self) -> Self::Seq<'_> {
1156 PackedSeqBase {
1157 seq: &self.seq[..self.len.div_ceil(Self::C8) + PADDING],
1158 offset: 0,
1159 len: self.len,
1160 }
1161 }
1162
1163 #[inline(always)]
1164 fn len(&self) -> usize {
1165 self.len
1166 }
1167
1168 #[inline(always)]
1169 fn is_empty(&self) -> bool {
1170 self.len == 0
1171 }
1172
1173 #[inline(always)]
1174 fn clear(&mut self) {
1175 self.seq.clear();
1176 self.len = 0;
1177 }
1178
1179 fn push_seq<'a>(&mut self, seq: PackedSeqBase<'_, B>) -> Range<usize> {
1180 let start = self.len.next_multiple_of(Self::C8) + seq.offset;
1181 let end = start + seq.len();
1182
1183 self.seq.resize(self.len.div_ceil(Self::C8), 0);
1185 self.seq.extend(seq.seq);
1187
1188 if B == 1 {
1190 if self.len % 8 != 0 {
1191 self.seq[self.len / 8] |= u8::MAX << (self.len % 8);
1192 }
1193 self.seq[self.len.div_ceil(8)] |= (1u8 << seq.offset) - 1;
1194 }
1195 self.len = end;
1196 start..end
1197 }
1198
1199 fn push_ascii(&mut self, seq: &[u8]) -> Range<usize> {
1210 match B {
1211 1 | 2 => {}
1212 _ => panic!(
1213 "Can only use ASCII input for 2-bit DNA packing, or 1-bit ambiguous indicators."
1214 ),
1215 }
1216
1217 if seq.is_empty() {
1218 return self.len..self.len;
1219 }
1220
1221 self.seq
1222 .resize((self.len + seq.len()).div_ceil(Self::C8) + PADDING, 0);
1223
1224 let start_aligned = self.len.next_multiple_of(Self::C8);
1225 let start = self.len;
1226 let len = seq.len();
1227 let mut idx = self.len / Self::C8;
1228
1229 let parse_base = |base| match B {
1230 1 => char_is_ambiguous(base),
1231 2 => pack_char_lossy(base),
1232 _ => unreachable!(),
1233 };
1234
1235 let unaligned = core::cmp::min(start_aligned - start, len);
1236 if unaligned > 0 {
1237 let mut packed_byte = self.seq[idx];
1238 for &base in &seq[..unaligned] {
1239 packed_byte |= parse_base(base) << ((self.len % Self::C8) * B);
1240 self.len += 1;
1241 }
1242 self.seq[idx] = packed_byte;
1243 idx += 1;
1244 }
1245
1246 #[allow(unused)]
1247 let mut last = unaligned;
1248
1249 if B == 2 {
1250 #[cfg(all(target_arch = "x86_64", target_feature = "bmi2"))]
1251 {
1252 last = unaligned + (len - unaligned) / 8 * 8;
1253
1254 for i in (unaligned..last).step_by(8) {
1255 let chunk =
1256 unsafe { seq.get_unchecked(i..i + 8).try_into().unwrap_unchecked() };
1257 let ascii = u64::from_le_bytes(chunk);
1258 let packed_bytes =
1259 unsafe { std::arch::x86_64::_pext_u64(ascii, 0x0606060606060606) } as u16;
1260 unsafe {
1261 self.seq
1262 .get_unchecked_mut(idx..(idx + 2))
1263 .copy_from_slice(&packed_bytes.to_le_bytes())
1264 };
1265 idx += 2;
1266 self.len += 8;
1267 }
1268 }
1269
1270 #[cfg(target_feature = "neon")]
1271 {
1272 use core::arch::aarch64::{
1273 vandq_u8, vdup_n_u8, vld1q_u8, vpadd_u8, vshlq_u8, vst1_u8,
1274 };
1275 use core::mem::transmute;
1276
1277 last = unaligned + (len - unaligned) / 16 * 16;
1278
1279 for i in (unaligned..last).step_by(16) {
1280 unsafe {
1281 let ascii = vld1q_u8(seq.as_ptr().add(i));
1282 let masked_bits = vandq_u8(ascii, transmute([6i8; 16]));
1283 let (bits_0, bits_1) = transmute(vshlq_u8(
1284 masked_bits,
1285 transmute([-1i8, 1, 3, 5, -1, 1, 3, 5, -1, 1, 3, 5, -1, 1, 3, 5]),
1286 ));
1287 let half_packed = vpadd_u8(bits_0, bits_1);
1288 let packed = vpadd_u8(half_packed, vdup_n_u8(0));
1289 vst1_u8(self.seq.as_mut_ptr().add(idx), packed);
1290 idx += Self::C8;
1291 self.len += 16;
1292 }
1293 }
1294 }
1295
1296 #[cfg(all(target_arch = "wasm32", target_feature = "simd128"))]
1297 {
1298 use core::mem::transmute as t;
1299 use wide::{u8x16, u32x4};
1300
1301 last = unaligned + (len - unaligned) / 16 * 16;
1302
1303 let pick_low_bytes = u8x16::new([
1304 0, 4, 8, 12, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255,
1305 ]);
1306
1307 for i in (unaligned..last).step_by(16) {
1308 unsafe {
1309 let ascii = seq.as_ptr().add(i).cast::<u8x16>().read_unaligned();
1310 let bits: u32x4 = t((ascii >> 1) & u8x16::splat(3));
1311 let bits = bits | (bits >> 6);
1312 let bits = bits | (bits >> 12);
1313 let packed: [u32; 4] =
1314 t(t::<_, u8x16>(bits).shuffle_zeroing(pick_low_bytes));
1315 self.seq
1316 .get_unchecked_mut(idx..(idx + 4))
1317 .copy_from_slice(&packed[0].to_le_bytes());
1318 idx += 4;
1319 self.len += 16;
1320 }
1321 }
1322 }
1323 }
1324
1325 if B == 1 {
1326 #[cfg(all(target_arch = "x86_64", target_feature = "avx2"))]
1327 {
1328 last = len;
1329 self.len += len - unaligned;
1330
1331 let mut last_i = unaligned;
1332
1333 for i in (unaligned..last).step_by(32) {
1334 use std::mem::transmute as t;
1335
1336 type S = wide::i8x32;
1338 let chars: S = unsafe { t(read_slice_32(seq, i)) };
1339 let upper_mask = !(b'a' - b'A');
1340 let chars = chars & S::splat(upper_mask as i8);
1342 let lossy_encoded = chars & S::splat(6);
1343 let table = unsafe { S::from(t::<_, S>(*b"AxCxTxGxxxxxxxxxAxCxTxGxxxxxxxxx")) };
1344 let lookup: S = unsafe {
1345 t(std::arch::x86_64::_mm256_shuffle_epi8(
1346 t(table),
1347 t(lossy_encoded),
1348 ))
1349 };
1350 let packed_bytes = !(chars.simd_eq(lookup).to_bitmask() as u32);
1351
1352 last_i = i;
1353 unsafe {
1354 self.seq
1355 .get_unchecked_mut(idx..(idx + 4))
1356 .copy_from_slice(&packed_bytes.to_le_bytes())
1357 };
1358 idx += 4;
1359 }
1360
1361 if unaligned < last {
1363 idx -= 4;
1364 let mut val = unsafe {
1365 u32::from_le_bytes(
1366 self.seq
1367 .get_unchecked(idx..(idx + 4))
1368 .try_into()
1369 .unwrap_unchecked(),
1370 )
1371 };
1372 let keep = last - last_i;
1374 val <<= 32 - keep;
1375 val >>= 32 - keep;
1376 unsafe {
1377 self.seq
1378 .get_unchecked_mut(idx..(idx + 4))
1379 .copy_from_slice(&val.to_le_bytes())
1380 };
1381 idx += keep.div_ceil(8);
1382 }
1383 }
1384
1385 #[cfg(target_feature = "neon")]
1386 {
1387 use core::arch::aarch64::*;
1388 use core::mem::transmute;
1389
1390 last = unaligned + (len - unaligned) / 64 * 64;
1391
1392 for i in (unaligned..last).step_by(64) {
1393 unsafe {
1394 let ptr = seq.as_ptr().add(i);
1395 let chars = vld4q_u8(ptr);
1396
1397 let upper_mask = vdupq_n_u8(!(b'a' - b'A'));
1398 let chars = neon::map_8x16x4(chars, |v| vandq_u8(v, upper_mask));
1399
1400 let two_bits_mask = vdupq_n_u8(6);
1401 let lossy_encoded = neon::map_8x16x4(chars, |v| vandq_u8(v, two_bits_mask));
1402
1403 let table = transmute(*b"AxCxTxGxxxxxxxxx");
1404 let lookup = neon::map_8x16x4(lossy_encoded, |v| vqtbl1q_u8(table, v));
1405
1406 let mask = neon::map_two_8x16x4(chars, lookup, |v1, v2| vceqq_u8(v1, v2));
1407 let packed_bytes = !neon::movemask_64(mask);
1408
1409 self.seq[idx..(idx + 8)].copy_from_slice(&packed_bytes.to_le_bytes());
1410 idx += 8;
1411 self.len += 64;
1412 }
1413 }
1414 }
1415
1416 #[cfg(all(target_arch = "wasm32", target_feature = "simd128"))]
1417 {
1418 use wide::u8x16;
1419
1420 last = unaligned + (len - unaligned) / 16 * 16;
1421
1422 let upper_mask = u8x16::splat(!(b'a' - b'A'));
1423 let two_bits_mask = u8x16::splat(6);
1424 let table = u8x16::new(*b"AxCxTxGxxxxxxxxx");
1425
1426 for i in (unaligned..last).step_by(16) {
1427 unsafe {
1428 let chars = seq.as_ptr().add(i).cast::<u8x16>().read_unaligned();
1429 let chars = chars & upper_mask;
1430 let lossy_encoded = chars & two_bits_mask;
1431 let lookup = table.shuffle_zeroing(lossy_encoded);
1432 let packed_bytes = !(chars.simd_eq(lookup).to_bitmask() as u16);
1433
1434 self.seq
1435 .get_unchecked_mut(idx..(idx + 2))
1436 .copy_from_slice(&packed_bytes.to_le_bytes());
1437 idx += 2;
1438 self.len += 16;
1439 }
1440 }
1441 }
1442 }
1443
1444 let mut packed_byte = 0;
1445 for &base in &seq[last..] {
1446 packed_byte |= parse_base(base) << ((self.len % Self::C8) * B);
1447 self.len += 1;
1448 if self.len % Self::C8 == 0 {
1449 self.seq[idx] = packed_byte;
1450 idx += 1;
1451 packed_byte = 0;
1452 }
1453 }
1454 if self.len % Self::C8 != 0 && last < len {
1455 self.seq[idx] = packed_byte;
1456 idx += 1;
1457 }
1458 assert_eq!(idx + PADDING, self.seq.len());
1459 start..start + len
1460 }
1461
1462 #[cfg(feature = "rand")]
1463 fn random(n: usize) -> Self {
1464 use rand::Rng;
1465
1466 let byte_len = n.div_ceil(Self::C8);
1467 let mut seq = vec![0; byte_len + PADDING];
1468 rand::make_rng::<rand::rngs::SmallRng>().fill_bytes(&mut seq[..byte_len]);
1469 if n % Self::C8 != 0 {
1471 seq[byte_len - 1] &= (1 << (B * (n % Self::C8))) - 1;
1472 }
1473
1474 Self { seq, len: n }
1475 }
1476}
1477
1478impl PackedSeqVecBase<2> {
1479 pub fn from_raw_parts(mut seq: Vec<u8>, len: usize) -> Self {
1484 assert!(len <= seq.len() * Self::C8);
1485 seq.resize(len.div_ceil(Self::C8) + PADDING, 0);
1486 Self { seq, len }
1487 }
1488
1489 #[cfg(feature = "file_io")]
1494 pub fn from_fastx(path: &std::path::Path) -> (Self, Vec<Range<usize>>) {
1495 use helicase::{HelicaseParser, input::FromFile};
1496
1497 let mut ranges = vec![];
1498
1499 const CONFIG: helicase::Config = helicase::ParserOptions::default()
1500 .ignore_headers()
1501 .ignore_quality()
1502 .dna_packed()
1503 .reject_non_actg()
1504 .config()
1505 | helicase::config::advanced::MERGE_RECORDS;
1506 let mut parser = helicase::FastxParser::<CONFIG>::from_file(&path).unwrap();
1507
1508 let mut start = 0;
1509 while let Some(event) = parser.next() {
1510 event.unwrap();
1511 let end = parser.get_dna_packed().len();
1512 ranges.push(start..end);
1513 start = end;
1514 }
1515 let (len, seq) = parser.get_dna_packed_owned().into_packed_bytes();
1516
1517 let seq = PackedSeqVec::from_raw_parts(seq, len);
1518 (seq, ranges)
1519 }
1520
1521 #[cfg(feature = "file_io")]
1529 pub fn from_fastx_parallel(
1530 path: &std::path::Path,
1531 num_threads: usize,
1532 ) -> (Self, Vec<Range<usize>>) {
1533 use helicase::HelicaseParser;
1534
1535 const CONFIG: helicase::Config = helicase::ParserOptions::default()
1536 .ignore_headers()
1537 .ignore_quality()
1538 .dna_packed()
1539 .reject_non_actg()
1540 .config()
1541 | helicase::config::advanced::MERGE_RECORDS;
1542 let chunks = std::sync::Mutex::new(Vec::new());
1543 helicase::ParallelFastxParser::<CONFIG>::from_file(path)
1545 .unwrap()
1546 .threads(num_threads)
1547 .for_each(|mut parser| {
1548 let mut ranges = Vec::new();
1549 let mut start = 0;
1550 while let Some(event) = parser.next() {
1551 event?;
1552 let end = parser.get_dna_packed().len();
1553 ranges.push(start..end);
1554 start = end;
1555 }
1556 let (len, mut seq) = parser.get_dna_packed_owned().into_packed_bytes();
1557 seq.truncate(len.div_ceil(Self::C8));
1558 chunks.lock().unwrap().push((seq, ranges));
1559 Ok(())
1560 })
1561 .unwrap();
1562
1563 let chunks = chunks.into_inner().unwrap();
1565 let total_bytes = chunks.iter().map(|(seq, _)| seq.len()).sum::<usize>();
1566 let total_records = chunks.iter().map(|(_, ranges)| ranges.len()).sum::<usize>();
1567
1568 let mut seq = vec![0; total_bytes + PADDING];
1570 let mut ranges = Vec::<Range<usize>>::with_capacity(total_records);
1571
1572 std::thread::scope(|scope| {
1574 let mut seq_out = &mut seq[..total_bytes];
1575 let mut ranges_out = &mut ranges.spare_capacity_mut()[..total_records];
1576 let mut offset = 0;
1577 let threads = num_threads.max(1);
1578 let groups = chunks.chunks(chunks.len().div_ceil(threads));
1579 for group in groups {
1580 let thread_bytes = group.iter().map(|(seq, _)| seq.len()).sum();
1581 let thread_records = group.iter().map(|(_, ranges)| ranges.len()).sum();
1582
1583 let (mut dest, rest) = seq_out.split_at_mut(thread_bytes);
1585 seq_out = rest;
1586 let (mut range_dest, rest) = ranges_out.split_at_mut(thread_records);
1587 ranges_out = rest;
1588
1589 let mut base_offset = offset * Self::C8;
1590 offset += thread_bytes;
1591
1592 scope.spawn(move || {
1593 for (chunk, chunk_ranges) in group {
1594 let (chunk_dest, rest) = dest.split_at_mut(chunk.len());
1596 dest = rest;
1597 let (chunk_range_dest, rest) = range_dest.split_at_mut(chunk_ranges.len());
1598 range_dest = rest;
1599
1600 chunk_dest.copy_from_slice(&chunk);
1602
1603 let chunk_offset = base_offset;
1605 base_offset += chunk.len() * Self::C8;
1606 for (dest, range) in chunk_range_dest.iter_mut().zip(chunk_ranges) {
1607 dest.write(chunk_offset + range.start..chunk_offset + range.end);
1608 }
1609 }
1610 });
1611 }
1612 });
1613 unsafe { ranges.set_len(total_records) };
1616 (Self::from_raw_parts(seq, total_bytes * Self::C8), ranges)
1617 }
1618}
1619
1620impl PackedSeqVecBase<1> {
1621 pub fn with_len(n: usize) -> Self {
1622 Self {
1623 seq: vec![0; n.div_ceil(Self::C8) + PADDING],
1624 len: n,
1625 }
1626 }
1627
1628 #[cfg(feature = "rand")]
1629 pub fn random(len: usize, n_frac: f32) -> Self {
1630 let byte_len = len.div_ceil(Self::C8);
1631 let mut seq = vec![0; byte_len + PADDING];
1632
1633 assert!(
1634 (0.0..=0.3).contains(&n_frac),
1635 "n_frac={} should be in [0, 0.3]",
1636 n_frac
1637 );
1638
1639 for _ in 0..(len as f32 * n_frac) as usize {
1640 let idx = rand::random::<u64>() as usize % len;
1641 let byte = idx / Self::C8;
1642 let offset = idx % Self::C8;
1643 seq[byte] |= 1 << offset;
1644 }
1645
1646 Self { seq, len }
1647 }
1648}
1649
1650impl<'s> PackedSeqBase<'s, 1> {
1651 #[inline(always)]
1655 pub fn iter_kmer_ambiguity(self, k: usize) -> impl ExactSizeIterator<Item = bool> + use<'s> {
1656 let this = self.normalize();
1657 assert!(k > 0);
1658 assert!(k <= Self::K64);
1659 (this.offset..this.offset + this.len.saturating_sub(k - 1))
1660 .map(move |i| self.read_kmer(k, i) != 0)
1661 }
1662
1663 #[inline(always)]
1671 pub fn par_iter_kmer_ambiguity(
1672 self,
1673 k: usize,
1674 context: usize,
1675 skip: usize,
1676 ) -> PaddedIt<impl ChunkIt<S> + use<'s>> {
1677 #[cfg(target_endian = "big")]
1678 panic!("Big endian architectures are not supported.");
1679
1680 assert!(k > 0, "par_iter_kmers requires k>0, but k={k}");
1681 assert!(k <= 96, "par_iter_kmers requires k<=96, but k={k}");
1682
1683 let this = self.normalize();
1684 let o = this.offset;
1685 assert!(o < Self::C8);
1686
1687 let delay = k - 1;
1688
1689 let it = self.par_iter_bp_delayed(context, Delay(delay));
1690
1691 let mut cnt = u32x8::ZERO;
1692
1693 it.map(
1694 #[inline(always)]
1695 move |(a, r)| {
1696 cnt += a;
1697 let out = cnt.simd_gt(S::ZERO);
1698 cnt -= r;
1699 out
1700 },
1701 )
1702 .advance(skip)
1703 }
1704
1705 #[inline(always)]
1706 pub fn par_iter_kmer_ambiguity_with_buf(
1707 self,
1708 k: usize,
1709 context: usize,
1710 skip: usize,
1711 buf: &'s mut Vec<S>,
1712 ) -> PaddedIt<impl ChunkIt<S> + use<'s>> {
1713 #[cfg(target_endian = "big")]
1714 panic!("Big endian architectures are not supported.");
1715
1716 assert!(k > 0, "par_iter_kmers requires k>0, but k={k}");
1717 assert!(k <= 96, "par_iter_kmers requires k<=96, but k={k}");
1718
1719 let this = self.normalize();
1720 let o = this.offset;
1721 assert!(o < Self::C8);
1722
1723 let delay = k - 1;
1724
1725 let it = self.par_iter_bp_delayed_with_buf(context, Delay(delay), buf);
1726
1727 let mut cnt = u32x8::ZERO;
1728
1729 it.map(
1730 #[inline(always)]
1731 move |(a, r)| {
1732 cnt += a;
1733 let out = cnt.simd_gt(S::ZERO);
1734 cnt -= r;
1735 out
1736 },
1737 )
1738 .advance(skip)
1739 }
1740}
1741
1742#[cfg(target_feature = "neon")]
1743mod neon {
1744 use core::arch::aarch64::*;
1745
1746 #[inline(always)]
1747 pub fn movemask_64(v: uint8x16x4_t) -> u64 {
1748 unsafe {
1750 let acc = vsriq_n_u8(vsriq_n_u8(v.3, v.2, 1), vsriq_n_u8(v.1, v.0, 1), 2);
1751 vget_lane_u64(
1752 vreinterpret_u64_u8(vshrn_n_u16(
1753 vreinterpretq_u16_u8(vsriq_n_u8(acc, acc, 4)),
1754 4,
1755 )),
1756 0,
1757 )
1758 }
1759 }
1760
1761 #[inline(always)]
1762 pub fn map_8x16x4<F>(v: uint8x16x4_t, mut f: F) -> uint8x16x4_t
1763 where
1764 F: FnMut(uint8x16_t) -> uint8x16_t,
1765 {
1766 uint8x16x4_t(f(v.0), f(v.1), f(v.2), f(v.3))
1767 }
1768
1769 #[inline(always)]
1770 pub fn map_two_8x16x4<F>(v1: uint8x16x4_t, v2: uint8x16x4_t, mut f: F) -> uint8x16x4_t
1771 where
1772 F: FnMut(uint8x16_t, uint8x16_t) -> uint8x16_t,
1773 {
1774 uint8x16x4_t(f(v1.0, v2.0), f(v1.1, v2.1), f(v1.2, v2.2), f(v1.3, v2.3))
1775 }
1776}