Skip to main content

bio_seq/seq/
iterators.rs

1// Copyright 2021, 2022 Jeff Knaggs
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
6use crate::codec::Codec;
7use crate::kmer::KmerIter;
8use crate::seq::{Seq, SeqSlice};
9use core::iter::Chain;
10use core::marker::PhantomData;
11
12/// An iterator over fixed-size non-overlapping chunks of a sequence
13pub struct SeqChunks<'a, A: Codec> {
14    slice: &'a SeqSlice<A>,
15    width: usize,
16    skip: usize,
17    index: usize,
18}
19
20/// An iterator over the elements of a sequence
21pub struct SeqIter<'a, A: Codec> {
22    slice: &'a SeqSlice<A>,
23    index: usize,
24}
25
26impl<'a, A: Codec> SeqSlice<A> {
27    pub fn chain(
28        self: &'a SeqSlice<A>,
29        second: &'a SeqSlice<A>,
30    ) -> Chain<SeqIter<'a, A>, SeqIter<'a, A>> {
31        self.into_iter().chain(second)
32    }
33
34    pub fn iter(&'a self) -> SeqIter<'a, A> {
35        <&Self as IntoIterator>::into_iter(self)
36    }
37
38    /// Iterate over sliding windows of length K
39    pub fn kmers<const K: usize>(&self) -> KmerIter<'_, A, K> {
40        KmerIter::<A, K> {
41            slice: self,
42            index: 0,
43            len: self.len(),
44            _p: PhantomData,
45        }
46    }
47
48    /// Iterate over the sequence in reverse order
49    pub fn rev_iter(&self) -> RevIter<'_, A> {
50        RevIter {
51            slice: self,
52            index: self.len(),
53        }
54    }
55
56    /// Iterate over the sequence in overlapping windows of a specified width
57    ///
58    /// ```
59    /// use bio_seq::prelude::*;
60    ///
61    /// let seq: Seq<Dna> = "ACTGATCG".try_into().unwrap();
62    /// let windows: Vec<String> = seq.windows(3).map(String::from).collect();
63    /// assert_eq!(windows, vec!["ACT", "CTG", "TGA", "GAT", "ATC", "TCG"]);
64    /// ```
65    ///
66    /// # Panics
67    ///
68    /// Panics if `width` is 0.
69    pub fn windows(&self, width: usize) -> SeqChunks<'_, A> {
70        assert!(width != 0, "window must be non-zero");
71        SeqChunks {
72            slice: self,
73            width,
74            skip: 1,
75            index: 0,
76        }
77    }
78
79    /// Iterate over the sequence in non-overlapping chunks of a specified width
80    ///
81    /// The last incomplete chunk will be excluded if the sequence length is not divisible by the specified
82    /// width.
83    ///
84    /// ```
85    /// use bio_seq::prelude::*;
86    ///
87    /// let seq: Seq<Dna> = "ACTGATCG".try_into().unwrap();
88    /// let chunks: Vec<Seq<Dna>> = seq.chunks(3).collect();
89    /// assert_eq!(chunks, vec![dna!("ACT"), dna!("GAT")]);
90    /// ```
91    ///
92    /// # Panics
93    ///
94    /// Panics if `width` is 0.
95    pub fn chunks(&self, width: usize) -> SeqChunks<'_, A> {
96        assert!(width != 0, "chunk must be non-zero");
97        SeqChunks {
98            slice: self,
99            width,
100            skip: width,
101            index: 0,
102        }
103    }
104}
105
106/// An iterator over the elements of a sequence in reverse order
107pub struct RevIter<'a, A: Codec> {
108    pub(crate) slice: &'a SeqSlice<A>,
109    pub(crate) index: usize,
110}
111
112impl<A: Codec> Iterator for RevIter<'_, A> {
113    type Item = A;
114    fn next(&mut self) -> Option<A> {
115        let i = self.index;
116
117        if self.index == 0 {
118            return None;
119        }
120        self.index -= 1;
121        Some(A::unsafe_from_bits(self.slice[i - 1].into()))
122    }
123
124    fn size_hint(&self) -> (usize, Option<usize>) {
125        (self.index, Some(self.index))
126    }
127}
128
129impl<A: Codec> ExactSizeIterator for RevIter<'_, A> {}
130
131impl<'a, A: Codec> Iterator for SeqChunks<'a, A> {
132    type Item = &'a SeqSlice<A>;
133
134    fn next(&mut self) -> Option<Self::Item> {
135        let i = self.index;
136        if i + self.width > self.slice.len() {
137            return None;
138        }
139        self.index += self.skip;
140        Some(&self.slice[i..i + self.width])
141    }
142
143    fn size_hint(&self) -> (usize, Option<usize>) {
144        let n = self
145            .slice
146            .len()
147            .saturating_sub(self.index)
148            .saturating_sub(self.width - 1)
149            .div_ceil(self.skip);
150        (n, Some(n))
151    }
152}
153
154impl<A: Codec> ExactSizeIterator for SeqChunks<'_, A> {}
155
156impl<'a, A: Codec> IntoIterator for &'a Seq<A> {
157    type Item = A;
158    type IntoIter = SeqIter<'a, A>;
159
160    fn into_iter(self) -> Self::IntoIter {
161        self.as_ref().into_iter()
162    }
163}
164
165impl<'a, A: Codec> IntoIterator for &'a SeqSlice<A> {
166    type Item = A;
167    type IntoIter = SeqIter<'a, A>;
168
169    fn into_iter(self) -> Self::IntoIter {
170        SeqIter {
171            slice: self,
172            index: 0,
173        }
174    }
175}
176
177impl<'a, A: Codec> FromIterator<&'a SeqSlice<A>> for Vec<Seq<A>> {
178    fn from_iter<T: IntoIterator<Item = &'a SeqSlice<A>>>(iter: T) -> Self {
179        iter.into_iter().map(ToOwned::to_owned).collect()
180    }
181}
182
183impl<A: Codec> Iterator for SeqIter<'_, A> {
184    type Item = A;
185    fn next(&mut self) -> Option<A> {
186        let i = self.index;
187        if self.index >= self.slice.len() {
188            return None;
189        }
190        self.index += 1;
191        Some(A::unsafe_from_bits(self.slice[i].into()))
192    }
193
194    fn size_hint(&self) -> (usize, Option<usize>) {
195        let n = self.slice.len() - self.index;
196        (n, Some(n))
197    }
198}
199
200impl<A: Codec> ExactSizeIterator for SeqIter<'_, A> {}
201
202#[cfg(test)]
203mod tests {
204    use crate::codec::dna::Dna::*;
205    use crate::prelude::*;
206
207    #[test]
208    fn seq_iter() {
209        let seq: Seq<Dna> = dna!("ACTGATCGATAC").into();
210        let elements: Vec<Dna> = seq.into_iter().collect();
211        assert_eq!(elements, vec![A, C, T, G, A, T, C, G, A, T, A, C]);
212        assert_ne!(elements, vec![A, C, T, G, A, T, C, G, A, T, A, C, A]);
213        assert_ne!(elements, vec![C, A, T, A, G, C, T, A, G, T, C, A]);
214    }
215
216    #[test]
217    fn rev_iter() {
218        let seq: Seq<Dna> = dna!("ACTGATCGATAC").into();
219        let rev_elements: Vec<Dna> = seq.rev_iter().collect();
220        assert_ne!(rev_elements, vec![A, C, T, G, A, T, C, G, A, T, A, C]);
221        assert_eq!(rev_elements, vec![C, A, T, A, G, C, T, A, G, T, C, A]);
222    }
223
224    #[test]
225    fn iterators() {
226        let seq: Seq<Dna> = dna!("ACTGATCGATAC").into();
227        let slice = &seq[2..9];
228        let elements: Vec<Dna> = slice.into_iter().collect();
229        assert_eq!(elements, vec![T, G, A, T, C, G, A]);
230    }
231
232    #[test]
233    fn chunks() {
234        let seq: Seq<Dna> = dna!("ACTGATCGATAC").into();
235        let cs: Vec<Seq<Dna>> = seq.chunks(5).collect();
236        assert_eq!(cs[0], dna!("ACTGA"));
237        assert_eq!(cs[1], dna!("TCGAT"));
238        assert_eq!(cs.len(), 2);
239    }
240
241    #[test]
242    fn test_chain() {
243        let seq1 = Seq::<Dna>::try_from("ATG").unwrap();
244        let seq2 = Seq::<Dna>::try_from("TAC").unwrap();
245
246        let chained = seq1.chain(&seq2);
247
248        let expected_seq = Seq::<Dna>::try_from("ATGTAC").unwrap();
249        for (a, b) in chained.zip(&expected_seq) {
250            assert_eq!(a, b);
251        }
252
253        let chained = seq1.chain(&seq2);
254        for (a, b) in chained.map(|b| b.to_comp()).zip(&expected_seq) {
255            assert_ne!(a, b);
256        }
257    }
258
259    #[test]
260    fn windows() {
261        let seq: Seq<Dna> = dna!("ACTGATACG").into();
262        let windows: Vec<String> = seq.windows(5).map(String::from).collect();
263        assert_eq!(windows, vec!["ACTGA", "CTGAT", "TGATA", "GATAC", "ATACG"]);
264    }
265
266    #[test]
267    #[should_panic(expected = "chunk must be non-zero")]
268    fn zero_width_chunk_panics() {
269        let _ = dna!("ACGTACGTGA").chunks(0);
270    }
271
272    #[test]
273    #[should_panic(expected = "window must be non-zero")]
274    fn zero_width_window_panics() {
275        let _ = dna!("ACGTACGTGA").windows(0);
276    }
277}