use std::{collections::VecDeque, path::Path};
use rust_htslib::bam::{FetchDefinition, IndexedReader, Read, Reader, Record};
use tracing::error;
use crate::common::coords::GenomeRegion;
pub struct ReaderEntireFile {
inner: Reader,
done: bool,
}
pub struct ReaderWithRegions {
inner: IndexedReader,
regions: VecDeque<GenomeRegion>,
region_done: bool,
}
pub enum BAMReader {
EntireFile(ReaderEntireFile),
Regions(ReaderWithRegions),
}
impl BAMReader {
pub fn entire_file(inner: Reader) -> Self {
Self::EntireFile(ReaderEntireFile { inner, done: false })
}
pub fn with_regions(inner: IndexedReader, regions: Vec<GenomeRegion>) -> Self {
Self::Regions(ReaderWithRegions {
inner,
regions: regions.into(),
region_done: true, })
}
pub fn set_reference<P: AsRef<Path>>(
&mut self,
path: P,
) -> Result<(), rust_htslib::errors::Error> {
match self {
BAMReader::EntireFile(reader_entire_file) => {
reader_entire_file.inner.set_reference(path)
}
BAMReader::Regions(reader_with_regions) => {
reader_with_regions.inner.set_reference(path)
}
}
}
pub fn done(&mut self) -> bool {
match self {
BAMReader::EntireFile(reader) => reader.done,
BAMReader::Regions(reader_with_regions) => {
reader_with_regions.region_done && reader_with_regions.regions.is_empty()
}
}
}
}
impl Read for BAMReader {
fn read(
&mut self,
record: &mut rust_htslib::bam::record::Record,
) -> Option<rust_htslib::tpool::Result<()>> {
match self {
BAMReader::EntireFile(reader_entire_file) => reader_entire_file.read(record),
BAMReader::Regions(reader_with_regions) => reader_with_regions.read(record),
}
}
fn records(&mut self) -> rust_htslib::bam::Records<'_, Self> {
unimplemented!()
}
fn rc_records(&mut self) -> rust_htslib::bam::RcRecords<'_, Self> {
unimplemented!()
}
fn pileup(&mut self) -> rust_htslib::bam::pileup::Pileups<'_, Self> {
unimplemented!()
}
fn htsfile(&self) -> *mut rust_htslib::htslib::htsFile {
unimplemented!()
}
fn header(&self) -> &rust_htslib::bam::HeaderView {
match self {
BAMReader::EntireFile(reader_entire_file) => reader_entire_file.inner.header(),
BAMReader::Regions(reader_with_regions) => reader_with_regions.inner.header(),
}
}
fn set_thread_pool(
&mut self,
_tpool: &rust_htslib::tpool::ThreadPool,
) -> rust_htslib::tpool::Result<()> {
unimplemented!()
}
}
impl ReaderEntireFile {
fn read(&mut self, target: &mut Record) -> Option<Result<(), rust_htslib::errors::Error>> {
let ret = self.inner.read(target);
if ret.is_none() {
self.done = true;
}
ret
}
}
impl ReaderWithRegions {
fn read(&mut self, target: &mut Record) -> Option<Result<(), rust_htslib::errors::Error>> {
if self.region_done {
if let Some(next) = self.regions.pop_front() {
let fd = match FetchDefinition::try_from(&next) {
Ok(fd) => fd,
Err(e) => {
error!("Skipping region: {e}");
return Some(Err(rust_htslib::tpool::Error::Fetch));
}
};
match self.inner.fetch(fd) {
Ok(()) => self.region_done = false,
Err(e) => return Some(Err(e)),
}
} else {
return None;
}
}
let ret = self.inner.read(target);
if ret.is_none() {
self.region_done = true;
}
ret
}
}