twitcher 0.1.8

Find template switch mutations in genomic data
use std::collections::VecDeque;

use rust_htslib::bcf::{IndexedReader, Read, Reader, Record};

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 VCFReader {
    EntireFile(ReaderEntireFile),
    Regions(ReaderWithRegions),
}

impl VCFReader {
    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 done(&mut self) -> bool {
        match self {
            VCFReader::EntireFile(reader) => reader.done,
            VCFReader::Regions(reader_with_regions) => {
                reader_with_regions.region_done && reader_with_regions.regions.is_empty()
            }
        }
    }
}

impl Read for VCFReader {
    fn read(
        &mut self,
        record: &mut rust_htslib::bcf::record::Record,
    ) -> Option<Result<(), rust_htslib::errors::Error>> {
        match self {
            VCFReader::EntireFile(reader) => reader.read(record),
            VCFReader::Regions(reader) => reader.read(record),
        }
    }

    fn records(&mut self) -> rust_htslib::bcf::Records<'_, Self> {
        unimplemented!()
    }

    fn header(&self) -> &rust_htslib::bcf::header::HeaderView {
        match self {
            VCFReader::EntireFile(reader) => reader.inner.header(),
            VCFReader::Regions(reader_with_regions) => reader_with_regions.inner.header(),
        }
    }

    fn empty_record(&self) -> Record {
        match self {
            VCFReader::EntireFile(reader) => reader.inner.empty_record(),
            VCFReader::Regions(reader_with_regions) => reader_with_regions.inner.empty_record(),
        }
    }

    fn set_threads(&mut self, n_threads: usize) -> rust_htslib::tpool::Result<()> {
        match self {
            VCFReader::EntireFile(reader) => reader.inner.set_threads(n_threads),
            VCFReader::Regions(reader_with_regions) => {
                reader_with_regions.inner.set_threads(n_threads)
            }
        }
    }
}

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 rid = match self.inner.header().name2rid(next.contig().as_ref()) {
                    Ok(rid) => rid,
                    Err(e) => return Some(Err(e)),
                };
                match self.inner.fetch(
                    rid,
                    next.start().position_0() as u64,
                    next.end_excl().map(|end| end.position_0() as u64),
                ) {
                    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
    }
}