gwseq_io/bbi/

extract.rs

//! The extraction kernels.
//!
//! Where the read time goes. They share one shape: walk the R-tree for the leaves the batch's loci touch,
//! inflate each leaf, walk its items forward, and fold every overlap into the
//! output.
//!
//! # The loc cursor
//!
//! A block hands out its items in increasing (chromosome, position) order, so a
//! locus the cursor has moved past cannot come back into range. Without the
//! cursor, every item rescans its block's loci from the first, which makes a
//! block of n items against m loci cost n·m. It is not an optimisation to leave
//! for later — it is the difference between linear and quadratic on a request
//! with many loci.
//!
//! # Where each worker writes
//!
//! A batch is a contiguous run of loci **in file order**, but a locus's output
//! slice follows the order the request asked in — so a batch's outputs are
//! scattered through the result, not contiguous in it. Each worker therefore
//! fills a compact buffer covering only its own loci, and the caller scatters
//! those into the result once every worker has finished. No two workers ever
//! touch the same bin, so no bin's accumulation order depends on scheduling.

use bytes::Bytes;

use crate::bbi::block::{DataInterval, DataIntervals};
use crate::bbi::rtree::LeafWalk;
use crate::error::{Error, Result};
use crate::genomic::{BinMode, BinStats, IndexedLoc, IndexedLocs, LocBatch, ValueStats};
use crate::progress::ProgressTracker;
use crate::source::ByteSource;

/// Everything a kernel needs that is the same for every batch.
pub(crate) struct Extraction<'a> {
    pub source: &'a dyn ByteSource,
    pub locs: &'a IndexedLocs,
    pub batches: &'a [LocBatch],
    /// Where the R-tree's root node begins — past the 48-byte index header.
    pub tree_root: u64,
    /// True when reading a zoom level rather than the full data.
    pub zoom: bool,
    pub uncompress_buffer_size: u32,
    pub tracker: &'a ProgressTracker,
}

impl Extraction<'_> {
    fn read_leaf(&self, offset: u64, size: u64) -> Result<Bytes> {
        let raw = self.source.read_exact_at(offset, size as usize)?;
        crate::bbi::block::decompress(raw, self.uncompress_buffer_size, self.source.path())
    }

    /// Walk every (leaf, loci) pair of one batch, handing each interval to
    /// `visit` together with the locus it overlaps and their overlap.
    ///
    /// The three value kernels differ only in what they do with that, so the
    /// walk — the R-tree descent, the block inflation, the cursor, the four
    /// overlap tests — lives here once.
    fn walk_batch(
        &self,
        batch: LocBatch,
        mut visit: impl FnMut(&DataInterval, usize, &IndexedLoc, i64, i64) -> Result<()>,
    ) -> Result<()> {
        let locs = &self.locs.locs;
        let leaves = LeafWalk::new(self.source, self.tree_root, locs, batch, self.tracker)?;
        for leaf in leaves {
            let (leaf, loc_range) = leaf?;
            let block = self.read_leaf(leaf.offset, leaf.size)?;
            let intervals = DataIntervals::new(
                block,
                self.zoom,
                locs,
                loc_range.clone(),
                self.source.path(),
            )?;
            let mut cursor = loc_range.start;
            // Indexed rather than iterated: `visit` is handed the locus's index
            // in the whole request, which is what tells a kernel where in its
            // output the locus belongs. Enumerating a subslice would report a
            // different number.
            #[allow(clippy::needless_range_loop)]
            for interval in intervals {
                let interval = interval?;
                cursor = advance_cursor(
                    locs,
                    cursor,
                    loc_range.end,
                    interval.chr_index,
                    interval.start,
                );
                for index in cursor..loc_range.end {
                    let loc = &locs[index];
                    // A tree item may straddle a chromosome boundary, so its
                    // loci are not all on the interval's chromosome. The cursor
                    // has already passed the ones behind it, so the first one
                    // beyond ends the scan.
                    if interval.chr_index != loc.chr_index as u32 {
                        break;
                    }
                    if interval.end <= loc.binned_start {
                        break;
                    }
                    // The grid left this locus narrower than a bin: it reads
                    // nothing, and its bin size is 0.
                    if loc.binned_end <= loc.binned_start {
                        continue;
                    }
                    if interval.start >= loc.binned_end {
                        continue;
                    }
                    let overlap_start = interval.start.max(loc.binned_start);
                    let overlap_end = interval.end.min(loc.binned_end);
                    visit(&interval, index, loc, overlap_start, overlap_end)?;
                }
            }
        }
        Ok(())
    }
}

/// The first locus at or after `cursor` that a value starting at
/// (`chr`, `start`), or any value after it, could still overlap.
#[inline]
fn advance_cursor(
    locs: &[IndexedLoc],
    mut cursor: usize,
    end: usize,
    chr: u32,
    start: i64,
) -> usize {
    while cursor < end {
        let loc = &locs[cursor];
        let loc_chr = loc.chr_index as u32;
        if loc_chr > chr {
            break;
        }
        if loc_chr == chr && loc.binned_end > start {
            break;
        }
        cursor += 1;
    }
    cursor
}

/// Values into a flat `(loci × bins)` buffer. The `read_values` kernel.
pub(crate) fn values(
    ex: &Extraction<'_>,
    executor: &crate::parallel::Executor,
    bin_mode: BinMode,
    def_value: f32,
) -> Result<Vec<f32>> {
    let bin_count = ex.locs.bin_count;
    let per_batch = executor.map_batches(ex.batches, |_, batch| {
        let mut stats = vec![BinStats::default(); batch.len() * bin_count];
        ex.walk_batch(*batch, |interval, index, loc, from, to| {
            let base = (index - batch.start) * bin_count;
            let bin_start = loc.bin_at(from);
            let bin_end = loc.bin_after(to);
            // Weighted by the bases of the bin the interval covers, not by the
            // interval itself: a record stands for a range, and at a zoom level
            // that range is a whole window. Counting it once per bin makes
            // `sum` and `count` track how the file cut its records. The two
            // agree where a record covers one base.
            for b in bin_start..bin_end {
                if b as usize >= bin_count {
                    break;
                }
                let covered = loc.bin_coverage(b, from, to);
                if covered <= 0.0 {
                    continue;
                }
                stats[base + b as usize].add(interval.value, covered);
            }
            Ok(())
        })?;
        Ok(stats)
    })?;

    let mut output = vec![def_value; ex.locs.output_len];
    for (batch, stats) in ex.batches.iter().zip(&per_batch) {
        for (offset, index) in (batch.start..batch.end).enumerate() {
            let loc = &ex.locs.locs[index];
            for b in 0..bin_count {
                let s = &stats[offset * bin_count + b];
                if s.count <= 0.0 {
                    continue;
                }
                output[loc.output_start + b] = s.apply(bin_mode);
            }
        }
    }
    ex.locs.reverse_output_rows(&mut output);
    Ok(output)
}

/// One [`ValueStats`] per locus, for `quantify`.
pub(crate) fn values_stats(
    ex: &Extraction<'_>,
    executor: &crate::parallel::Executor,
) -> Result<Vec<ValueStats>> {
    let bin_count = ex.locs.bin_count;
    let per_batch = executor.map_batches(ex.batches, |_, batch| {
        let mut stats = vec![ValueStats::default(); batch.len()];
        ex.walk_batch(*batch, |interval, index, _loc, from, to| {
            let overlap = to - from;
            // Weighted by the bases the interval has data for, not by the bases
            // it spans: a zoom record summarises a fixed window that may be
            // mostly empty, so weighting by the span would grow every sum and
            // count with the zoom level. The identity at full resolution.
            let span = interval.end - interval.start;
            let covered = if span > 0 && interval.valid_count != span {
                (interval.valid_count as f64 * overlap as f64 / span as f64).round() as i64
            } else {
                overlap
            };
            // The record's own sum of squares, prorated over the part this
            // window takes. Squaring the mean would keep only the variance
            // *between* windows and drop the variance inside each, collapsing
            // `sd` as the zoom level rises.
            let fraction = if interval.valid_count > 0 {
                covered as f64 / interval.valid_count as f64
            } else {
                0.0
            };
            stats[index - batch.start].add_aggregate(
                interval.min_value,
                interval.max_value,
                interval.value as f64 * covered as f64,
                interval.sum_squared * fraction,
                covered,
            );
            Ok(())
        })?;
        Ok(stats)
    })?;

    let mut output = vec![ValueStats::default(); ex.locs.locs.len()];
    for (batch, stats) in ex.batches.iter().zip(&per_batch) {
        for (offset, index) in (batch.start..batch.end).enumerate() {
            output[ex.locs.locs[index].row(bin_count)] = stats[offset];
        }
    }
    Ok(output)
}

/// The bin of a locus a profile walk is currently filling.
#[derive(Debug, Clone, Copy)]
struct OpenBin {
    /// -1 before the first value.
    bin: i64,
    stats: BinStats,
    /// The locus is read from its end, so the bin lands mirrored in the
    /// profile. Carried here rather than looked up per value: the walk holds
    /// one of these per locus of its batch for as long as it lasts.
    reverse: bool,
}

/// One [`ValueStats`] per bin, folded across loci, for `profile`.
///
/// The one kernel whose workers all contribute to every bin, so each keeps a
/// private `bin_count` row and they are merged at the end **in batch order**.
/// That merge order is behaviour: it fixes the `f32` accumulation and so the
/// last bits of the result.
///
/// A bin is folded into its column as soon as the walk leaves it, so a locus
/// never holds more than the bin it is filling. The walk only moves forward, so
/// bins fill left to right and none is ever reopened — holding them all would
/// be the whole (locus × bin) matrix.
pub(crate) fn values_profile(
    ex: &Extraction<'_>,
    executor: &crate::parallel::Executor,
    bin_mode: BinMode,
) -> Result<Vec<ValueStats>> {
    let bin_count = ex.locs.bin_count;
    let per_batch = executor.map_batches(ex.batches, |_, batch| {
        let mut column = vec![ValueStats::default(); bin_count];
        let mut open: Vec<OpenBin> = (batch.start..batch.end)
            .map(|i| OpenBin {
                bin: -1,
                stats: BinStats::default(),
                reverse: ex.locs.locs[i].reverse,
            })
            .collect();

        let close = |open: &mut OpenBin, column: &mut Vec<ValueStats>| {
            if open.stats.count <= 0.0 {
                return;
            }
            let value = open.stats.apply(bin_mode);
            // A locus read from its end has its bins mirrored into the profile.
            // Only here, and not in `open.bin` itself, which the walk compares
            // against the bin it is filling.
            let col = if open.reverse {
                bin_count as i64 - 1 - open.bin
            } else {
                open.bin
            };
            if col >= 0 && (col as usize) < column.len() {
                // One bin's value, folded in at full width. Squaring in `f32`
                // and widening afterwards throws away half the mantissa of the
                // square before the sum ever sees it, and gives a profile of
                // constant data a standard deviation of 1.13 where the answer
                // is 0.
                column[col as usize].add(value);
            }
            open.stats = BinStats::default();
        };

        ex.walk_batch(*batch, |interval, index, loc, from, to| {
            let bin_start = loc.bin_at(from);
            let bin_end = loc.bin_after(to);
            let slot = &mut open[index - batch.start];
            // Weighted by covered bases, as `values` weights its own — a
            // profile is read_values by column, so the two have to bin alike.
            for b in bin_start..bin_end {
                if b as usize >= bin_count {
                    break;
                }
                let covered = loc.bin_coverage(b, from, to);
                if covered <= 0.0 {
                    continue;
                }
                if b != slot.bin {
                    close(slot, &mut column);
                    slot.bin = b;
                }
                slot.stats.add(interval.value, covered);
            }
            Ok(())
        })?;
        for slot in &mut open {
            close(slot, &mut column);
        }
        Ok(column)
    })?;

    let mut output = vec![ValueStats::default(); bin_count];
    for column in &per_batch {
        for (col, batch_stats) in column.iter().enumerate() {
            if batch_stats.count == 0 {
                continue;
            }
            output[col].merge(batch_stats);
        }
    }
    Ok(output)
}

/// Walk every (leaf, loci) pair of one batch's bed blocks.
///
/// The bed twin of [`Extraction::walk_batch`]: same R-tree descent, same
/// cursor, same four overlap tests, but the block decodes into entries rather
/// than intervals. `visit` is handed the entry and the index of every locus it
/// overlaps.
impl Extraction<'_> {
    pub(crate) fn walk_bed_batch(
        &self,
        batch: LocBatch,
        auto_sql: &indexmap::IndexMap<String, String>,
        col_count: usize,
        mut visit: impl FnMut(&mut BedRecord, &[usize]) -> Result<()>,
    ) -> Result<()> {
        let locs = &self.locs.locs;
        let leaves = LeafWalk::new(self.source, self.tree_root, locs, batch, self.tracker)?;
        let mut matched: Vec<usize> = Vec::new();
        for leaf in leaves {
            let (leaf, loc_range) = leaf?;
            let block = self.read_leaf(leaf.offset, leaf.size)?;
            let records = super::block::BedRecords::new(
                block,
                auto_sql,
                col_count,
                locs,
                loc_range.clone(),
                self.source.path(),
            )?;
            let mut cursor = loc_range.start;
            for record in records {
                let (chr_index, start, end, fields) = record?;
                cursor = advance_cursor(locs, cursor, loc_range.end, chr_index, start);
                matched.clear();
                // A zero-length entry occupies the single base it names. See
                // `BedRecords::next`.
                let reach = end.max(start + 1);
                // Indexed rather than enumerated: `visit` is handed the locus's
                // index in the whole request, which is what tells a kernel where
                // in its output the locus belongs.
                #[allow(clippy::needless_range_loop)]
                for index in cursor..loc_range.end {
                    let loc = &locs[index];
                    if chr_index != loc.chr_index as u32 {
                        break;
                    }
                    if reach <= loc.binned_start {
                        break;
                    }
                    if start >= loc.binned_end {
                        continue;
                    }
                    matched.push(index);
                }
                if matched.is_empty() {
                    continue;
                }
                let mut entry = BedRecord {
                    chr_index,
                    start,
                    end,
                    fields,
                };
                visit(&mut entry, &matched)?;
            }
        }
        Ok(())
    }
}

/// A bed record as the walk hands it over: coordinates plus whatever columns
/// were asked for, before the chromosome index becomes a name.
pub(crate) struct BedRecord {
    pub chr_index: u32,
    pub start: i64,
    pub end: i64,
    pub fields: Vec<(String, String)>,
}

/// bigBed entries per locus.
///
/// An entry belongs to every locus it overlaps, so two overlapping loci both
/// report the entries they share. The fields are moved into the **last** locus
/// that claims the entry and cloned into the rest, which is what keeps the
/// common case — an entry in exactly one locus — free of a copy.
pub(crate) fn entries(
    ex: &Extraction<'_>,
    executor: &crate::parallel::Executor,
    auto_sql: &indexmap::IndexMap<String, String>,
    chr_names: &[String],
    col_count: usize,
) -> Result<Vec<Vec<super::BedEntry>>> {
    let bin_count = ex.locs.bin_count;
    let per_batch = executor.map_batches(ex.batches, |_, batch| {
        let mut out: Vec<Vec<super::BedEntry>> = vec![Vec::new(); batch.len()];
        ex.walk_bed_batch(*batch, auto_sql, col_count, |entry, matched| {
            let chr = chr_names
                .get(entry.chr_index as usize)
                .cloned()
                .unwrap_or_default();
            for (n, index) in matched.iter().enumerate() {
                let fields = if n + 1 == matched.len() {
                    std::mem::take(&mut entry.fields)
                } else {
                    entry.fields.clone()
                };
                out[index - batch.start].push(super::BedEntry {
                    chr: chr.clone(),
                    start: entry.start,
                    end: entry.end,
                    fields,
                });
            }
            Ok(())
        })?;
        Ok(out)
    })?;

    let mut output: Vec<Vec<super::BedEntry>> = vec![Vec::new(); ex.locs.locs.len()];
    for (batch, lists) in ex.batches.iter().zip(per_batch) {
        for (offset, list) in lists.into_iter().enumerate() {
            output[ex.locs.locs[batch.start + offset].row(bin_count)] = list;
        }
    }
    // A locus reads its entries in block order, and a block straddling two loci
    // hands them over per block rather than per position, so the list has to be
    // put back in coordinate order.
    for entries in &mut output {
        entries.sort_by(|a, b| (&a.chr, a.start, a.end).cmp(&(&b.chr, b.start, b.end)));
    }
    Ok(output)
}

/// Depth of coverage of a bigBed's entries, binned — what a bigBed means by a
/// "value".
///
/// The depth over the bin, not the number of entries touching it: an entry
/// covering a third of a bin raises its depth by a third, so a bin holds the
/// mean depth over the bases it spans whatever its width. The same at the
/// default bin size of one.
pub(crate) fn entries_pileup(
    ex: &Extraction<'_>,
    executor: &crate::parallel::Executor,
    auto_sql: &indexmap::IndexMap<String, String>,
    def_value: f32,
) -> Result<Vec<f32>> {
    let bin_count = ex.locs.bin_count;
    let per_batch = executor.map_batches(ex.batches, |_, batch| {
        let mut depth = vec![0.0f32; batch.len() * bin_count];
        // A pileup counts coverage, so it asks for the 3 coordinate columns
        // alone and never pays for the fields it would drop.
        ex.walk_bed_batch(*batch, auto_sql, 3, |entry, matched| {
            for index in matched {
                let loc = &ex.locs.locs[*index];
                if loc.binned_end <= loc.binned_start {
                    continue;
                }
                let from = entry.start.max(loc.binned_start);
                let to = entry.end.min(loc.binned_end);
                let base = (index - batch.start) * bin_count;
                for b in loc.bin_at(from)..loc.bin_after(to) {
                    if b as usize >= bin_count {
                        break;
                    }
                    let fraction = loc.bin_fraction(b, from, to);
                    if fraction <= 0.0 {
                        continue;
                    }
                    depth[base + b as usize] += fraction as f32;
                }
            }
            Ok(())
        })?;
        Ok(depth)
    })?;

    let mut output = vec![0.0f32; ex.locs.output_len];
    for (batch, depth) in ex.batches.iter().zip(&per_batch) {
        for (offset, index) in (batch.start..batch.end).enumerate() {
            let loc = &ex.locs.locs[index];
            output[loc.output_start..loc.output_end]
                .copy_from_slice(&depth[offset * bin_count..(offset + 1) * bin_count]);
        }
    }
    // A pileup of 0 is exactly a bin no entry reached, which is what def_value
    // stands for.
    if def_value != 0.0 {
        for value in &mut output {
            if *value == 0.0 {
                *value = def_value;
            }
        }
    }
    // Mirrored here rather than at each of the three callers: the profile reads
    // this by column and the quantification by row, and both are right once the
    // rows themselves are the way they were asked for.
    ex.locs.reverse_output_rows(&mut output);
    Ok(output)
}

/// The pileup reduced to one [`ValueStats`] per locus, for `quantify`.
///
/// Over the bins the request asked for, not one per locus. Pinning `bin_count`
/// to 1 would pile every entry of a locus into a single bin and leave the
/// reduction nothing to run over — mean, sum, min and max would all come back
/// as the number of overlapping entries.
///
/// A bin no entry reached holds `def_value`, put there by the pileup, so every
/// bin counts and the reduction sees the numbers `read_values` would show. A
/// NaN `def_value` is how a caller asks for the uncovered bins to be left out,
/// and they are skipped rather than poisoning every statistic of the locus.
pub(crate) fn pileup_stats(locs: &IndexedLocs, pileup: &[f32]) -> Vec<ValueStats> {
    let mut output = vec![ValueStats::default(); locs.locs.len()];
    for loc in &locs.locs {
        let stats = &mut output[loc.row(locs.bin_count)];
        for value in &pileup[loc.output_start..loc.output_end] {
            if value.is_nan() {
                continue;
            }
            stats.add(*value);
        }
    }
    output
}

/// The pileup reduced to one [`ValueStats`] per bin, for `profile`.
///
/// Every locus holds a value in every bin — one no entry reached piles up to
/// `def_value` — so every one of them counts, and dividing the sum by anything
/// less would inflate it. A NaN `def_value` is again how a caller asks for the
/// loci whose bin held nothing to be left out.
pub(crate) fn pileup_profile(locs: &IndexedLocs, pileup: &[f32]) -> Vec<ValueStats> {
    let mut output = vec![ValueStats::default(); locs.bin_count];
    for (col, stats) in output.iter_mut().enumerate() {
        for loc in &locs.locs {
            let value = pileup[loc.output_start + col];
            if value.is_nan() {
                continue;
            }
            stats.add(value);
        }
    }
    output
}

// ---------------------------------------------------------------------------
// Whole-file iterators
// ---------------------------------------------------------------------------
//
// Plain `Iterator`s, exhausted after one pass, exposing their windows' regions
// up front. A window never spans two chromosomes and no bin straddles a window
// boundary, so concatenating a chromosome's windows gives exactly what a
// whole-chromosome read gives at the same bin size.

/// Region of one window: chromosome, start, end.
pub type WindowLoc = (String, i64, i64);

/// Bytes a piece has to stand to read before splitting a window is worth it.
///
/// A piece costs a thread and an index descent, which one holding next to
/// nothing never earns back: on a sparse file, where a megabase window is a few
/// hundred bytes, a walk split 24 ways measures several times slower than one
/// read window by window.
const MIN_PIECE_DATA_SIZE: f64 = 16384.0;

/// What both whole-file walks share: the windows, where the walk has got to,
/// and how finely a window is worth splitting.
///
/// `locs` is behind an [`Arc`](std::sync::Arc) so that a walk can be restarted
/// without rebuilding it — see [`Walk::restarted`]. It is the only field with
/// a heap allocation, and for a whole-genome walk it is a few thousand
/// entries, so sharing it makes a restart cost a refcount bump instead of a
/// copy.
#[derive(Debug)]
struct Walk {
    locs: std::sync::Arc<Vec<WindowLoc>>,
    next: usize,
    parallel: usize,
    /// Bytes of the file's data one base pair of the genome holds on average,
    /// which is what a window's share of the file is estimated from.
    bytes_per_bp: f64,
    total_coverage: u64,
    done_coverage: u64,
}

impl Walk {
    fn new(locs: Vec<WindowLoc>, parallel: usize, data_size: u64, genome_size: i64) -> Self {
        let total_coverage = locs.iter().map(|(_, s, e)| (e - s).max(0) as u64).sum();
        let bytes_per_bp = if genome_size < 1 || data_size < 1 {
            0.0
        } else {
            data_size as f64 / genome_size as f64
        };
        Self {
            locs: std::sync::Arc::new(locs),
            next: 0,
            parallel: parallel.max(1),
            bytes_per_bp,
            total_coverage,
            done_coverage: 0,
        }
    }

    /// The same plan, back at the first window.
    ///
    /// Shares `locs` rather than copying it, so this is a refcount bump and
    /// two integers. What resets is the cursor and the progress tally: a
    /// second pass reports its own progress from zero, which is what a caller
    /// watching it expects.
    fn restarted(&self) -> Self {
        Self {
            locs: self.locs.clone(),
            next: 0,
            parallel: self.parallel,
            bytes_per_bp: self.bytes_per_bp,
            total_coverage: self.total_coverage,
            done_coverage: 0,
        }
    }

    /// Split a window into the pieces the threads share out.
    ///
    /// `units` is what the window is measured in and what the pieces divide:
    /// bins for a values walk, base pairs for an entries walk. Returns how many
    /// units a piece holds and how many pieces that takes.
    ///
    /// A piece is never empty of units, so a window with fewer units than
    /// pieces comes back as fewer pieces rather than as empty ones. The last
    /// piece is the one that may reach past the window — the units rarely
    /// divide evenly — and the caller drops what it read out there.
    fn split(&self, units: i64, coverage: i64) -> (i64, i64) {
        let window_data = self.bytes_per_bp * coverage as f64;
        let worth = (window_data / MIN_PIECE_DATA_SIZE) as i64;
        let pieces = worth.min(self.parallel as i64).max(1);
        let piece_units = ((units + pieces - 1) / pieces).max(1);
        (piece_units, (units + piece_units - 1) / piece_units)
    }

    /// Report the window the walk stands on and move past it.
    ///
    /// Reported as a window is handed over rather than as one is read, so that
    /// progress counts what the caller has seen.
    fn take(&mut self, progress: Option<&crate::progress::ProgressFn>) -> usize {
        let index = self.next;
        self.next += 1;
        let (_, start, end) = &self.locs[index];
        self.done_coverage += (end - start).max(0) as u64;
        if let Some(report) = progress {
            report(self.done_coverage, self.total_coverage);
        }
        index
    }

    /// The final report, once the last window has been handed over — or
    /// straight away for a walk over no window at all.
    fn finish(&mut self, progress: Option<&crate::progress::ProgressFn>) {
        if let Some(report) = progress {
            if self.done_coverage < self.total_coverage {
                self.done_coverage = self.total_coverage;
                report(self.total_coverage, self.total_coverage);
            }
        }
    }
}

/// Successive windows of values over whole chromosomes.
///
/// Exhausted after one pass. `locs()` gives the region of each window up
/// front, so `iter.locs().to_vec()` before the walk is how both are had at
/// once.
pub struct ValuesWalk {
    walk: Walk,
    bin_size: i64,
    /// Bins each window hands back, which for the last window of a chromosome
    /// is fewer than a full window holds. Shared with any restart of this
    /// walk, as `Walk::locs` is.
    bins: std::sync::Arc<Vec<i64>>,
    bin_mode: BinMode,
    def_value: f32,
    zoom: super::Zoom,
    progress: Option<crate::progress::ProgressFn>,
}

/// Hand-written: a progress callback is a boxed closure and cannot be `Debug`.
impl std::fmt::Debug for ValuesWalk {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("ValuesWalk")
            .field("windows", &self.walk.locs.len())
            .field("next", &self.walk.next)
            .field("bin_size", &self.bin_size)
            .finish()
    }
}

impl ValuesWalk {
    /// The same walk, back at its first window.
    ///
    /// What makes the Python iterator re-iterable: `__iter__` hands back a
    /// fresh cursor over the plan that is already built, so a second `for`
    /// loop walks the file again instead of yielding nothing. Costs two
    /// refcount bumps — the windows and the per-window bin counts are shared,
    /// not copied — and re-reads the file, which is what a second pass is.
    pub fn restarted(&self) -> Self {
        Self {
            walk: self.walk.restarted(),
            bin_size: self.bin_size,
            bins: self.bins.clone(),
            bin_mode: self.bin_mode,
            def_value: self.def_value,
            zoom: self.zoom,
            progress: self.progress.clone(),
        }
    }

    /// Resolve the whole walk up front: the windows, the zoom level, and every
    /// argument that could be refused.
    ///
    /// A request that cannot be read says so when it is made rather than part
    /// way through the iteration.
    pub fn plan(reader: &super::BbiReader, req: &super::ValuesRequest, span: i64) -> Result<Self> {
        if span < 1 {
            return Err(Error::invalid(format!(
                "span must be positive (got {span})"
            )));
        }
        // A walk lays its windows out on the grid the whole genome shares, so
        // the bin size has to be a whole number of base pairs — which
        // `BinPlan::new` now requires of every request, not only of this one.
        let bin_size =
            crate::genomic::BinPlan::new(req.common.bin_size, None, req.common.full_bin)?
                .whole_bin_size();
        if reader.kind().is_bigbed() && !matches!(req.common.zoom, super::Zoom::Full) {
            return Err(Error::invalid("zoom is only supported for bigwig files"));
        }
        // Resolved here rather than at the first step, so a request that cannot
        // be read says so when it is made instead of part way through the walk.
        let level = reader.select_zoom(req.common.bin_size, req.common.zoom)?;

        // Counted in bins rather than base pairs, so no bin straddles a window
        // boundary; `span` is rounded up to a whole number of them.
        let bins_per_window = ((span + bin_size - 1) / bin_size).max(1);
        let mut locs = Vec::new();
        let mut bins = Vec::new();
        for chr in reader.chr_sizes().select(&req.common.locs.chr_ids)? {
            // The bins of the chromosome, which the windows share out. Its last
            // is partial, and only `full_bin` keeps it — as it is only
            // `full_bin` that keeps it in a read of the whole chromosome.
            let chr_bins = if req.common.full_bin {
                (chr.size + bin_size - 1) / bin_size
            } else {
                chr.size / bin_size
            };
            let mut bin = 0;
            while bin < chr_bins {
                let start = bin * bin_size;
                let window_bins = bins_per_window.min(chr_bins - bin);
                locs.push((
                    chr.id.clone(),
                    start,
                    (start + window_bins * bin_size).min(chr.size),
                ));
                bins.push(window_bins);
                bin += bins_per_window;
            }
        }

        let walk = Walk::new(
            locs,
            reader.parallel(),
            reader.data_size(level),
            reader.genome_size(),
        );
        Ok(Self {
            walk,
            bin_size,
            bins: std::sync::Arc::new(bins),
            bin_mode: req.bin_mode,
            def_value: req.common.def_value,
            zoom: req.common.zoom,
            progress: req.common.progress.clone(),
        })
    }

    /// Number of windows, so `len(iterator)` works on the Python side.
    pub fn len(&self) -> usize {
        self.walk.locs.len()
    }

    /// The whole number of base pairs each bin covers.
    ///
    /// Resolved at `plan` from the request's `f64`, so a caller laying window
    /// values out on coordinates — as the exporters do — reads the grid the
    /// walk actually used rather than rounding the request a second time.
    pub fn bin_size(&self) -> i64 {
        self.bin_size
    }

    pub fn is_empty(&self) -> bool {
        self.walk.locs.is_empty()
    }

    /// The region of each window: the nth array covers `locs()[n]`.
    pub fn locs(&self) -> &[WindowLoc] {
        &self.walk.locs
    }

    /// Read one window as pieces of itself, laid end to end.
    fn read(&self, reader: &super::BbiReader, index: usize) -> Result<Vec<f32>> {
        let (chr, start, end) = &self.walk.locs[index];
        let (piece_bins, piece_count) = self.walk.split(self.bins[index], end - start);
        let piece_span = piece_bins * self.bin_size;

        let chr_ids = vec![chr.clone(); piece_count as usize];
        let starts: Vec<i64> = (0..piece_count).map(|i| start + i * piece_span).collect();
        let ends: Vec<i64> = starts.iter().map(|s| s + piece_span).collect();

        // Every piece is read at the same width, which is what settles the read
        // on `piece_bins` bins for all of them and lays the window out as one
        // buffer. `full_bin` is false whatever the walk was asked for: the
        // pieces are on the grid already — it is the windows themselves that a
        // full_bin walk lays out one bin further.
        //
        // Progress is not passed down: it is reported per window handed over, so
        // a step never reports a window the caller has not seen.
        let request =
            super::ValuesRequest::new(crate::genomic::Locs::spans(&chr_ids, &starts, &ends)?)
                .bin_size(self.bin_size as f64)
                .bin_count(piece_bins as usize)
                .bin_mode(self.bin_mode)
                .def_value(self.def_value)
                .zoom(self.zoom);
        let values = reader.read_values(&request)?;
        Ok(values.into_raw_vec_and_offset().0)
    }

    /// The next window, or `None` once the walk is spent.
    ///
    /// Takes the reader on every step rather than holding it: a `#[pyclass]`
    /// cannot carry a borrow, and the Python iterator drives exactly this.
    pub fn next_window(
        &mut self,
        reader: &super::BbiReader,
    ) -> Option<Result<ndarray::Array1<f32>>> {
        if self.walk.next >= self.walk.locs.len() {
            self.walk.finish(self.progress.as_ref());
            return None;
        }
        // Read before the window is taken; `take` is what moves the walk on.
        let index = self.walk.next;
        let mut values = match self.read(reader, index) {
            Ok(v) => v,
            // 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 reporting the walk spent. That is what makes closing
            // a reader under a walk raise every time instead of once and then
            // stopping, as `LocusWalk` does and as the "Non-obvious
            // constraints" section of ARCHITECTURE requires.
            Err(e) => return Some(Err(e)),
        };
        self.walk.take(self.progress.as_ref());
        // The pieces rarely divide the window evenly, so the last one may reach
        // past it; an array standing for more than the window it names is not
        // one to hand over.
        values.truncate(self.bins[index] as usize);
        Some(Ok(ndarray::Array1::from_vec(values)))
    }
}

/// [`ValuesWalk`] as a plain [`Iterator`], for Rust callers.
///
/// Exhausted after one pass.
#[derive(Debug)]
pub struct ValuesWindows<'a> {
    reader: &'a super::BbiReader,
    walk: ValuesWalk,
}

impl<'a> ValuesWindows<'a> {
    pub(crate) fn plan(
        reader: &'a super::BbiReader,
        req: &super::ValuesRequest,
        span: i64,
    ) -> Result<Self> {
        Ok(Self {
            reader,
            walk: ValuesWalk::plan(reader, req, span)?,
        })
    }

    pub fn len(&self) -> usize {
        self.walk.len()
    }
    pub fn is_empty(&self) -> bool {
        self.walk.is_empty()
    }
    /// The region of each window: the nth array covers `locs()[n]`.
    pub fn locs(&self) -> &[WindowLoc] {
        self.walk.locs()
    }
    /// The whole number of base pairs each bin covers. See
    /// [`ValuesWalk::bin_size`].
    pub fn bin_size(&self) -> i64 {
        self.walk.bin_size()
    }
}

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

/// Successive windows of bed entries over whole chromosomes.
pub struct EntryWalk {
    walk: Walk,
    col_count: usize,
    progress: Option<crate::progress::ProgressFn>,
}

/// Hand-written: a progress callback is a boxed closure and cannot be `Debug`.
impl std::fmt::Debug for EntryWalk {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("EntryWalk")
            .field("windows", &self.walk.locs.len())
            .field("next", &self.walk.next)
            .field("col_count", &self.col_count)
            .finish()
    }
}

impl EntryWalk {
    /// The same walk, back at its first window. See
    /// [`ValuesWalk::restarted`].
    pub fn restarted(&self) -> Self {
        Self {
            walk: self.walk.restarted(),
            col_count: self.col_count,
            progress: self.progress.clone(),
        }
    }

    pub fn plan(reader: &super::BbiReader, req: &super::EntriesRequest, span: i64) -> Result<Self> {
        if !reader.kind().is_bigbed() {
            return Err(Error::invalid("iter_all_entries only for bigbed"));
        }
        if span < 1 {
            return Err(Error::invalid(format!(
                "span must be positive (got {span})"
            )));
        }
        reader.check_col_count(req.col_count, 3)?;

        let mut locs = Vec::new();
        for chr in reader.chr_sizes().select(&req.common.locs.chr_ids)? {
            let mut start = 0;
            while start < chr.size {
                locs.push((chr.id.clone(), start, (start + span).min(chr.size)));
                start += span;
            }
        }
        // No reordering: the windows are built in chromosome then coordinate
        // order already, which is the order to read them in.
        let walk = Walk::new(
            locs,
            reader.parallel(),
            reader.data_size(None),
            reader.genome_size(),
        );
        Ok(Self {
            walk,
            col_count: req.col_count,
            progress: req.common.progress.clone(),
        })
    }

    pub fn len(&self) -> usize {
        self.walk.locs.len()
    }

    pub fn is_empty(&self) -> bool {
        self.walk.locs.is_empty()
    }

    pub fn locs(&self) -> &[WindowLoc] {
        &self.walk.locs
    }

    fn read(&self, reader: &super::BbiReader, index: usize) -> Result<Vec<super::BedEntry>> {
        let (chr, start, end) = &self.walk.locs[index];
        let coverage = end - start;
        let (piece_span, piece_count) = self.walk.split(coverage, coverage);

        let chr_ids = vec![chr.clone(); piece_count as usize];
        let mut starts = Vec::with_capacity(piece_count as usize);
        let mut ends = Vec::with_capacity(piece_count as usize);
        let mut piece_min_starts = Vec::with_capacity(piece_count as usize);
        for i in 0..piece_count {
            let piece_start = start + i * piece_span;
            piece_min_starts.push(piece_start);
            // Read from a base before the piece, so an entry covering no base —
            // a zero-length one at the piece's very first base — is not missed
            // by a test that asks it to overlap. The piece still *reports* from
            // its own start, which is what keeps an entry reaching over a
            // boundary to the one piece it starts in.
            starts.push((piece_start - 1).max(0));
            ends.push((piece_start + piece_span).min(*end));
        }

        let request =
            super::EntriesRequest::new(crate::genomic::Locs::spans(&chr_ids, &starts, &ends)?)
                .col_count(self.col_count);
        let pieces = reader.read_entries(&request)?;

        let mut out = Vec::new();
        for (piece, min_start) in pieces.into_iter().zip(piece_min_starts) {
            out.extend(piece.into_iter().filter(|e| e.start >= min_start));
        }
        Ok(out)
    }

    /// The next window, or `None` once the walk is spent.
    pub fn next_window(
        &mut self,
        reader: &super::BbiReader,
    ) -> Option<Result<Vec<super::BedEntry>>> {
        if self.walk.next >= self.walk.locs.len() {
            self.walk.finish(self.progress.as_ref());
            return None;
        }
        let index = self.walk.next;
        let entries = match self.read(reader, index) {
            Ok(e) => e,
            // Left where it was on failure. See `ValuesWalk::next_window`.
            Err(e) => return Some(Err(e)),
        };
        self.walk.take(self.progress.as_ref());
        Some(Ok(entries))
    }
}

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

impl<'a> EntryWindows<'a> {
    pub(crate) fn plan(
        reader: &'a super::BbiReader,
        req: &super::EntriesRequest,
        span: i64,
    ) -> Result<Self> {
        Ok(Self {
            reader,
            walk: EntryWalk::plan(reader, req, span)?,
        })
    }

    pub fn len(&self) -> usize {
        self.walk.len()
    }
    pub fn is_empty(&self) -> bool {
        self.walk.is_empty()
    }
    pub fn locs(&self) -> &[WindowLoc] {
        self.walk.locs()
    }
}

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

#[cfg(test)]
mod tests {
    use super::*;

    fn loc(chr: usize, start: i64, end: i64) -> IndexedLoc {
        IndexedLoc {
            chr_index: chr,
            start,
            end,
            binned_start: start,
            binned_end: end,
            bin_size: 1.0,
            reverse: false,
            output_start: 0,
            output_end: 1,
        }
    }

    #[test]
    fn the_cursor_skips_loci_the_block_has_passed() {
        let locs = [loc(0, 0, 10), loc(0, 20, 30), loc(0, 40, 50)];
        // A value at 25 cannot reach the first locus.
        assert_eq!(advance_cursor(&locs, 0, 3, 0, 25), 1);
        // One at 45 cannot reach the first two.
        assert_eq!(advance_cursor(&locs, 0, 3, 0, 45), 2);
        // One at 5 reaches everything from the start.
        assert_eq!(advance_cursor(&locs, 0, 3, 0, 5), 0);
    }

    #[test]
    fn the_cursor_stops_at_a_higher_chromosome() {
        let locs = [loc(0, 0, 10), loc(1, 0, 10), loc(2, 0, 10)];
        // A value on chromosome 1 passes the chromosome-0 locus and stops.
        assert_eq!(advance_cursor(&locs, 0, 3, 1, 5), 1);
        // A value on chromosome 0 never advances past its own.
        assert_eq!(advance_cursor(&locs, 0, 3, 0, 5), 0);
    }

    #[test]
    fn the_cursor_never_goes_backwards() {
        let locs = [loc(0, 0, 10), loc(0, 20, 30)];
        assert_eq!(advance_cursor(&locs, 1, 2, 0, 0), 1);
    }
}