Skip to main content

debruijn/
dna_string.rs

1// Copyright 2014 Johannes Köster and 10x Genomics
2// Licensed under the MIT license (http://opensource.org/licenses/MIT)
3// This file may not be copied, modified, or distributed
4// except according to those terms.
5
6//! A 2-bit encoding of arbitrary length DNA sequences.
7//!
8//! Store arbitrary-length DNA strings in a packed 2-bit encoding. Individual base values are encoded
9//! as the integers 0,1,2,3 corresponding to A,C,G,T.
10//!
11//! # Example
12//! ```
13//! use debruijn::Kmer;
14//! use debruijn::dna_string::*;
15//! use debruijn::kmer::Kmer16;
16//! use debruijn::Vmer;
17//!
18//! // Construct a new DNA string
19//! let dna_string1 = DnaString::from_dna_string("ACAGCAGCAGCACGTATGACAGATAGTGACAGCAGTTTGTGACCGCAAGAGCAGTAATATGATG");
20//!
21//! // Get an immutable view into the sequence
22//! let slice1 = dna_string1.slice(10, 40);
23//!
24//! // Get a kmer from the DNA string
25//! let first_kmer: Kmer16 = slice1.get_kmer(0);
26//! assert_eq!(first_kmer, Kmer16::from_ascii(b"CACGTATGACAGATAG"))
27
28use itertools::Itertools;
29use serde_derive::{Deserialize, Serialize};
30use std::borrow::Borrow;
31use std::cmp::min;
32use std::collections::hash_map::DefaultHasher;
33use std::error::Error;
34use std::fmt::{self, Display};
35use std::hash::{Hash, Hasher};
36
37use crate::{base_to_bits, base_to_bits_checked};
38use crate::bits_to_ascii;
39use crate::bits_to_base;
40use crate::dna_only_base_to_bits;
41
42use crate::Kmer;
43use crate::Mer;
44use crate::MerIter;
45use crate::Vmer;
46
47const BLOCK_BITS: usize = 64;
48const WIDTH: usize = 2;
49
50const MASK: u64 = 0x3;
51
52/// A container for sequence of DNA bases.
53/// ```
54/// use debruijn::dna_string::DnaString;
55/// use debruijn::kmer::Kmer8;
56/// use debruijn::{Mer, Vmer};
57///
58/// let dna_string = DnaString::from_dna_string("ATCGTACGTACGTAGTC");
59///
60/// // Iterate over 8-mers
61/// for k in dna_string.iter_kmers::<Kmer8>() {
62///     println!("{:?}", k);
63/// }
64///
65/// // Get a base, encoded as a byte in 0-3 range
66/// assert_eq!(dna_string.get(0), 0);
67/// assert_eq!(dna_string.get(1), 3);
68///
69/// // Make a read-only 'slice' of a DnaString
70/// let slc = dna_string.slice(1, 10);
71///
72///  assert_eq!(slc.iter_kmers::<Kmer8>().next(), dna_string.iter_kmers::<Kmer8>().skip(1).next());
73/// ```
74#[derive(Ord, PartialOrd, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
75pub struct DnaString {
76    storage: Vec<u64>,
77    len: usize,
78}
79
80impl Mer for DnaString {
81    fn len(&self) -> usize {
82        self.len
83    }
84
85    fn is_empty(&self) -> bool {
86        self.len == 0
87    }
88
89    /// Get the value at position `i`.
90    #[inline(always)]
91    fn get(&self, i: usize) -> u8 {
92        let (block, bit) = self.addr(i);
93        self.get_by_addr(block, bit)
94    }
95
96    /// Set the value as position `i`.
97    fn set_mut(&mut self, i: usize, value: u8) {
98        let (block, bit) = self.addr(i);
99        self.set_by_addr(block, bit, value);
100    }
101
102    fn set_slice_mut(&mut self, _: usize, _: usize, _: u64) {
103        unimplemented!()
104    }
105
106    fn rc(&self) -> DnaString {
107        let mut dna_string = DnaString::new();
108        let rc = (0..self.len()).rev().map(|i| 3 - self.get(i));
109
110        dna_string.extend(rc);
111        dna_string
112    }
113}
114
115impl Vmer for DnaString {
116    fn new(len: usize) -> Self {
117        Self::blank(len)
118    }
119
120    fn max_len() -> usize {
121        <usize>::MAX
122    }
123
124    /// Get the kmer starting at position pos
125    fn get_kmer<K: Kmer>(&self, pos: usize) -> K {
126        assert!(self.len() - pos >= K::k());
127
128        // Which block has the first base
129        let (mut block, _) = self.addr(pos);
130
131        // Where we are in the kmer
132        let mut kmer_pos = 0;
133
134        // Where are in the block
135        let mut block_pos = pos % 32;
136
137        let mut kmer = K::empty();
138
139        while kmer_pos < K::k() {
140            // get relevent bases for current block
141            let nb = min(K::k() - kmer_pos, 32 - block_pos);
142
143            let v = self.storage[block];
144            let val = v << (2 * block_pos);
145            kmer.set_slice_mut(kmer_pos, nb, val);
146
147            // move to next block, move ahead in kmer.
148            block += 1;
149            kmer_pos += nb;
150            // alway start a beginning of next block
151            block_pos = 0;
152        }
153
154        kmer
155    }
156}
157
158impl DnaString {
159    /// Create an empty DNA string
160    pub fn new() -> DnaString {
161        DnaString {
162            storage: Vec::new(),
163            len: 0,
164        }
165    }
166
167    /// Length of the sequence
168    pub fn len(&self) -> usize {
169        self.len
170    }
171
172    /// Create a new instance with a given capacity.
173    pub fn with_capacity(n: usize) -> Self {
174        let blocks = ((n * WIDTH) >> 6) + (if (n * WIDTH) & 0x3F > 0 { 1 } else { 0 });
175        let storage = Vec::with_capacity(blocks);
176
177        DnaString { storage, len: 0 }
178    }
179
180    /// Create a DnaString of length n initialized to all A's
181    pub fn blank(n: usize) -> Self {
182        let blocks = ((n * WIDTH) >> 6) + (if (n * WIDTH) & 0x3F > 0 { 1 } else { 0 });
183        let storage = vec![0; blocks];
184
185        DnaString { storage, len: n }
186    }
187
188    /// Create a DnaString corresponding to an ACGT-encoded str.
189    pub fn from_dna_string(dna: &str) -> DnaString {
190        let mut dna_string = DnaString {
191            storage: Vec::new(),
192            len: 0,
193        };
194
195        dna_string.extend(dna.chars().map(|c| base_to_bits(c as u8)));
196        dna_string
197    }
198
199    /// Create a DnaString corresponding to an ACGT-encoded str.
200    pub fn from_dna_only_string(dna: &str) -> Vec<DnaString> {
201        let mut dna_vector: Vec<DnaString> = Vec::new();
202        let mut dna_string = DnaString::new();
203
204        for c in dna.chars() {
205            match dna_only_base_to_bits(c as u8) {
206                Some(bit) => {
207                    dna_string.push(bit);
208                }
209                None => {
210                    if !dna_string.is_empty() {
211                        dna_vector.push(dna_string);
212                        dna_string = DnaString::new();
213                    }
214                }
215            }
216        }
217        if !dna_string.is_empty() {
218            dna_vector.push(dna_string);
219        }
220
221        dna_vector
222    }
223
224    /// Create a DnaString from an ASCII ACGT-encoded byte slice.
225    /// Non ACGT positions will be converted to 'A'
226    pub fn from_acgt_bytes(bytes: &[u8]) -> DnaString {
227        let mut dna_string = DnaString::with_capacity(bytes.len());
228
229        // Accelerated avx2 mode. Should run on most machines made since 2013.
230        #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
231        {
232            if is_x86_feature_detected!("avx2") {
233                for chunk in bytes.chunks(32) {
234                    if chunk.len() == 32 {
235                        let (conv_chunk, _) = unsafe { crate::bitops_avx2::convert_bases(chunk) };
236                        let packed = unsafe { crate::bitops_avx2::pack_32_bases(conv_chunk) };
237                        dna_string.storage.push(packed);
238                    } else {
239                        let b = chunk.iter().map(|c| base_to_bits(*c));
240                        dna_string.extend(b);
241                    }
242                }
243
244                dna_string.len = bytes.len();
245                return dna_string;
246            }
247        }
248
249        let b = bytes.iter().map(|c| base_to_bits(*c));
250        dna_string.extend(b);
251        dna_string
252    }
253
254    /// Create a DnaString from an ASCII ACGT-encoded byte slice.
255    /// Will return `None` if there are ambiguous bases in the DnaString
256    pub fn from_acgt_bytes_checked(bytes: &[u8]) -> Result<DnaString, AmbiguousBasesError> {
257        let mut dna_string = DnaString::with_capacity(bytes.len());
258
259        // Accelerated avx2 mode. Should run on most machines made since 2013.
260        #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
261        {
262            if is_x86_feature_detected!("avx2") {
263                for chunk in bytes.chunks(32) {
264                    if chunk.len() == 32 {
265                        let (conv_chunk, correct) = unsafe { crate::bitops_avx2::convert_bases(chunk) };
266                        if !correct { return Err(AmbiguousBasesError {  }) }
267                        let packed = unsafe { crate::bitops_avx2::pack_32_bases(conv_chunk) };
268                        dna_string.storage.push(packed);
269                    } else {
270                        let (b, corrects): (Vec<u8>, Vec<bool>) = chunk.iter().map(|c| base_to_bits_checked(*c)).collect();
271                        let correct = !corrects.iter().contains(&false);
272                        if !correct { return Err(AmbiguousBasesError {  }) }
273                        dna_string.extend(b.into_iter());
274                    }
275                }
276
277                dna_string.len = bytes.len();
278                return Ok(dna_string);
279            }
280        }
281        
282        let (b, corrects): (Vec<u8>, Vec<bool>) = bytes.iter().map(|c| base_to_bits_checked(*c)).collect();
283        let correct = corrects.iter().contains(&false);
284        if !correct { return Err(AmbiguousBasesError {  }) }
285        dna_string.extend(b.into_iter());
286
287        Ok(dna_string)
288    }
289
290    /// Create a DnaString from an ACGT-encoded byte slice,
291    /// Non ACGT positions will be converted to repeatable random base determined
292    /// by a hash of the read name and the position within the string.
293    pub fn from_acgt_bytes_hashn(bytes: &[u8], read_name: &[u8]) -> DnaString {
294        let mut hasher = DefaultHasher::new();
295        read_name.hash(&mut hasher);
296
297        let mut dna_string = DnaString::with_capacity(bytes.len());
298
299        for (pos, c) in bytes.iter().enumerate() {
300            let v = match c {
301                b'A' | b'a' => 0u8,
302                b'C' | b'c' => 1u8,
303                b'G' | b'g' => 2u8,
304                b'T' | b't' => 3u8,
305                _ => {
306                    let mut hasher_clone = hasher.clone();
307                    pos.hash(&mut hasher_clone);
308                    (hasher_clone.finish() % 4) as u8
309                }
310            };
311
312            dna_string.push(v);
313        }
314
315        dna_string
316    }
317
318    /// Create a DnaString from a 0-4 encoded byte slice
319    pub fn from_bytes(bytes: &[u8]) -> DnaString {
320        let mut dna_string = DnaString {
321            storage: Vec::new(),
322            len: 0,
323        };
324
325        dna_string.extend(bytes.iter().cloned());
326        dna_string
327    }
328
329    /// Convert sequence to a Vector of 0-4 encoded bytes
330    pub fn to_bytes(&self) -> Vec<u8> {
331        self.iter().collect()
332    }
333
334    /// Convert sequence to a Vector of ascii-encoded bytes
335    pub fn to_ascii_vec(&self) -> Vec<u8> {
336        self.iter().map(bits_to_ascii).collect()
337    }
338
339    /// Append a 0-4 encoded base.
340    #[inline]
341    pub fn push(&mut self, value: u8) {
342        let (block, bit) = self.addr(self.len);
343        if bit == 0 && block >= self.storage.len() {
344            self.storage.push(0);
345        }
346        self.set_by_addr(block, bit, value);
347        self.len += 1;
348    }
349
350    pub fn extend(&mut self, mut bytes: impl Iterator<Item = u8>) {
351        // fill the last incomplete u64 block
352        while self.len % 32 != 0 {
353            match bytes.next() {
354                Some(b) => self.push(b),
355                None => return,
356            }
357        }
358
359        let mut bytes = bytes.peekable();
360
361        // chunk the remaining items into groups of at most 32 and handle them together
362        while bytes.peek().is_some() {
363            let mut val: u64 = 0;
364            let mut offset = 62;
365            let mut n_added = 0;
366
367            for _ in 0..32 {
368                if let Some(b) = bytes.next() {
369                    assert!(b < 4);
370                    val |= (b as u64) << offset;
371                    offset -= 2;
372                    n_added += 1;
373                } else {
374                    break;
375                }
376            }
377
378            self.storage.push(val);
379            self.len += n_added;
380        }
381    }
382
383    /// Push 0-4 encoded bases from a byte array.
384    ///
385    /// # Arguments
386    /// `bytes`: byte array to read values from
387    /// `seq_length`: how many values to read from the byte array. Note that this
388    /// is number of values not number of elements of the byte array.
389    pub fn push_bytes(&mut self, bytes: &[u8], seq_length: usize) {
390        assert!(
391            seq_length <= bytes.len() * 8 / WIDTH,
392            "Number of elements to push exceeds array length"
393        );
394
395        for i in 0..seq_length {
396            let byte_index = (i * WIDTH) / 8;
397            let byte_slot = (i * WIDTH) % 8;
398
399            let v = bytes[byte_index];
400            let bits = (v >> byte_slot) & (MASK as u8);
401
402            self.push(bits);
403        }
404    }
405
406    /// Iterate over stored values (values will be unpacked into bytes).
407    pub fn iter(&self) -> DnaStringIter<'_> {
408        DnaStringIter {
409            dna_string: self,
410            i: 0,
411        }
412    }
413
414    /// Clear the sequence.
415    pub fn clear(&mut self) {
416        self.storage.clear();
417        self.len = 0;
418    }
419
420    #[inline(always)]
421    fn get_by_addr(&self, block: usize, bit: usize) -> u8 {
422        ((self.storage[block] >> (62 - bit)) & MASK) as u8
423    }
424
425    #[inline(always)]
426    fn set_by_addr(&mut self, block: usize, bit: usize, value: u8) {
427        let mask = MASK << (62 - bit);
428        self.storage[block] |= mask;
429        self.storage[block] ^= mask;
430        self.storage[block] |= (value as u64 & MASK) << (62 - bit);
431    }
432
433    #[inline(always)]
434    fn addr(&self, i: usize) -> (usize, usize) {
435        let k = i * WIDTH;
436        (k / BLOCK_BITS, k % BLOCK_BITS)
437    }
438
439    pub fn is_empty(&self) -> bool {
440        self.len == 0
441    }
442
443    /// Get the length `k` prefix of the DnaString
444    pub fn prefix(&self, k: usize) -> DnaStringSlice<'_> {
445        assert!(k <= self.len, "Prefix size exceeds number of elements.");
446        DnaStringSlice {
447            dna_string: self,
448            start: 0,
449            length: k,
450            is_rc: false,
451        }
452    }
453
454    /// Get the length `k` suffix of the DnaString
455    pub fn suffix(&self, k: usize) -> DnaStringSlice<'_> {
456        assert!(k <= self.len, "Suffix size exceeds number of elements.");
457
458        DnaStringSlice {
459            dna_string: self,
460            start: self.len() - k,
461            length: k,
462            is_rc: false,
463        }
464    }
465
466    /// Get slice containing the interval [`start`, `end`) of `self`
467    pub fn slice(&self, start: usize, end: usize) -> DnaStringSlice<'_> {
468        assert!(start <= self.len, "coordinate exceeds number of elements.");
469        assert!(end <= self.len, "coordinate exceeds number of elements.");
470
471        DnaStringSlice {
472            dna_string: self,
473            start,
474            length: end - start,
475            is_rc: false,
476        }
477    }
478
479    /// Create a fresh DnaString containing the reverse of `self`
480    pub fn reverse(&self) -> DnaString {
481        let values: Vec<u8> = self.iter().collect();
482        let mut dna_string = DnaString::new();
483        for v in values.iter().rev() {
484            dna_string.push(*v);
485        }
486        dna_string
487    }
488
489    /// Compute Hamming distance between this DnaString and another DnaString. The two strings must have the same length.
490    pub fn hamming_distance(&self, other: &DnaString) -> usize {
491        ndiffs(self, other)
492    }
493
494    // pub fn complement(&self) -> DnaString {
495    //    assert!(self.width == 2, "Complement only supported for 2bit encodings.");
496    //    let values: Vec<u32> = Vec::with_capacity(self.len());
497    //    for i, v in self.storage.iter() {
498    //        values[i] = v;
499    //    }
500    //    values[values.len() - 1] =
501    // }
502
503    /// shrink the storage of the `DnaString` to fit its contents
504    pub fn shrink_to_fit(&mut self) {
505        self.storage.shrink_to_fit();
506    }
507}
508
509impl fmt::Display for DnaString {
510    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
511        for v in self.iter() {
512            write!(f, "{}", bits_to_base(v))?;
513        }
514        Ok(())
515    }
516}
517
518impl fmt::Debug for DnaString {
519    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
520        let mut s = String::new();
521        for pos in 0..self.len() {
522            s.push(bits_to_base(self.get(pos)))
523        }
524
525        write!(f, "{}", s)
526    }
527}
528
529impl Default for DnaString {
530    fn default() -> Self {
531        Self::new()
532    }
533}
534
535/// Iterator over values of a DnaStringoded sequence (values will be unpacked into bytes).
536pub struct DnaStringIter<'a> {
537    dna_string: &'a DnaString,
538    i: usize,
539}
540
541impl Iterator for DnaStringIter<'_> {
542    type Item = u8;
543
544    fn next(&mut self) -> Option<u8> {
545        if self.i < self.dna_string.len() {
546            let value = self.dna_string.get(self.i);
547            self.i += 1;
548            Some(value)
549        } else {
550            None
551        }
552    }
553}
554
555impl<'a> IntoIterator for &'a DnaString {
556    type Item = u8;
557    type IntoIter = DnaStringIter<'a>;
558
559    fn into_iter(self) -> DnaStringIter<'a> {
560        self.iter()
561    }
562}
563
564/// count Hamming distance between 2 2-bit DNA packed u64s
565#[inline]
566fn count_diff_2_bit_packed(a: u64, b: u64) -> u32 {
567    let bit_diffs = a ^ b;
568    let two_bit_diffs = (bit_diffs | bit_diffs >> 1) & 0x5555555555555555;
569    two_bit_diffs.count_ones()
570}
571
572/// Compute the number of base positions at which two DnaStrings differ, assuming
573/// that they have the same length.
574pub fn ndiffs(b1: &DnaString, b2: &DnaString) -> usize {
575    assert_eq!(b1.len(), b2.len());
576    let mut diffs = 0;
577    let (s1, s2) = (&b1.storage, &b2.storage);
578    for i in 0..s1.len() {
579        diffs += count_diff_2_bit_packed(s1[i], s2[i])
580    }
581    diffs as usize
582}
583
584/// An immutable slice into a DnaString
585#[derive(Clone)]
586pub struct DnaStringSlice<'a> {
587    pub dna_string: &'a DnaString,
588    pub start: usize,
589    pub length: usize,
590    pub is_rc: bool,
591}
592
593impl PartialEq for DnaStringSlice<'_> {
594    fn eq(&self, other: &DnaStringSlice) -> bool {
595        if other.length != self.length {
596            return false;
597        }
598        for i in 0..self.length {
599            if self.get(i) != other.get(i) {
600                return false;
601            }
602        }
603        true
604    }
605}
606impl Eq for DnaStringSlice<'_> {}
607
608impl<'a> Mer for DnaStringSlice<'a> {
609    #[inline(always)]
610    fn len(&self) -> usize {
611        self.length
612    }
613
614    fn is_empty(&self) -> bool {
615        self.length == 0
616    }
617
618    /// Get the base at position `i`.
619    #[inline(always)]
620    fn get(&self, i: usize) -> u8 {
621        if !self.is_rc {
622            self.dna_string.get(i + self.start)
623        } else {
624            crate::complement(self.dna_string.get(self.start + self.length - 1 - i))
625        }
626    }
627
628    /// Set the base as position `i`.
629    fn set_mut(&mut self, _: usize, _: u8) {
630        unimplemented!()
631        //debug_assert!(i < self.length);
632        //self.dna_string.set(i + self.start, value);
633    }
634
635    fn set_slice_mut(&mut self, _: usize, _: usize, _: u64) {
636        unimplemented!();
637    }
638
639    fn rc(&self) -> DnaStringSlice<'a> {
640        DnaStringSlice {
641            dna_string: self.dna_string,
642            start: self.start,
643            length: self.length,
644            is_rc: !self.is_rc,
645        }
646    }
647}
648
649impl Vmer for DnaStringSlice<'_> {
650    fn new(_: usize) -> Self {
651        unimplemented!()
652    }
653
654    fn max_len() -> usize {
655        <usize>::MAX
656    }
657
658    /// Get the kmer starting at position pos
659    fn get_kmer<K: Kmer>(&self, pos: usize) -> K {
660        debug_assert!(pos + K::k() <= self.length);
661        if !self.is_rc {
662            self.dna_string.get_kmer(self.start + pos)
663        } else {
664            let k = self
665                .dna_string
666                .get_kmer(self.start + self.length - K::k() - pos);
667            K::rc(&k)
668        }
669    }
670}
671
672impl DnaStringSlice<'_> {
673    pub fn is_palindrome(&self) -> bool {
674        unimplemented!();
675    }
676
677    pub fn bytes(&self) -> Vec<u8> {
678        let mut v = Vec::with_capacity(self.length);
679        for pos in 0..self.length {
680            v.push(self.get(pos));
681        }
682        v
683    }
684
685    pub fn ascii(&self) -> Vec<u8> {
686        let mut v = Vec::with_capacity(self.length);
687        for pos in 0..self.length {
688            v.push(bits_to_ascii(self.get(pos)));
689        }
690        v
691    }
692
693    pub fn to_dna_string(&self) -> String {
694        let mut dna: String = String::with_capacity(self.length);
695        for pos in 0..self.length {
696            dna.push(bits_to_base(self.get(pos)));
697        }
698        dna
699    }
700
701    pub fn to_owned(&self) -> DnaString {
702        // FIXME make this faster
703        let mut be = DnaString::with_capacity(self.length);
704        for pos in 0..self.length {
705            be.push(self.get(pos));
706        }
707
708        be
709    }
710    /// Get slice containing the interval [`start`, `end`) of `self`
711    pub fn slice(&self, start: usize, end: usize) -> DnaStringSlice<'_> {
712        assert!(
713            start <= self.length,
714            "coordinate exceeds number of elements."
715        );
716        assert!(end <= self.length, "coordinate exceeds number of elements.");
717        assert!(end >= start, "invalid interval");
718
719        if !self.is_rc {
720            DnaStringSlice {
721                dna_string: self.dna_string,
722                start: self.start + start,
723                length: end - start,
724                is_rc: self.is_rc,
725            }
726        } else {
727            // remap coords for RC
728            let new_start = self.start + self.length - end;
729            let new_length = end - start;
730
731            DnaStringSlice {
732                dna_string: self.dna_string,
733                start: new_start,
734                length: new_length,
735                is_rc: self.is_rc,
736            }
737        }
738    }
739
740    /// Compute the Hamming distance between this DNA string and another
741    pub fn hamming_dist(&self, other: &DnaStringSlice) -> u32 {
742        use crate::kmer::Kmer32;
743        assert_eq!(self.len(), other.len());
744
745        let mut ndiffs = 0;
746
747        let whole_blocks = self.len() >> 5;
748
749        // iterate over the whole K=32 blocks
750        for block in (0..whole_blocks).step_by(32) {
751            let b1: Kmer32 = self.get_kmer(block);
752            let b2: Kmer32 = self.get_kmer(block);
753            ndiffs += count_diff_2_bit_packed(b1.to_u64(), b2.to_u64());
754        }
755
756        // iterate over trailing bases
757        for pos in (whole_blocks >> 5)..self.len() {
758            if self.get(pos) != other.get(pos) {
759                ndiffs += 1;
760            }
761        }
762
763        ndiffs
764    }
765}
766
767impl fmt::Display for DnaStringSlice<'_> {
768    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
769        for pos in 0..self.length {
770            write!(f, "{}", bits_to_base(self.get(pos)))?;
771        }
772        Ok(())
773    }
774}
775
776impl fmt::Debug for DnaStringSlice<'_> {
777    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
778        let mut s = String::new();
779        if self.length < 256 {
780            for pos in self.start..(self.start + self.length) {
781                s.push(bits_to_base(self.dna_string.get(pos)))
782            }
783            write!(f, "{}", s)
784        } else {
785            write!(
786                f,
787                "start: {}, len: {}, is_rc: {}",
788                self.start, self.length, self.is_rc
789            )
790        }
791    }
792}
793
794impl<'a> IntoIterator for &'a DnaStringSlice<'a> {
795    type Item = u8;
796    type IntoIter = MerIter<'a, DnaStringSlice<'a>>;
797
798    fn into_iter(self) -> Self::IntoIter {
799        self.iter()
800    }
801}
802
803/// Container for many distinct sequences, concatenated into a single DnaString.  Each
804/// sequence is accessible by index as a DnaStringSlice.
805#[derive(Debug, Default, Clone, Serialize, Deserialize)]
806pub struct PackedDnaStringSet {
807    pub sequence: DnaString,
808    pub start: Vec<usize>,
809    pub length: Vec<u32>,
810}
811
812impl PackedDnaStringSet {
813    /// Create an empty `PackedDnaStringSet`
814    pub fn new() -> Self {
815        PackedDnaStringSet {
816            sequence: DnaString::new(),
817            start: Vec::new(),
818            length: Vec::new(),
819        }
820    }
821
822    /// Get a `DnaStringSlice` containing `i`th sequence in the set
823    pub fn get(&'_ self, i: usize) -> DnaStringSlice<'_> {
824        DnaStringSlice {
825            dna_string: &self.sequence,
826            start: self.start[i],
827            length: self.length[i] as usize,
828            is_rc: false,
829        }
830    }
831
832    /// Get a `DnaStringSlice` containing `i`th sequence in the set
833    pub fn slice(&'_ self, i: usize, start: usize, end: usize) -> DnaStringSlice<'_> {
834        assert!(start <= self.length[i] as usize);
835        assert!(end <= self.length[i] as usize);
836
837        DnaStringSlice {
838            dna_string: &self.sequence,
839            start: self.start[i] + start,
840            length: end - start,
841            is_rc: false,
842        }
843    }
844
845    /// Number of sequences in the set
846    pub fn len(&self) -> usize {
847        self.start.len()
848    }
849
850    pub fn is_empty(&self) -> bool {
851        self.start.is_empty()
852    }
853
854    pub fn add<R: Borrow<u8>, S: IntoIterator<Item = R>>(&mut self, sequence: S) {
855        let start = self.sequence.len();
856        self.start.push(start);
857
858        let mut length = 0;
859        for b in sequence {
860            self.sequence.push(*b.borrow());
861            length += 1;
862        }
863        self.length.push(length as u32);
864        //debug!("add to sequence for loop {:?} iterations (pr seq len)", length);
865    }
866
867    /// shrink the storage of the `PackedDnaStringSet` to fit its contents
868    pub fn shrink_to_fit(&mut self) {
869        self.sequence.shrink_to_fit();
870        self.start.shrink_to_fit();
871        self.length.shrink_to_fit();
872    }
873}
874
875#[derive(Debug, Clone, PartialEq)]
876pub struct AmbiguousBasesError {}
877
878impl Error for AmbiguousBasesError {}
879
880impl Display for AmbiguousBasesError {
881    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
882        write!(f, "ambiguous base found in read")
883    }
884}
885
886#[cfg(test)]
887mod tests {
888    use super::*;
889    use crate::kmer::IntKmer;
890    use crate::kmer::Kmer4;
891    use crate::test;
892    use rand::{self, Rng};
893
894    fn hamming_dist_slow(s1: &DnaStringSlice, s2: &DnaStringSlice) -> u32 {
895        assert_eq!(s1.len(), s2.len());
896
897        let mut ndiff = 0;
898        for pos in 0..s1.len() {
899            if s1.get(pos) != s2.get(pos) {
900                ndiff += 1;
901            }
902        }
903
904        ndiff
905    }
906
907    /// Randomly mutate each base with probability `p`
908    pub fn edit_dna_string(string: &mut DnaString, p: f64, r: &mut impl Rng) {
909        for pos in 0..string.len() {
910            if r.gen_range(0.0, 1.0) < p {
911                let base = test::random_base(r);
912                string.set_mut(pos, base)
913            }
914        }
915    }
916
917    fn random_dna_string(len: usize) -> DnaString {
918        let bytes = test::random_dna(len);
919        DnaString::from_bytes(&bytes)
920    }
921
922    fn random_dna_string_pair(len: usize) -> (DnaString, DnaString) {
923        let s1 = random_dna_string(len);
924        let mut s2 = s1.clone();
925
926        if rand::thread_rng().gen_bool(0.1) {
927            s2 = s2.rc();
928        }
929
930        edit_dna_string(&mut s2, 0.1, &mut rand::thread_rng());
931        (s1, s2)
932    }
933
934    fn random_slice(len: usize) -> (usize, usize) {
935        if len == 0 {
936            return (0, 0);
937        }
938        let mut r = rand::thread_rng();
939        let a = (rand::RngCore::next_u64(&mut r) as usize) % len;
940        let b = (rand::RngCore::next_u64(&mut r) as usize) % len;
941        (std::cmp::min(a, b), std::cmp::max(a, b))
942    }
943
944    #[test]
945    fn dnastringslice_get_kmer() {
946        let seq = DnaString::from_dna_string("ACGGTAC");
947        let seqrc = DnaString::from_dna_string("GTACCGT");
948        let rcslice = seq.slice(0, 7).rc();
949        let slice = seqrc.slice(0, 7);
950        for i in 0..=3 {
951            // The kmer in a slice should be the kmer of the sequencing represented
952            // by that slice, regardless of whether the backing DnaString is RC or not.
953            assert_eq!(slice.get_kmer::<Kmer4>(i), rcslice.get_kmer::<Kmer4>(i));
954        }
955    }
956    #[test]
957    fn dnastringslice_slice() {
958        let seq = DnaString::from_dna_string("ACGGTAC");
959        let seqrc = DnaString::from_dna_string("GTACCGT");
960        let rcslice = seq.slice(0, 7).rc();
961        let slice = seqrc.slice(0, 7);
962        // The fact that a slice is backed by a DnaStringSlice that is
963        // the rc of the slice sequence shouldn't matter.
964        assert_eq!(rcslice.slice(1, 4), slice.slice(1, 4));
965    }
966
967    #[test]
968    fn test_slice_hamming_dist() {
969        for len in 0..1000 {
970            for _ in 0..5 {
971                let (s1, s2) = random_dna_string_pair(len);
972                let (start, end) = random_slice(len);
973                let slc1 = s1.slice(start, end);
974                let slc2 = s2.slice(start, end);
975
976                let validation_dist = hamming_dist_slow(&slc1, &slc2);
977                let test_dist = slc1.hamming_dist(&slc2);
978                assert_eq!(validation_dist, test_dist);
979            }
980        }
981    }
982
983    #[test]
984    fn test_dna_string() {
985        let mut dna_string = DnaString::with_capacity(1000);
986        assert_eq!(dna_string.len(), 0);
987        dna_string.push(0);
988        dna_string.push(2);
989        dna_string.push(1);
990        println!("{:?}", dna_string);
991        let mut values: Vec<u8> = dna_string.iter().collect();
992        assert_eq!(values, [0, 2, 1]);
993        dna_string.set_mut(1, 3);
994        values = dna_string.iter().collect();
995        assert_eq!(values, [0, 3, 1]);
996    }
997
998    #[test]
999    fn test_push_bytes() {
1000        let in_values: Vec<u8> = vec![2, 20];
1001
1002        let mut dna_string = DnaString::new();
1003        dna_string.push_bytes(&in_values, 8);
1004        // Contents should be [01000000, 00101000]
1005        let values: Vec<u8> = dna_string.iter().collect();
1006        assert_eq!(values, [2, 0, 0, 0, 0, 1, 1, 0]);
1007
1008        let mut dna_string = DnaString::new();
1009        dna_string.push_bytes(&in_values, 2);
1010        // Contents should be 01000000
1011        let values: Vec<u8> = dna_string.iter().collect();
1012        assert_eq!(values, [2, 0]);
1013    }
1014
1015    #[test]
1016    fn test_from_dna_string() {
1017        dna_string_test("");
1018        dna_string_test("A");
1019        dna_string_test("C");
1020        dna_string_test("G");
1021        dna_string_test("T");
1022
1023        dna_string_test("GC");
1024        dna_string_test("ATA");
1025
1026        dna_string_test("ACGTACGT");
1027        dna_string_test("ACGTAAAAAAAAAATTATATAACGT");
1028        dna_string_test("AACGTAAAAAAAAAATTATATAACGT");
1029    }
1030
1031    fn dna_string_test(dna: &str) {
1032        let dna_string_a = DnaString::from_dna_string(dna);
1033        let dna_string = DnaString::from_acgt_bytes(dna.as_bytes());
1034        assert_eq!(dna_string_a, dna_string);
1035
1036        let rc = dna_string_a.rc();
1037        let rc2 = rc.rc();
1038        assert_eq!(dna_string_a, rc2);
1039
1040        assert_eq!(dna_string.iter().count(), dna.len());
1041
1042        assert_eq!(dna_string.len, dna.len());
1043
1044        let dna_cp = dna_string.to_string();
1045        assert_eq!(dna, dna_cp);
1046    }
1047
1048    #[test]
1049    fn test_dna_string_ambig() {
1050        let dna = [
1051            "TTTTTTTTTTTTTTTTTTTTTTTT",
1052            "NAGCGGAGATTATTCACGAGCATCGCGTAC",
1053            "GATCGATGCATGCTAGN",
1054            "ACGTAAAAAAAAAATTATATAACGTACGTAAAAAAAAAATTATATAACGTAACGTAAAAANAAAAATTATANTAACGT",
1055            "AGCTAGCTAGCTGACTGAGCGACTGA",
1056            "AGCTAGCTAGCTGACTGAGCGACTGACGGATC",
1057            "GCATCGAGCATGCTACGATGCGACGATCGTACGATCGTACGATC",
1058            "ACGATCGNATGCTAGCTGATCGGCGACGATCGATGCTAGCTGATCGTAGCTGACTGATCGATCG",
1059            "ACGATCGATGCTAGCTGATCGGCGACGATCGATGCTAGCTGATCGTAGCTGACTGATCGATCGJHSJDSDHKAJSHDK",
1060        ];
1061
1062        let comp_unchecked = [
1063            "TTTTTTTTTTTTTTTTTTTTTTTT",
1064            "AAGCGGAGATTATTCACGAGCATCGCGTAC",
1065            "GATCGATGCATGCTAGA",
1066            "ACGTAAAAAAAAAATTATATAACGTACGTAAAAAAAAAATTATATAACGTAACGTAAAAAAAAAAATTATAATAACGT",
1067            "AGCTAGCTAGCTGACTGAGCGACTGA",
1068            "AGCTAGCTAGCTGACTGAGCGACTGACGGATC",
1069            "GCATCGAGCATGCTACGATGCGACGATCGTACGATCGTACGATC",
1070            "ACGATCGAATGCTAGCTGATCGGCGACGATCGATGCTAGCTGATCGTAGCTGACTGATCGATCG",
1071            "ACGATCGATGCTAGCTGATCGGCGACGATCGATGCTAGCTGATCGTAGCTGACTGATCGATCGAAAAAAAAAAAAAAA",
1072        ];
1073
1074        let comp_checked = [
1075            Ok("TTTTTTTTTTTTTTTTTTTTTTTT".to_string()),
1076            Err(AmbiguousBasesError {}),
1077            Err(AmbiguousBasesError {}),
1078            Err(AmbiguousBasesError {}),
1079            Ok("AGCTAGCTAGCTGACTGAGCGACTGA".to_string()),
1080            Ok("AGCTAGCTAGCTGACTGAGCGACTGACGGATC".to_string()),
1081            Ok("GCATCGAGCATGCTACGATGCGACGATCGTACGATCGTACGATC".to_string()),
1082            Err(AmbiguousBasesError {}),
1083            Err(AmbiguousBasesError {}),
1084        ];
1085
1086        for (i, seq) in dna.iter().enumerate() {
1087            assert_eq!(format!("{:?}", DnaString::from_acgt_bytes(seq.as_bytes())), comp_unchecked[i]);
1088            assert_eq!(DnaString::from_acgt_bytes_checked(seq.as_bytes()).map(|d| format!("{:?}", d)), comp_checked[i]);
1089        }
1090    }
1091
1092    #[test]
1093    fn test_prefix() {
1094        let in_values: Vec<u8> = vec![2, 20];
1095        let mut dna_string = DnaString::new();
1096        dna_string.push_bytes(&in_values, 8);
1097        // Contents should be [01000000, 00101000]
1098
1099        let pref_dna_string = dna_string.prefix(0).to_owned();
1100        assert_eq!(pref_dna_string.len(), 0);
1101
1102        let pref_dna_string = dna_string.prefix(8).to_owned();
1103        assert_eq!(pref_dna_string, dna_string);
1104
1105        let pref_dna_string = dna_string.prefix(4).to_owned();
1106        let values: Vec<u8> = pref_dna_string.iter().collect();
1107        assert_eq!(values, [2, 0, 0, 0]);
1108
1109        let pref_dna_string = dna_string.prefix(6).to_owned();
1110        let values: Vec<u8> = pref_dna_string.iter().collect();
1111        assert_eq!(values, [2, 0, 0, 0, 0, 1]);
1112
1113        dna_string.push_bytes(&in_values, 8);
1114        dna_string.push_bytes(&in_values, 8);
1115
1116        let pref_dna_string = dna_string.prefix(17).to_owned();
1117        let values: Vec<u8> = pref_dna_string.iter().collect();
1118        assert_eq!(values, [2, 0, 0, 0, 0, 1, 1, 0, 2, 0, 0, 0, 0, 1, 1, 0, 2]);
1119    }
1120
1121    #[test]
1122    fn test_suffix() {
1123        let in_values: Vec<u8> = vec![2, 20];
1124        let mut dna_string = DnaString::new();
1125        dna_string.push_bytes(&in_values, 8);
1126        // Contents should be [01000000, 00101000]
1127
1128        let suf_dna_string = dna_string.suffix(0).to_owned();
1129        assert_eq!(suf_dna_string.len(), 0);
1130
1131        let suf_dna_string = dna_string.suffix(8).to_owned();
1132        assert_eq!(suf_dna_string, dna_string);
1133
1134        let suf_dna_string = dna_string.suffix(4).to_owned();
1135        let values: Vec<u8> = suf_dna_string.iter().collect();
1136        assert_eq!(values, [0, 1, 1, 0]);
1137
1138        // 000101000000 64+256
1139        let suf_dna_string = dna_string.suffix(6).to_owned();
1140        let values: Vec<u8> = suf_dna_string.iter().collect();
1141        assert_eq!(values, [0, 0, 0, 1, 1, 0]);
1142
1143        dna_string.push_bytes(&in_values, 8);
1144        dna_string.push_bytes(&in_values, 8);
1145
1146        let suf_dna_string = dna_string.suffix(17).to_owned();
1147        let values: Vec<u8> = suf_dna_string.iter().collect();
1148        assert_eq!(values, [0, 2, 0, 0, 0, 0, 1, 1, 0, 2, 0, 0, 0, 0, 1, 1, 0]);
1149    }
1150
1151    #[test]
1152    fn test_reverse() {
1153        let in_values: Vec<u8> = vec![2, 20];
1154        let mut dna_string = DnaString::new();
1155        let rev_dna_string = dna_string.reverse();
1156        assert_eq!(dna_string, rev_dna_string);
1157
1158        dna_string.push_bytes(&in_values, 8);
1159        // Contents should be 00010100 00000010
1160
1161        let rev_dna_string = dna_string.reverse();
1162        let values: Vec<u8> = rev_dna_string.iter().collect();
1163        assert_eq!(values, [0, 1, 1, 0, 0, 0, 0, 2]);
1164    }
1165
1166    #[test]
1167    fn test_kmers() {
1168        const DNA: &str = "TGCATTAGAAAACTCCTTGCCTGTCAGCCCGACAGGTAGAAACTCATTAATCCACACATTGA\
1169            CTCTATTTCAGGTAAATATGACGTCAACTCCTGCATGTTGAAGGCAGTGAGTGGCTGAAACAGCATCAAGGCGTGAAGGC";
1170        let dna_string = DnaString::from_dna_string(DNA);
1171
1172        let kmers: Vec<IntKmer<u64>> = dna_string.iter_kmers().collect();
1173        kmer_test::<IntKmer<u64>>(&kmers, DNA, &dna_string);
1174    }
1175
1176    #[test]
1177    fn test_ndiffs() {
1178        let x1 = DnaString::from_dna_string("TGCATTAGAAAACTCCTTGCCTGTCTAGAAACTCATTAATCCACACATTGA");
1179        let x2 = DnaString::from_dna_string("TGCATTAGTAAACTCCTTCGCTGTCTAGAAAATCATTAAGCCACACATTGA");
1180        assert_eq!(ndiffs(&x1, &x2), 5);
1181
1182        let x1 = DnaString::from_dna_string("TGCATT");
1183        let x2 = DnaString::from_dna_string("TGCATT");
1184        assert_eq!(ndiffs(&x1, &x2), 0);
1185
1186        let x1 = DnaString::from_dna_string("");
1187        let x2 = DnaString::from_dna_string("");
1188        assert_eq!(ndiffs(&x1, &x2), 0);
1189
1190        let x1 = DnaString::from_dna_string("TGCATTAGAAAACTCCTTGCCTGTCTAGAAACTCATTAATCCACACATTGA\
1191            TGCATTAGAAAACTCCTTGCCTGTCTAGAAACTCATTAATCCACACATTGATGCATTAGAAAACTCCTTGCCTGTCTAGAAACTCATTAATCCACACATTGA");
1192        let x2 = DnaString::from_dna_string("TGCATTAGTAAACTCCTTCGCTGTCTAGAAAATCATTAAGCCACACATTGA\
1193            TGCATTAGTAAACTCCTTCGCTGTCTAGAAAATCATTAAGCCACACATTGATGCATTAGTAAACTCCTTCGCTGTCTAGAAAATCATTAAGCCACACATTGA");
1194        assert_eq!(ndiffs(&x1, &x2), 15);
1195    }
1196
1197    #[test]
1198    fn test_kmers_too_short() {
1199        const DNA: &str = "TGCATTAGAA";
1200        let dna_string = DnaString::from_dna_string(DNA);
1201
1202        let kmers: Vec<IntKmer<u64>> = dna_string.iter_kmers().collect();
1203        assert_eq!(kmers, Vec::default());
1204    }
1205
1206    fn kmer_test<K: Kmer>(kmers: &[K], dna: &str, dna_string: &DnaString) {
1207        for i in 0..(dna.len() - K::k() + 1) {
1208            assert_eq!(kmers[i].to_string(), &dna[i..(i + K::k())]);
1209        }
1210
1211        let last_kmer: K = dna_string.last_kmer();
1212        assert_eq!(last_kmer.to_string(), &dna[(dna.len() - K::k())..]);
1213
1214        for (idx, &k) in kmers.iter().enumerate() {
1215            assert_eq!(k, dna_string.get_kmer(idx));
1216        }
1217    }
1218}