gwseq_io/bam/

reader.rs

//! The BAM reader.
//!
//! Four read paths over the same machinery: the
//! index turns a locus into chunks, BGZF turns a chunk into records, and the
//! filter decides which of them the caller asked for.
//!
//! # Why a locus is split by chunk, not by sub-region
//!
//! Cutting a window into sub-regions costs far more than it saves: the index
//! reaches an alignment through runs of compressed blocks rather than through a
//! region, so neighbouring sub-regions land in the same blocks and each pays to
//! decompress them — three times the work at 24 pieces, measured on a 10 M
//! read file. The chunks of a locus are disjoint, so sharing *those* out
//! splits the work where the format divides it, and only the block a run
//! boundary falls in is read twice.

use std::sync::Arc;

use crate::bam::bai::{BamIndex, MAX_MERGE_SPAN};
use crate::bam::bgzf::Chunk;
use crate::bam::header::SamHeader;
use crate::bam::record::{decode_block, BamRecord, EntryFilter, RecordFilter};
use crate::error::{Error, Result};
use crate::genomic::{ChrMap, Locs};
use crate::parallel::Executor;
use crate::progress::{ProgressFn, ProgressTracker};
use crate::source::ByteSource;

/// Chunks a cursor keeps decompressed, so a walk that revisits one does not
/// inflate it again.
const CHUNK_CACHE_SIZE: usize = 4;

/// And the bytes it keeps between them, since a chunk has no size limit worth
/// the name — `chunks` merges neighbours up to `MAX_MERGE_SPAN` compressed. A
/// walk holds one cursor per worker for its whole life, so this is what bounds
/// what a walk costs to have open.
const CURSOR_CACHE_BYTES: usize = 8 << 20;

/// A run of chunks under this many compressed bytes costs more to hand a thread
/// than to read.
const MIN_RUN_SIZE: u64 = 64 * 1024;

#[derive(Debug)]
struct Inner {
    source: Arc<dyn ByteSource>,
    executor: Executor,
    index: Option<BamIndex>,
}

pub struct BamReader {
    inner: Option<Inner>,
    path: String,
    index_path: String,
    header: SamHeader,
    chr_map: ChrMap,
    /// Reference names by the index the file gives them, shared with every
    /// record so turning an index back into a name costs a clone of an `Arc`.
    chr_names: Arc<Vec<String>>,
    /// Why the index is absent, when it is. Empty when it loaded — and empty as
    /// well when the file simply has none, which is a distinction the API keeps.
    index_error: String,
    /// Whether an index loaded at open. Kept here rather than read off the
    /// index itself, which `close()` drops, so that `is_indexed` still answers
    /// afterwards: what kind of file was opened does not stop being true
    /// because the handle was released.
    indexed: bool,
}

impl std::fmt::Debug for BamReader {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("BamReader")
            .field("path", &self.path)
            .field("references", &self.chr_map.len())
            .field("indexed", &self.is_indexed())
            .field("closed", &self.is_closed())
            .finish()
    }
}

/// One worker's state for a walk: a decompressed-chunk cache it carries from
/// locus to locus.
///
/// The cache is the caller's, not the call's. A walk that reads locus after
/// locus hands the same cursors back every time, which is what lets a cursor's
/// blocks serve more than the read that filled them — and rebuilding them per
/// call is what makes reading a few loci at a time cost more than reading many.
/// See [`Cursors`], which is what a walk holds.
#[derive(Debug, Default)]
pub struct Cursor {
    cache: Vec<(Chunk, bytes::Bytes)>,
    next: usize,
    /// Decompressed bytes resident, so the budget below is a subtraction rather
    /// than a walk of the cache.
    bytes: usize,
}

/// One [`Cursor`] per worker, carried across the steps of a walk.
///
/// A run takes the cursor its batch index names, so two workers never want the
/// same one and the lock is uncontended — a `Mutex` each rather than one shared
/// is what keeps it that way. Living on the walk rather than on the call is the
/// point: it is what makes `sort_locations` mean anything, since neighbouring
/// loci touch the same BGZF blocks and a cursor rebuilt per call throws them
/// away before the next locus can ask.
#[derive(Debug, Default)]
pub struct Cursors(Vec<parking_lot::Mutex<Cursor>>);

impl Cursors {
    pub fn new(count: usize) -> Self {
        Self(
            (0..count.max(1))
                .map(|_| parking_lot::Mutex::new(Cursor::default()))
                .collect(),
        )
    }

    /// The cursor for a worker, by its batch index.
    ///
    /// Wrapped rather than bounded: `chunk_runs` never makes more runs than
    /// there are workers, so this is exact in practice and safe if that ever
    /// stops being true.
    fn get(&self, index: usize) -> parking_lot::MutexGuard<'_, Cursor> {
        self.0[index % self.0.len()].lock()
    }
}

impl BamReader {
    pub fn open(
        path: &str,
        index_path: Option<&str>,
        parallel: i64,
        block_size: Option<u64>,
        max_blocks: Option<usize>,
    ) -> Result<Self> {
        let source = crate::source::open(path, block_size, max_blocks)?;
        Self::from_source(source, path, index_path, parallel, block_size, max_blocks)
    }

    pub(crate) fn from_source(
        source: Arc<dyn ByteSource>,
        path: &str,
        index_path: Option<&str>,
        parallel: i64,
        block_size: Option<u64>,
        max_blocks: Option<usize>,
    ) -> Result<Self> {
        // Before the header, so a truncated file is refused rather than opened
        // and read short.
        super::bgzf::check_eof(source.as_ref())?;
        let (header, chr_map) = super::header::read(source.as_ref())?;

        let mut names = vec![String::new(); chr_map.iter().map(|e| e.index + 1).max().unwrap_or(0)];
        for entry in chr_map.iter() {
            names[entry.index] = entry.id.clone();
        }

        let index_path = index_path
            .map(str::to_string)
            .unwrap_or_else(|| format!("{path}.bai"));
        // An index is optional, so a missing one is no error — but a corrupt one
        // is worth reporting separately, and used to reach the caller as the
        // same "not indexed". A local index that is simply absent is checked for
        // before opening, since opening fails the same way a corrupt one does.
        let (index, index_error) =
            if !crate::source::is_url(&index_path) && !std::path::Path::new(&index_path).exists() {
                (None, String::new())
            } else {
                match crate::source::open(&index_path, block_size, max_blocks)
                    .and_then(|s| BamIndex::read(s.as_ref()))
                {
                    Ok(index) => (Some(index), String::new()),
                    Err(e) => (None, e.to_string()),
                }
            };

        Ok(Self {
            indexed: index.is_some(),
            inner: Some(Inner {
                source,
                executor: Executor::new(parallel)?,
                index,
            }),
            path: path.to_string(),
            index_path,
            header,
            chr_map,
            chr_names: Arc::new(names),
            index_error,
        })
    }

    pub fn header(&self) -> &SamHeader {
        &self.header
    }
    pub fn chr_sizes(&self) -> &ChrMap {
        &self.chr_map
    }
    pub fn index_error(&self) -> &str {
        &self.index_error
    }
    pub fn is_indexed(&self) -> bool {
        self.indexed
    }
    pub fn is_closed(&self) -> bool {
        self.inner.is_none()
    }
    pub fn path(&self) -> &str {
        &self.path
    }
    pub fn parallel(&self) -> usize {
        self.inner.as_ref().map_or(0, |i| i.executor.parallel())
    }

    pub fn close(&mut self) {
        if let Some(inner) = self.inner.take() {
            inner.source.close();
        }
    }

    fn inner(&self) -> Result<&Inner> {
        self.inner.as_ref().ok_or_else(|| Error::Closed {
            path: self.path.clone(),
        })
    }

    /// The index, or why there is none.
    ///
    /// A missing index and a corrupt one are different failures, and used to
    /// reach the caller as the same "not indexed".
    fn index<'a>(&self, inner: &'a Inner) -> Result<&'a BamIndex> {
        inner.index.as_ref().ok_or_else(|| {
            Error::invalid(if self.index_error.is_empty() {
                format!("bam file is not indexed ({} not found)", self.index_path)
            } else {
                format!(
                    "bam index {} could not be read: {}",
                    self.index_path, self.index_error
                )
            })
        })
    }

    /// One list per locus, in the order the loci were given.
    ///
    /// Each locus gets its own full list, so two overlapping loci both report
    /// the alignments they share.
    pub fn read_entries(&self, req: &EntriesRequest) -> Result<Vec<Vec<BamRecord>>> {
        let inner = self.inner()?;
        let resolved = self.resolve(&req.locs)?;
        let coverage = resolved.iter().map(|(_, s, e)| (e - s).max(0) as u64).sum();
        let tracker = ProgressTracker::with_callback(coverage, req.progress.clone());
        let out = self.read_loci(inner, &resolved, req, &tracker)?;
        tracker.done_report();
        Ok(out)
    }

    /// Every alignment on the named references, in reference then coordinate
    /// order.
    ///
    /// Unplaced alignments are left out: the index reaches an alignment only
    /// through the reference it sits on.
    pub fn read_all_entries(&self, req: &EntriesRequest) -> Result<Vec<BamRecord>> {
        let locs = Locs::whole_chromosomes(&self.chr_map, &req.locs.chr_ids)?;
        let whole = EntriesRequest {
            locs,
            ..req.clone()
        };
        Ok(self.read_entries(&whole)?.into_iter().flatten().collect())
    }

    /// Per-locus, lazily.
    pub fn iter_entries(&self, req: &EntriesRequest) -> Result<LocusEntries<'_>> {
        LocusEntries::plan(self, req)
    }

    /// Successive windows over whole references.
    ///
    /// A window never spans two references, and an alignment reaching over a
    /// boundary is reported by the window it starts in.
    pub fn iter_all_entries(&self, req: &EntriesRequest, window: i64) -> Result<WindowEntries<'_>> {
        WindowEntries::plan(self, req, window)
    }

    /// Resolve each locus to (reference index, start, end).
    fn resolve(&self, locs: &Locs) -> Result<Vec<(usize, i64, i64)>> {
        (0..locs.len())
            .map(|i| {
                let entry = self.chr_map.resolve(&locs.chr_ids[i])?;
                Ok((entry.index, locs.starts[i], locs.ends[i]))
            })
            .collect()
    }

    /// Read a stretch of loci, spread over the executor's workers.
    ///
    /// Split into runs of consecutive loci rather than dealt round-robin, so
    /// each worker walks the request forward and the blocks it decompressed for
    /// one locus serve the ones after it.
    fn read_loci(
        &self,
        inner: &Inner,
        loci: &[(usize, i64, i64)],
        req: &EntriesRequest,
        tracker: &ProgressTracker,
    ) -> Result<Vec<Vec<BamRecord>>> {
        if loci.is_empty() {
            return Ok(Vec::new());
        }
        let index = self.index(inner)?;
        let workers = inner.executor.parallel().min(loci.len()).max(1);
        let per_worker = loci.len().div_ceil(workers);
        let batches: Vec<(usize, usize)> = (0..workers)
            .map(|w| (w * per_worker, ((w + 1) * per_worker).min(loci.len())))
            .filter(|(from, to)| from < to)
            .collect();

        let lists = inner.executor.map_batches(&batches, |_, (from, to)| {
            let mut cursor = Cursor::default();
            let mut out = Vec::with_capacity(to - from);
            for (chr, start, end) in &loci[*from..*to] {
                let chunks = index.chunks(*chr, *start, *end, Some(MAX_MERGE_SPAN))?;
                let mut records = Vec::new();
                self.read_chunks(
                    inner,
                    &mut cursor,
                    &chunks,
                    0..chunks.len(),
                    (*chr, *start, *end),
                    req,
                    &mut records,
                )?;
                out.push(records);
                tracker.add((end - start).max(0) as u64);
            }
            Ok(out)
        })?;
        Ok(lists.into_iter().flatten().collect())
    }

    /// Read one locus, its chunks shared out over the executor's workers.
    ///
    /// The runs come back in file order, which for a coordinate-sorted file is
    /// coordinate order, so laying them end to end gives exactly what reading
    /// the locus on one thread gives.
    fn read_locus_split(
        &self,
        inner: &Inner,
        locus: (usize, i64, i64),
        req: &EntriesRequest,
        cursors: &Cursors,
    ) -> Result<Vec<BamRecord>> {
        let index = self.index(inner)?;
        let (chr, start, end) = locus;
        let chunks = index.chunks(chr, start, end, Some(MAX_MERGE_SPAN))?;
        let runs = chunk_runs(&chunks, inner.executor.parallel());
        if runs.is_empty() {
            return Ok(Vec::new());
        }
        if runs.len() == 1 {
            // Read where it stands rather than through the pool, which for a
            // window holding one run of chunks is every window of a sparse file.
            // This is also the path a scattered request spends its time on, and
            // so the one whose cursor has to survive the call: cursor 0 carries
            // the blocks from one locus to the next.
            let mut cursor = cursors.get(0);
            let mut out = Vec::new();
            self.read_chunks(
                inner,
                &mut cursor,
                &chunks,
                runs[0].clone(),
                locus,
                req,
                &mut out,
            )?;
            return Ok(out);
        }
        let lists = inner.executor.map_batches(&runs, |index, run| {
            let mut cursor = cursors.get(index);
            let mut out = Vec::new();
            self.read_chunks(
                inner,
                &mut cursor,
                &chunks,
                run.clone(),
                locus,
                req,
                &mut out,
            )?;
            Ok(out)
        })?;
        Ok(lists.into_iter().flatten().collect())
    }

    /// Append the alignments of a run of chunks that overlap the locus.
    ///
    /// The region still reaches the decoder as the filter, so a run holds only
    /// the alignments of its own chunks that the locus overlaps. The chunks of a
    /// locus are disjoint, so a run is disjoint from every other and an
    /// alignment is read by exactly one of them.
    #[allow(clippy::too_many_arguments)]
    fn read_chunks(
        &self,
        inner: &Inner,
        cursor: &mut Cursor,
        chunks: &[Chunk],
        run: std::ops::Range<usize>,
        locus: (usize, i64, i64),
        req: &EntriesRequest,
        out: &mut Vec<BamRecord>,
    ) -> Result<()> {
        let (chr, start, end) = locus;
        let filter = EntryFilter {
            chr_index: Some(chr as i32),
            start,
            end: Some(end),
            standard_flags: req.filter.enabled,
        };
        for chunk in &chunks[run] {
            let data = cursor.get_or_read(inner.source.as_ref(), *chunk, &self.path)?;
            out.extend(decode_block(
                &data,
                req.parse_tags,
                &filter,
                &self.chr_names,
                &self.path,
            )?);
        }
        Ok(())
    }
}

impl Cursor {
    /// The decompressed bytes of a chunk, from the cache or from the file.
    ///
    /// Keyed by the *pair* of virtual offsets, not by the start alone: two loci
    /// touch different bins, so their chunk runs merge differently and two
    /// chunks may start together and end apart.
    fn get_or_read(
        &mut self,
        source: &dyn ByteSource,
        chunk: Chunk,
        path: &str,
    ) -> Result<bytes::Bytes> {
        if let Some((_, data)) = self.cache.iter().find(|(c, _)| *c == chunk) {
            return Ok(data.clone());
        }
        let data = super::bgzf::decompress_chunk(source, chunk, path)?;
        // A count alone is not a bound: `chunks` merges neighbours up to
        // `MAX_MERGE_SPAN` of *compressed* bytes, so four entries can be a lot
        // of memory — and a cursor now lives as long as the walk that holds it
        // rather than as long as one call. A chunk too large to share is still
        // returned; it is simply not kept.
        if data.len() <= CURSOR_CACHE_BYTES {
            if self.cache.len() < CHUNK_CACHE_SIZE {
                self.bytes += data.len();
                self.cache.push((chunk, data.clone()));
            } else {
                self.bytes -= self.cache[self.next].1.len();
                self.bytes += data.len();
                self.cache[self.next] = (chunk, data.clone());
                self.next = (self.next + 1) % CHUNK_CACHE_SIZE;
            }
            while self.bytes > CURSOR_CACHE_BYTES && self.cache.len() > 1 {
                // Oldest first, which with a ring of `next` is the entry it is
                // about to overwrite anyway.
                let oldest = self.next % self.cache.len();
                self.bytes -= self.cache[oldest].1.len();
                self.cache.remove(oldest);
                self.next = oldest.min(self.cache.len().saturating_sub(1));
            }
        }
        Ok(data)
    }
}

/// Split a locus's chunks into the consecutive runs the threads share out.
///
/// Fewer runs than `workers` where the chunks hold too little to be worth that
/// many, and a single run where they hold less than one. A chunk is never split,
/// so a locus reaching one chunk is read on one thread however many there are.
/// This is the loci batcher's rule, over compressed bytes rather than base
/// pairs.
fn chunk_runs(chunks: &[Chunk], workers: usize) -> Vec<std::ops::Range<usize>> {
    if chunks.is_empty() || workers < 1 {
        return Vec::new();
    }
    let total: u64 = chunks.iter().map(|c| c.compressed_size()).sum();
    let wanted = (workers as u64).min(total / MIN_RUN_SIZE).max(1);
    // Rounded up, so filling every run to it leaves no run in excess.
    let per_run = total.div_ceil(wanted).max(1);

    #[allow(clippy::single_range_in_vec_init)]
    let mut runs = vec![0usize..0];
    let mut size = 0u64;
    for (i, chunk) in chunks.iter().enumerate().take(chunks.len() - 1) {
        size += chunk.compressed_size();
        if size < per_run {
            continue;
        }
        // Once the count is reached everything left joins the last run, which is
        // where the rounding already expects the remainder to go.
        if runs.len() as u64 >= wanted {
            continue;
        }
        runs.last_mut().expect("pushed one above").end = i + 1;
        runs.push(i + 1..i + 1);
        size = 0;
    }
    runs.last_mut().expect("pushed one above").end = chunks.len();
    runs
}

#[derive(Clone)]
pub struct EntriesRequest {
    pub locs: Locs,
    pub filter: RecordFilter,
    /// Optional fields are only decoded when read, so keeping them costs one
    /// small slice per alignment.
    pub parse_tags: bool,
    /// Read the loci in reference and position order, and report them in that
    /// order. Loci close together then share the blocks a read decompressed —
    /// 4 to 7 times faster on a scattered request.
    pub sort_locations: bool,
    pub progress: Option<ProgressFn>,
}

impl std::fmt::Debug for EntriesRequest {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("EntriesRequest")
            .field("loci", &self.locs.len())
            .field("filter", &self.filter.enabled)
            .field("parse_tags", &self.parse_tags)
            .field("sort_locations", &self.sort_locations)
            .finish()
    }
}

impl EntriesRequest {
    pub fn new(locs: Locs) -> Self {
        Self {
            locs,
            filter: RecordFilter::default(),
            parse_tags: true,
            sort_locations: false,
            progress: None,
        }
    }
    pub fn filter(mut self, enabled: bool) -> Self {
        self.filter.enabled = enabled;
        self
    }
    pub fn parse_tags(mut self, v: bool) -> Self {
        self.parse_tags = v;
        self
    }
    pub fn sort_locations(mut self, v: bool) -> Self {
        self.sort_locations = v;
        self
    }
    pub fn progress(mut self, f: ProgressFn) -> Self {
        self.progress = Some(f);
        self
    }
}

/// A walk over loci, one list at a time.
///
/// Owns its plan and takes the reader on each step, so a `#[pyclass]` can drive
/// it — the same shape as the bbi walks.
/// The plan a [`LocusWalk`] walks: everything decided when it was made, and
/// nothing that changes as it runs.
///
/// Split out and shared behind an [`Arc`] so that restarting a walk costs a
/// refcount bump rather than a copy of every locus. A per-locus walk over a
/// hundred thousand requested regions holds a hundred thousand of each of
/// these, and a caller looping twice should not pay for them twice.
struct WalkPlan {
    /// The loci in the order they will be read.
    loci: Vec<(usize, i64, i64)>,
    /// Request index of each, so a sorted walk can say which locus a list is.
    order: Vec<usize>,
    /// Lowest start an alignment may have to be reported by its locus, or
    /// `None` to report every alignment the locus overlaps.
    ///
    /// What makes a whole-file walk group alignments by the window they *start*
    /// in: an alignment reaching over a boundary overlaps two windows and would
    /// otherwise be handed over by both.
    min_starts: Vec<Option<i64>>,
    request: EntriesRequest,
    /// Bases the whole walk covers, for a restart's fresh tracker.
    coverage: u64,
}

pub struct LocusWalk {
    plan: Arc<WalkPlan>,
    next: usize,
    tracker: Arc<ProgressTracker>,
    /// Carried across steps, which is what makes `sort_locations` worth asking
    /// for: neighbouring loci share BGZF blocks, and a cursor built per call
    /// has thrown them away before the next locus arrives.
    cursors: Cursors,
}

impl std::fmt::Debug for LocusWalk {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("LocusWalk")
            .field("loci", &self.plan.loci.len())
            .field("next", &self.next)
            .finish()
    }
}

impl LocusWalk {
    /// A per-locus walk: every alignment its locus overlaps, boundaries and
    /// all.
    pub fn plan(reader: &BamReader, req: &EntriesRequest) -> Result<Self> {
        Self::plan_with(reader, req, false)
    }

    /// `from_locus_start` reports only alignments starting at or after each
    /// locus's own start, which is what tiling windows need and what a
    /// per-locus walk must not do.
    fn plan_with(reader: &BamReader, req: &EntriesRequest, from_locus_start: bool) -> Result<Self> {
        let inner = reader.inner()?;
        reader.index(inner)?;
        let resolved = reader.resolve(&req.locs)?;
        let mut order: Vec<usize> = (0..resolved.len()).collect();
        if req.sort_locations {
            order.sort_by_key(|i| resolved[*i]);
        }
        let loci: Vec<(usize, i64, i64)> = order.iter().map(|i| resolved[*i]).collect();
        // Progress is counted in base pairs, as every other reader counts it.
        let coverage = loci.iter().map(|(_, s, e)| (e - s).max(0) as u64).sum();
        let min_starts = if from_locus_start {
            loci.iter().map(|(_, start, _)| Some(*start)).collect()
        } else {
            vec![None; loci.len()]
        };
        Ok(Self {
            tracker: Arc::new(ProgressTracker::with_callback(
                coverage,
                req.progress.clone(),
            )),
            plan: Arc::new(WalkPlan {
                min_starts,
                loci,
                order,
                request: req.clone(),
                coverage,
            }),
            next: 0,
            cursors: Cursors::new(reader.parallel()),
        })
    }

    /// The same walk, back at its first locus.
    ///
    /// Shares the plan, so this is one refcount bump however many loci it
    /// holds. The progress tracker is fresh: a second pass reports its own
    /// progress from zero, which is what a caller watching it expects.
    pub fn restarted(&self) -> Self {
        Self {
            plan: self.plan.clone(),
            next: 0,
            tracker: Arc::new(ProgressTracker::with_callback(
                self.plan.coverage,
                self.plan.request.progress.clone(),
            )),
            // Fresh, not shared: a second pass reads the file again, as the
            // docs say, and two cursors of held blocks is not what a caller
            // asking for a second `for` loop is asking to pay for.
            cursors: Cursors::new(self.cursors.0.len()),
        }
    }

    /// The same plan, for a walk that tiles whole references.
    pub fn plan_windows(reader: &BamReader, req: &EntriesRequest, span: i64) -> Result<Self> {
        if span < 1 {
            return Err(Error::invalid(format!(
                "span must be positive (got {span})"
            )));
        }
        let locs = window_locs(&reader.chr_map, &req.locs.chr_ids, span)?;
        let windowed = EntriesRequest {
            locs,
            sort_locations: false,
            ..req.clone()
        };
        Self::plan_with(reader, &windowed, true)
    }

    pub fn len(&self) -> usize {
        self.plan.loci.len()
    }
    pub fn is_empty(&self) -> bool {
        self.plan.loci.is_empty()
    }
    /// Request index of each yielded list: the nth list belongs to locus
    /// `order()[n]`.
    pub fn order(&self) -> &[usize] {
        &self.plan.order
    }

    pub fn next_window(&mut self, reader: &BamReader) -> Option<Result<Vec<BamRecord>>> {
        if self.next >= self.plan.loci.len() {
            self.tracker.done_report();
            return None;
        }
        let index = self.next;
        let locus = self.plan.loci[index];
        let outcome = reader.inner().and_then(|inner| {
            reader.read_locus_split(inner, locus, &self.plan.request, &self.cursors)
        });
        match outcome {
            // A failed step is not a step at all: the walk stays where it was,
            // so the call after it comes back here and raises the same error
            // rather than reading past a locus nothing was read for. That is
            // what makes closing a reader under a walk raise cleanly, every
            // time, instead of once.
            Err(e) => Some(Err(e)),
            Ok(mut records) => {
                self.next += 1;
                if let Some(min_start) = self.plan.min_starts[index] {
                    records.retain(|r| r.start() >= min_start);
                }
                let (_, start, end) = locus;
                self.tracker.add((end - start).max(0) as u64);
                Some(Ok(records))
            }
        }
    }
}

/// [`LocusWalk`] as a plain [`Iterator`], for Rust callers.
#[derive(Debug)]
pub struct LocusEntries<'a> {
    reader: &'a BamReader,
    walk: LocusWalk,
}

impl<'a> LocusEntries<'a> {
    fn plan(reader: &'a BamReader, req: &EntriesRequest) -> Result<Self> {
        Ok(Self {
            reader,
            walk: LocusWalk::plan(reader, req)?,
        })
    }
    pub fn len(&self) -> usize {
        self.walk.len()
    }
    pub fn is_empty(&self) -> bool {
        self.walk.is_empty()
    }
    pub fn order(&self) -> &[usize] {
        self.walk.order()
    }
}

impl Iterator for LocusEntries<'_> {
    type Item = Result<Vec<BamRecord>>;
    fn next(&mut self) -> Option<Self::Item> {
        self.walk.next_window(self.reader)
    }
}

/// [`LocusWalk`] over windows tiling whole references.
#[derive(Debug)]
pub struct WindowEntries<'a> {
    reader: &'a BamReader,
    walk: LocusWalk,
}

/// The windows a whole-file walk tiles the references with.
pub fn window_locs(map: &ChrMap, chr_ids: &[String], span: i64) -> Result<Locs> {
    let mut ids = Vec::new();
    let mut starts = Vec::new();
    let mut ends = Vec::new();
    for chr in map.select(chr_ids)? {
        let mut start = 0;
        while start < chr.size {
            ids.push(chr.id.clone());
            starts.push(start);
            ends.push((start + span).min(chr.size));
            start += span;
        }
    }
    Locs::spans(&ids, &starts, &ends)
}

impl<'a> WindowEntries<'a> {
    fn plan(reader: &'a BamReader, req: &EntriesRequest, span: i64) -> Result<Self> {
        if span < 1 {
            return Err(Error::invalid(format!(
                "span must be positive (got {span})"
            )));
        }
        let locs = window_locs(&reader.chr_map, &req.locs.chr_ids, span)?;
        let windowed = EntriesRequest {
            locs,
            sort_locations: false,
            ..req.clone()
        };
        let walk = LocusWalk::plan_with(reader, &windowed, true)?;
        Ok(Self { reader, walk })
    }
    pub fn len(&self) -> usize {
        self.walk.len()
    }
    pub fn is_empty(&self) -> bool {
        self.walk.is_empty()
    }
}

impl Iterator for WindowEntries<'_> {
    type Item = Result<Vec<BamRecord>>;
    fn next(&mut self) -> Option<Self::Item> {
        self.walk.next_window(self.reader)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::bam::bgzf::VirtualOffset;

    fn chunk(a: u64, b: u64) -> Chunk {
        Chunk {
            begin: VirtualOffset::new(a, 0),
            end: VirtualOffset::new(b, 0),
        }
    }

    #[test]
    fn no_chunks_means_no_runs() {
        assert!(chunk_runs(&[], 4).is_empty());
        assert!(chunk_runs(&[chunk(0, 100)], 0).is_empty());
    }

    #[test]
    fn chunks_holding_too_little_are_read_as_one_run() {
        // Well under MIN_RUN_SIZE: splitting would cost more than it saves.
        let chunks = [chunk(0, 100), chunk(100, 200), chunk(200, 300)];
        #[allow(clippy::single_range_in_vec_init)]
        let one_run = [0..3];
        assert_eq!(chunk_runs(&chunks, 8), one_run);
    }

    #[test]
    fn a_big_locus_splits_into_at_most_one_run_per_worker() {
        let big = MIN_RUN_SIZE * 4;
        let chunks: Vec<Chunk> = (0..8).map(|i| chunk(i * big, (i + 1) * big)).collect();
        let runs = chunk_runs(&chunks, 4);
        assert_eq!(runs.len(), 4);
        // Contiguous, covering every chunk exactly once.
        assert_eq!(runs[0].start, 0);
        assert_eq!(runs.last().unwrap().end, chunks.len());
        for pair in runs.windows(2) {
            assert_eq!(pair[0].end, pair[1].start);
        }
    }

    #[test]
    fn a_run_is_never_empty_and_a_chunk_is_never_split() {
        let big = MIN_RUN_SIZE * 100;
        let chunks = [chunk(0, big), chunk(big, big * 2)];
        // Far more workers than chunks: two chunks cannot make eight runs.
        let runs = chunk_runs(&chunks, 8);
        assert!(runs.len() <= chunks.len());
        assert!(runs.iter().all(|r| r.start < r.end));
    }
}

#[cfg(test)]
mod cursor_tests {
    use super::*;
    use crate::bam::bgzf::VirtualOffset;
    use crate::source::testing::MemorySource;
    use std::sync::atomic::{AtomicUsize, Ordering};

    /// A source that counts the reads reaching it, so a cache hit is visible as
    /// a read that never happened.
    #[derive(Debug)]
    struct CountingSource {
        inner: MemorySource,
        reads: AtomicUsize,
    }

    impl ByteSource for CountingSource {
        fn path(&self) -> &str {
            self.inner.path()
        }
        fn len(&self) -> Result<u64> {
            self.inner.len()
        }
        fn read_at(&self, offset: u64, len: usize) -> Result<bytes::Bytes> {
            self.reads.fetch_add(1, Ordering::SeqCst);
            self.inner.read_at(offset, len)
        }
    }

    /// One BGZF block wrapping `payload`, the way `crate::fuzz` builds them.
    fn bgzf_block(payload: &[u8]) -> Vec<u8> {
        use std::io::Write as _;
        let mut encoder =
            flate2::write::DeflateEncoder::new(Vec::new(), flate2::Compression::new(6));
        encoder.write_all(payload).expect("deflate to a Vec");
        let deflated = encoder.finish().expect("deflate to a Vec");
        let total = 18 + deflated.len() + 8;
        let mut out = Vec::with_capacity(total);
        out.extend_from_slice(&[0x1f, 0x8b, 8, 4, 0, 0, 0, 0, 0, 0xff]);
        out.extend_from_slice(&6u16.to_le_bytes());
        out.extend_from_slice(b"BC");
        out.extend_from_slice(&2u16.to_le_bytes());
        out.extend_from_slice(&((total - 1) as u16).to_le_bytes());
        out.extend_from_slice(&deflated);
        let mut crc = flate2::Crc::new();
        crc.update(payload);
        out.extend_from_slice(&crc.sum().to_le_bytes());
        out.extend_from_slice(&(payload.len() as u32).to_le_bytes());
        out
    }

    /// `count` blocks of `size` bytes each, and where each one begins.
    fn blocks(count: usize, size: usize) -> (CountingSource, Vec<u64>) {
        let mut bytes = Vec::new();
        let mut offsets = Vec::new();
        for i in 0..count {
            offsets.push(bytes.len() as u64);
            bytes.extend_from_slice(&bgzf_block(&vec![(i % 251) as u8; size]));
        }
        offsets.push(bytes.len() as u64);
        (
            CountingSource {
                inner: MemorySource::new(bytes),
                reads: AtomicUsize::new(0),
            },
            offsets,
        )
    }

    fn chunk(from: u64, to: u64) -> Chunk {
        Chunk {
            begin: VirtualOffset::new(from, 0),
            end: VirtualOffset::new(to, 0),
        }
    }

    /// What carrying a cursor across a walk's steps buys, without a stopwatch:
    /// a chunk asked for twice is inflated once. The end-to-end effect of that
    /// — `sort_locations` going from doing nothing to a 3.5x to 20x speedup,
    /// depending on how close the loci sit — is a measurement, recorded in
    /// `local/plan.md`; this is the mechanism it rests on.
    #[test]
    fn a_chunk_a_cursor_has_already_read_is_not_read_again() {
        let (source, offsets) = blocks(4, 4096);
        let mut cursor = Cursor::default();
        let first = chunk(offsets[0], offsets[1]);

        let a = cursor.get_or_read(&source, first, "x.bam").unwrap();
        assert_eq!(source.reads.load(Ordering::SeqCst), 1);
        let b = cursor.get_or_read(&source, first, "x.bam").unwrap();
        assert_eq!(a, b);
        assert_eq!(
            source.reads.load(Ordering::SeqCst),
            1,
            "the second read of one chunk reached the file"
        );

        // A different chunk is a different entry, and the first is still there.
        let second = chunk(offsets[1], offsets[2]);
        cursor.get_or_read(&source, second, "x.bam").unwrap();
        let before = source.reads.load(Ordering::SeqCst);
        cursor.get_or_read(&source, first, "x.bam").unwrap();
        assert_eq!(source.reads.load(Ordering::SeqCst), before);

        // A fresh cursor is what the walk used to build per call, and it knows
        // nothing.
        let mut fresh = Cursor::default();
        fresh.get_or_read(&source, first, "x.bam").unwrap();
        assert!(source.reads.load(Ordering::SeqCst) > before);
    }

    /// The count alone is no bound: chunks merge up to `MAX_MERGE_SPAN`
    /// compressed, and a cursor now lives as long as its walk.
    #[test]
    fn a_cursor_stays_under_its_byte_budget() {
        // Four chunks whose payloads together are well past the budget.
        let each = CURSOR_CACHE_BYTES / 2 + 1;
        let (source, offsets) = blocks(4, each);
        let mut cursor = Cursor::default();
        for i in 0..4 {
            cursor
                .get_or_read(&source, chunk(offsets[i], offsets[i + 1]), "x.bam")
                .unwrap();
            let held: usize = cursor.cache.iter().map(|(_, d)| d.len()).sum();
            assert_eq!(held, cursor.bytes, "the running total drifted");
            assert!(
                cursor.bytes <= CURSOR_CACHE_BYTES || cursor.cache.len() == 1,
                "after {} chunks the cursor holds {} bytes",
                i + 1,
                cursor.bytes
            );
        }
        // And a chunk too big to share is returned without being kept.
        let (big_source, big_offsets) = blocks(1, CURSOR_CACHE_BYTES * 2);
        let mut cursor = Cursor::default();
        let data = cursor
            .get_or_read(&big_source, chunk(big_offsets[0], big_offsets[1]), "x.bam")
            .unwrap();
        assert_eq!(data.len(), CURSOR_CACHE_BYTES * 2);
        assert!(cursor.cache.is_empty(), "an oversized chunk was cached");
    }
}