twitcher 0.1.8

Find template switch mutations in genomic data
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, // forces loading of the first region when reading the first record
        })
    }

    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
    }
}