debruijn/
vmer.rs

1// Copyright 2017 10x Genomics
2
3//! Variable-length DNA strings packed into fixed-size structs.
4
5use std::cmp::{max, min};
6use std::fmt;
7use std::hash::Hash;
8
9use crate::bits_to_base;
10use crate::kmer::{IntHelp, IntKmer};
11use crate::Kmer;
12use crate::Mer;
13use crate::Vmer;
14use serde_derive::{Deserialize, Serialize};
15
16fn block_set(kmer: u64, pos: usize, val: u8) -> u64 {
17    let offset = (31 - pos) * 2;
18    let mask = !(3 << offset);
19
20    (kmer & mask) | ((val as u64) << offset)
21}
22
23fn block_get(kmer: u64, pos: usize) -> u8 {
24    let offset = (31 - pos) * 2;
25    ((kmer >> offset) & 3) as u8
26}
27
28pub type Lmer1 = Lmer<[u64; 1]>;
29pub type Lmer2 = Lmer<[u64; 2]>;
30pub type Lmer3 = Lmer<[u64; 3]>;
31
32/// Store a variable-length DNA sequence in a packed 2-bit encoding, up 92bp in length
33/// The length of the sequence is stored in the lower 8 bits of storage
34#[derive(Hash, Copy, Clone, PartialEq, PartialOrd, Eq, Ord, Serialize, Deserialize)]
35pub struct Lmer<A: Array> {
36    storage: A,
37}
38
39impl<A: Array<Item = u64> + Copy + Eq + Ord + Hash> Mer for Lmer<A> {
40    /// The length of the DNA string
41    fn len(&self) -> usize {
42        (self.storage.as_slice()[A::size() - 1] & 0xff) as usize
43    }
44
45    fn is_empty(&self) -> bool {
46        self.len() == 0
47    }
48
49    /// Get the base at position pos
50    fn get(&self, pos: usize) -> u8 {
51        let block = pos / 32;
52        let offset = pos % 32;
53        block_get(self.storage.as_slice()[block], offset)
54    }
55
56    /// Return a new Lmer with position pos set to base val
57    fn set_mut(&mut self, pos: usize, val: u8) {
58        let block = pos / 32;
59        let offset = pos % 32;
60
61        let block_val = block_set(self.storage.as_slice()[block], offset, val);
62        self.storage.as_mut_slice()[block] = block_val;
63    }
64
65    fn set_slice_mut(&mut self, pos: usize, n_bases: usize, value: u64) {
66        let slc = self.storage.as_mut_slice();
67        let b0 = pos / 32;
68        let block_pos = pos % 32;
69        let top_mask = IntKmer::<u64>::top_mask(block_pos);
70        let mut bottom_mask =
71            IntKmer::<u64>::bottom_mask(max(0, 32usize.saturating_sub(block_pos + n_bases)));
72        if b0 == A::size() - 1 {
73            bottom_mask |= 0xFF
74        }
75        let mask = top_mask | bottom_mask;
76
77        let nb0 = 32 - block_pos;
78        let value_top = value >> (block_pos * 2);
79        let v0 = (slc[b0] & mask) | (value_top & !mask);
80        slc[b0] = v0;
81
82        if n_bases > nb0 {
83            let b1 = b0 + 1;
84            let nb1 = n_bases - nb0;
85            let bottom_mask = IntKmer::<u64>::bottom_mask(32 - nb1);
86            let value_bottom = value << (nb0 * 2);
87            let v1 = (slc[b1] & bottom_mask) | (value_bottom & !bottom_mask);
88            slc[b1] = v1
89        }
90    }
91
92    fn rc(&self) -> Self {
93        let slc = self.storage.as_slice();
94
95        let mut new_lmer = Self::new(self.len());
96        let mut block = 0;
97        let mut pos = 0;
98
99        while pos < self.len() {
100            let n_bases = min(32, self.len() - pos);
101
102            let mut v = slc[block];
103            // Mask the length field packed into the last block
104            if block == A::size() - 1 {
105                v &= !0xFF
106            }
107
108            let v_rc = !v.reverse_by_twos() << (64 - n_bases * 2);
109            new_lmer.set_slice_mut(self.len() - pos - n_bases, n_bases, v_rc);
110            block += 1;
111            pos += n_bases;
112        }
113
114        new_lmer
115    }
116}
117
118impl<A: Array<Item = u64> + Copy + Eq + Ord + Hash> Vmer for Lmer<A> {
119    fn max_len() -> usize {
120        (A::size() * 64 - 8) / 2
121    }
122
123    /// Initialize an blank Lmer of length len.
124    /// Will initially represent all A's.
125    fn new(len: usize) -> Lmer<A> {
126        let mut arr = A::new();
127        {
128            let slc = arr.as_mut_slice();
129
130            // Write the length into the last 8 bits
131            slc[A::size() - 1] = (len as u64) & 0xff;
132        }
133        Lmer { storage: arr }
134    }
135
136    /// Get the kmer starting at position pos
137    fn get_kmer<K: Kmer>(&self, pos: usize) -> K {
138        assert!(self.len() - pos >= K::k());
139        let slc = self.storage.as_slice();
140
141        // Which block has the first base
142        let mut block = pos / 32;
143
144        // Where we are in the kmer
145        let mut kmer_pos = 0;
146
147        // Where are in the block
148        let mut block_pos = pos % 32;
149
150        let mut kmer = K::empty();
151
152        while kmer_pos < K::k() {
153            // get relevent bases for current block
154            let nb = min(K::k() - kmer_pos, 32 - block_pos);
155
156            let val = slc[block] << (2 * block_pos);
157            kmer.set_slice_mut(kmer_pos, nb, val);
158
159            // move to next block, move ahead in kmer.
160            block += 1;
161            kmer_pos += nb;
162            // alway start a beginning of next block
163            block_pos = 0;
164        }
165
166        kmer
167    }
168}
169
170impl<A: Array<Item = u64> + Copy + Eq + Ord + Hash> fmt::Debug for Lmer<A> {
171    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
172        let mut s = String::new();
173        for pos in 0..self.len() {
174            s.push(bits_to_base(self.get(pos)))
175        }
176
177        write!(f, "{}", s)
178    }
179}
180
181/// Types that can be used as the backing store for a SmallVec
182pub trait Array {
183    type Item;
184    fn new() -> Self;
185    fn size() -> usize;
186
187    fn as_slice(&self) -> &[Self::Item];
188    fn as_mut_slice(&mut self) -> &mut [Self::Item];
189}
190
191macro_rules! impl_array(
192    ($($size:expr),+) => {
193        $(
194            impl<T: Default + Copy + Eq + Ord> Array for [T; $size] {
195                type Item = T;
196                fn new() -> [T; $size] { [T::default(); $size] }
197                fn size() -> usize { $size }
198                fn as_slice(&self) -> &[T] { self }
199                fn as_mut_slice(&mut self) -> &mut [T] { self }
200
201            }
202        )+
203    }
204);
205
206impl_array!(1, 2, 3, 4, 5, 6);
debruijn/vmer.rs

debruijn/
vmer.rs
