noodles_cram/io/reader/
query.rs

1use std::{
2    io::{self, Read, Seek, SeekFrom},
3    slice, vec,
4};
5
6use noodles_core::region::Interval;
7use noodles_sam as sam;
8
9use super::{Container, Reader};
10use crate::crai;
11
12/// An iterator over records that intersect a given region.
13///
14/// This is created by calling [`Reader::query`].
15pub struct Query<'a, R>
16where
17    R: Read + Seek,
18{
19    reader: &'a mut Reader<R>,
20
21    header: &'a sam::Header,
22
23    index: slice::Iter<'a, crai::Record>,
24
25    reference_sequence_id: usize,
26    interval: Interval,
27
28    records: vec::IntoIter<sam::alignment::RecordBuf>,
29}
30
31impl<'a, R> Query<'a, R>
32where
33    R: Read + Seek,
34{
35    pub(super) fn new(
36        reader: &'a mut Reader<R>,
37        header: &'a sam::Header,
38        index: &'a crai::Index,
39        reference_sequence_id: usize,
40        interval: Interval,
41    ) -> Self {
42        Self {
43            reader,
44
45            header,
46
47            index: index.iter(),
48
49            reference_sequence_id,
50            interval,
51
52            records: Vec::new().into_iter(),
53        }
54    }
55
56    fn read_next_container(&mut self) -> Option<io::Result<()>> {
57        let index_record = self.index.next()?;
58
59        if index_record.reference_sequence_id() != Some(self.reference_sequence_id) {
60            return Some(Ok(()));
61        }
62
63        if let Err(e) = self.reader.seek(SeekFrom::Start(index_record.offset())) {
64            return Some(Err(e));
65        }
66
67        let mut container = Container::default();
68
69        match self.reader.read_container(&mut container) {
70            Ok(0) => return None,
71            Ok(_) => {}
72            Err(e) => return Some(Err(e)),
73        };
74
75        let compression_header = match container.compression_header() {
76            Ok(compression_header) => compression_header,
77            Err(e) => return Some(Err(e)),
78        };
79
80        let records = container
81            .slices()
82            .map(|result| {
83                let slice = result?;
84
85                let (core_data_src, external_data_srcs) = slice.decode_blocks()?;
86
87                slice
88                    .records(
89                        self.reader.reference_sequence_repository.clone(),
90                        self.header,
91                        &compression_header,
92                        &core_data_src,
93                        &external_data_srcs,
94                    )
95                    .and_then(|records| {
96                        records
97                            .into_iter()
98                            .map(|record| {
99                                sam::alignment::RecordBuf::try_from_alignment_record(
100                                    self.header,
101                                    &record,
102                                )
103                            })
104                            .collect::<io::Result<Vec<_>>>()
105                    })
106            })
107            .collect::<Result<Vec<_>, _>>();
108
109        let records = match records {
110            Ok(records) => records,
111            Err(e) => return Some(Err(e)),
112        };
113
114        self.records = records
115            .into_iter()
116            .flatten()
117            .collect::<Vec<_>>()
118            .into_iter();
119
120        Some(Ok(()))
121    }
122}
123
124impl<R> Iterator for Query<'_, R>
125where
126    R: Read + Seek,
127{
128    type Item = io::Result<sam::alignment::RecordBuf>;
129
130    fn next(&mut self) -> Option<Self::Item> {
131        loop {
132            match self.records.next() {
133                Some(r) => {
134                    if let (Some(start), Some(end)) = (r.alignment_start(), r.alignment_end()) {
135                        let alignment_interval = (start..=end).into();
136
137                        if self.interval.intersects(alignment_interval) {
138                            return Some(Ok(r));
139                        }
140                    }
141                }
142                None => match self.read_next_container() {
143                    Some(Ok(())) => {}
144                    Some(Err(e)) => return Some(Err(e)),
145                    None => return None,
146                },
147            }
148        }
149    }
150}
noodles_cram/io/reader/query.rs

noodles_cram/io/reader/
query.rs
