d4_hts/
lib.rs

1pub mod alignment;
2
3pub use alignment::*;
4
5use std::cmp::Reverse;
6use std::collections::BinaryHeap;
7
8pub struct DepthIter<'a, R: AlignmentReader<'a>> {
9    iter: AlignmentIter<'a, R>,
10    cur_chrom: i32,
11    cur_pos: usize,
12    heap: BinaryHeap<Reverse<usize>>,
13    next_read: Option<(i32, usize, usize)>,
14    filter: Option<Box<dyn Fn(&Alignment<'_>) -> bool + 'a>>,
15}
16
17impl<'a, R: AlignmentReader<'a>> DepthIter<'a, R> {
18    pub fn with_filter<F: Fn(&Alignment<'_>) -> bool + 'a>(reader: R, filter: F) -> Self {
19        let (chrom, pos) = reader.start();
20        let iter = reader.into_alignment_iter();
21
22        let mut ret = Self {
23            iter,
24            next_read: None,
25            cur_chrom: chrom as i32,
26            cur_pos: pos as usize,
27            heap: BinaryHeap::new(),
28            filter: Some(Box::new(filter)),
29        };
30
31        ret.load_next();
32        ret
33    }
34
35    fn load_next(&mut self) {
36        self.next_read = loop {
37            if let Some(Ok(read)) = self.iter.next() {
38                if self.filter.as_ref().map_or(true, |predict| predict(&read)) {
39                    break Some(read);
40                }
41            } else {
42                break None;
43            }
44        }
45        .map(|read| (read.ref_id(), read.ref_begin(), read.ref_end()));
46    }
47}
48
49impl<'a, R: AlignmentReader<'a>> Iterator for DepthIter<'a, R> {
50    type Item = (i32, usize, u32);
51    fn next(&mut self) -> Option<Self::Item> {
52        if self.next_read.is_none() && self.heap.is_empty() {
53            return None;
54        }
55
56        let ret = (self.cur_chrom, self.cur_pos, self.heap.len() as u32);
57
58        self.cur_pos += 1;
59
60        while let Some((tid, left, right)) = self.next_read {
61            if tid != self.cur_chrom {
62                if self.heap.is_empty() {
63                    self.cur_chrom = tid;
64                    self.cur_pos = 0;
65                }
66                break;
67            }
68            if left > self.cur_pos {
69                break;
70            }
71            self.heap.push(Reverse(right));
72            self.load_next();
73        }
74
75        while self.heap.peek().map_or(false, |x| x.0 < self.cur_pos) {
76            self.heap.pop();
77        }
78
79        Some(ret)
80    }
81}
82
83#[cfg(test)]
84mod tests {
85    #[test]
86    fn it_works() {
87        assert_eq!(2 + 2, 4);
88    }
89}