Skip to main content

gwseq_io/bbi/
block.rs

1//! Decoding a decompressed data block.
2//!
3//! Wig section headers and their three item encodings,
4//! zoom records, and the [`DataIntervals`] walk that turns either into one
5//! stream of [`DataInterval`].
6//!
7//! Everything borrows. A block is inflated once into a [`Bytes`], and every
8//! item is read out of it by offset, so walking a block of a thousand items
9//! allocates once for the block and not at all per item.
10
11use bytes::Bytes;
12
13use crate::bytes::LeCursor;
14use crate::error::{Error, Result};
15use crate::genomic::IndexedLoc;
16
17/// Byte size of the header a wig section starts with (Supp. Table 13).
18pub const WIG_HEADER_SIZE: usize = 24;
19/// Byte size of a zoom record (Supp. Table 19).
20pub const ZOOM_RECORD_SIZE: usize = 32;
21/// Bytes of a bed record before its tab-separated tail (Supp. Table 12).
22pub const BED_RECORD_HEADER_SIZE: usize = 12;
23
24/// The three wig item encodings, in the order the format numbers them.
25#[derive(Debug, Clone, Copy, PartialEq, Eq)]
26pub enum WigEncoding {
27    /// Twelve bytes an item: start, end, value.
28    BedGraph = 1,
29    /// Eight: start, value. One span for the section.
30    VarStep = 2,
31    /// Four: value. One span and one step for the section.
32    FixedStep = 3,
33}
34
35impl WigEncoding {
36    pub fn from_u8(v: u8) -> Option<Self> {
37        match v {
38            1 => Some(WigEncoding::BedGraph),
39            2 => Some(WigEncoding::VarStep),
40            3 => Some(WigEncoding::FixedStep),
41            _ => None,
42        }
43    }
44
45    /// Bytes one item of this encoding occupies.
46    pub fn item_size(self) -> usize {
47        match self {
48            WigEncoding::BedGraph => 12,
49            WigEncoding::VarStep => 8,
50            WigEncoding::FixedStep => 4,
51        }
52    }
53}
54
55#[derive(Debug, Clone, Copy)]
56pub struct WigSectionHeader {
57    pub chr_index: u32,
58    pub chr_start: i64,
59    /// Decoded but not read: the item walk is bounded by `item_count`, not by
60    /// this. Kept so the struct is the on-disk header.
61    #[allow(dead_code)]
62    pub chr_end: i64,
63    pub item_step: i64,
64    pub item_span: i64,
65    pub encoding: WigEncoding,
66    pub item_count: u16,
67}
68
69pub fn read_wig_header(block: &[u8], path: &str) -> Result<WigSectionHeader> {
70    // Checked here rather than by the caller, which needs the item count out of
71    // these bytes before it can check anything else about the block.
72    if block.len() < WIG_HEADER_SIZE {
73        return Err(Error::corrupt(
74            path,
75            0,
76            format!(
77                "wig section header needs {WIG_HEADER_SIZE} bytes, its block holds {}",
78                block.len()
79            ),
80        ));
81    }
82    let mut c = LeCursor::new(block, 0, path);
83    let chr_index = c.read_u32()?;
84    let chr_start = c.read_u32()? as i64;
85    let chr_end = c.read_u32()? as i64;
86    let item_step = c.read_u32()? as i64;
87    let item_span = c.read_u32()? as i64;
88    let type_byte = c.read_u8()?;
89    c.skip(1)?; // reserved
90    let item_count = c.read_u16()?;
91    let encoding = WigEncoding::from_u8(type_byte)
92        .ok_or_else(|| Error::corrupt(path, 20, format!("wig data type {type_byte} invalid")))?;
93    Ok(WigSectionHeader {
94        chr_index,
95        chr_start,
96        chr_end,
97        item_step,
98        item_span,
99        encoding,
100        item_count,
101    })
102}
103
104/// One value of the file over a genomic range.
105///
106/// `valid_count` is the bases of the range that actually carry data: equal to
107/// the span at full resolution, but a zoom record summarises a fixed-width
108/// window that may be mostly empty, and weighting by the span instead would
109/// inflate every summed or counted quantification as the zoom level rises.
110///
111/// `min_value`/`max_value`/`sum_squared` are the record's **own**, not derived
112/// from `value`. At a zoom level `value` is the window mean, so taking extremes
113/// and spread from it would give the largest and smallest *window* rather than
114/// base — both pulled towards the middle, and further with every level — and
115/// would leave `sd` with only the variance between windows.
116#[derive(Debug, Clone, Copy, PartialEq)]
117pub struct DataInterval {
118    pub chr_index: u32,
119    pub start: i64,
120    pub end: i64,
121    pub value: f32,
122    pub valid_count: i64,
123    pub min_value: f32,
124    pub max_value: f32,
125    pub sum_squared: f64,
126}
127
128/// Read item `index` of a wig section.
129///
130/// Pulled out of the walk so the writer's encoder has exactly one decoder to be
131/// the inverse of, and so the zoom pass can walk a section without standing up
132/// a set of loci for it.
133pub fn read_wig_item(
134    block: &[u8],
135    header: &WigSectionHeader,
136    index: usize,
137    path: &str,
138) -> Result<DataInterval> {
139    let item_size = header.encoding.item_size();
140    let offset = WIG_HEADER_SIZE + index * item_size;
141    let mut c = LeCursor::new(block, 0, path);
142    c.seek(offset)?;
143    let (start, end, value) = match header.encoding {
144        WigEncoding::BedGraph => {
145            let start = c.read_u32()? as i64;
146            let end = c.read_u32()? as i64;
147            (start, end, c.read_f32()?)
148        }
149        WigEncoding::VarStep => {
150            let start = c.read_u32()? as i64;
151            (start, start + header.item_span, c.read_f32()?)
152        }
153        WigEncoding::FixedStep => {
154            let start = header.chr_start + index as i64 * header.item_step;
155            (start, start + header.item_span, c.read_f32()?)
156        }
157    };
158    let valid_count = end - start;
159    Ok(DataInterval {
160        chr_index: header.chr_index,
161        start,
162        end,
163        value,
164        valid_count,
165        // One value over the whole range, so it is its own minimum and maximum
166        // and the squares sum to it repeated over every base it covers.
167        min_value: value,
168        max_value: value,
169        sum_squared: value as f64 * value as f64 * valid_count as f64,
170    })
171}
172
173/// A zoom record: a pre-summarised window.
174#[derive(Debug, Clone, Copy)]
175pub struct ZoomRecord {
176    pub chr_index: u32,
177    pub chr_start: i64,
178    pub chr_end: i64,
179    pub valid_count: i64,
180    pub min_value: f32,
181    pub max_value: f32,
182    pub sum_data: f32,
183    pub sum_squared: f32,
184}
185
186pub fn read_zoom_record(block: &[u8], offset: usize, path: &str) -> Result<ZoomRecord> {
187    let mut c = LeCursor::new(block, 0, path);
188    c.seek(offset)?;
189    Ok(ZoomRecord {
190        chr_index: c.read_u32()?,
191        chr_start: c.read_u32()? as i64,
192        chr_end: c.read_u32()? as i64,
193        valid_count: c.read_u32()? as i64,
194        min_value: c.read_f32()?,
195        max_value: c.read_f32()?,
196        sum_data: c.read_f32()?,
197        sum_squared: c.read_f32()?,
198    })
199}
200
201/// A bigBed entry. Always carries its coordinates; the rest of the columns are
202/// present only up to the request's `col_count`.
203#[derive(Debug, Clone)]
204pub struct BedEntry {
205    pub chr: String,
206    pub start: i64,
207    pub end: i64,
208    /// The remaining tab-separated columns, in file order, named by the file's
209    /// autoSql. Empty when `col_count` was 3.
210    pub fields: Vec<(String, String)>,
211}
212
213/// How much one block is allowed to inflate to.
214///
215/// A hard limit, not a hint. Without it a block whose deflate stream expands
216/// far beyond what any real writer produces — a corrupt file, or a deliberate
217/// one — is inflated until the process runs out of memory, and Rust aborts on
218/// a failed allocation rather than returning an error.
219const MAX_INFLATED_SIZE: usize = 1 << 30;
220
221/// How much of the output buffer is reserved up front, whatever the header
222/// says. The declared size is a hint from the file and a corrupt one can name
223/// four gigabytes; reserving that aborts the process before a single byte is
224/// inflated. Reserving less costs a few reallocations on a block that really is
225/// larger, and `read_to_end` grows past it either way.
226const MAX_INFLATE_RESERVE: usize = 1 << 20;
227
228/// Inflate a data block.
229///
230/// `uncompress_buffer_size` of 0 means the file stores blocks uncompressed, and
231/// the block is handed back as it is. Otherwise it is the size the writer
232/// declared its largest block inflates to, which is what the output buffer is
233/// sized from — but it is a hint from the file, so the decoder is allowed to
234/// exceed it rather than truncating silently.
235pub fn decompress(block: Bytes, uncompress_buffer_size: u32, path: &str) -> Result<Bytes> {
236    decompress_limited(block, uncompress_buffer_size, path, MAX_INFLATED_SIZE)
237}
238
239/// The same with the cap as an argument, so the tests can reach it without
240/// inflating a gibibyte to prove a limit they can prove with a kilobyte.
241fn decompress_limited(
242    block: Bytes,
243    uncompress_buffer_size: u32,
244    path: &str,
245    max_size: usize,
246) -> Result<Bytes> {
247    if uncompress_buffer_size == 0 {
248        return Ok(block);
249    }
250    use std::io::Read;
251    let reserve = (uncompress_buffer_size as usize).min(MAX_INFLATE_RESERVE.min(max_size));
252    let mut out = Vec::with_capacity(reserve);
253    // `take` is what enforces the cap: the decoder stops one byte past the
254    // limit and the length says whether it got there, so a stream that would
255    // have gone on is refused while inflating rather than after.
256    flate2::read::ZlibDecoder::new(&block[..])
257        .take(max_size as u64 + 1)
258        .read_to_end(&mut out)
259        .map_err(|e| Error::corrupt(path, 0, format!("could not inflate data block: {e}")))?;
260    if out.len() > max_size {
261        return Err(Error::corrupt(
262            path,
263            0,
264            format!("decompressed data exceeds limit ({max_size})"),
265        ));
266    }
267    Ok(Bytes::from(out))
268}
269
270/// The bounds of the loci a block is being read for, as (chromosome, base)
271/// pairs.
272///
273/// A block groups records by count rather than by chromosome, so comparing bare
274/// coordinates would drop data at every boundary a block straddles.
275#[derive(Debug, Clone, Copy)]
276struct LocBounds {
277    min_chr: u32,
278    min_start: i64,
279    max_chr: u32,
280    max_end: i64,
281}
282
283impl LocBounds {
284    fn of(locs: &[IndexedLoc], range: std::ops::Range<usize>) -> Self {
285        let first = &locs[range.start];
286        let last = &locs[range.end - 1];
287        // Loci are sorted by chromosome then position, so the last one sits on
288        // the highest chromosome; the furthest reach on that chromosome is the
289        // widest end among the loci that share it.
290        let max_chr = last.chr_index as u32;
291        let mut max_end = last.binned_end;
292        for loc in locs[range.start..range.end - 1].iter().rev() {
293            if loc.chr_index as u32 != max_chr {
294                break;
295            }
296            max_end = max_end.max(loc.binned_end);
297        }
298        Self {
299            min_chr: first.chr_index as u32,
300            min_start: first.binned_start,
301            max_chr,
302            max_end,
303        }
304    }
305}
306
307/// Walks a block's items as [`DataInterval`]s, skipping the ones no locus of
308/// the batch can reach and stopping once the block has run past them all.
309#[derive(Debug)]
310pub struct DataIntervals {
311    block: Bytes,
312    path: String,
313    bounds: LocBounds,
314    /// `None` for full data, `Some(_)` for a zoom block.
315    header: Option<WigSectionHeader>,
316    count: usize,
317    index: usize,
318}
319
320impl DataIntervals {
321    /// `zoom` selects the record layout: zoom blocks are a bare run of 32-byte
322    /// records, full-data blocks a wig section header followed by its items.
323    pub fn new(
324        block: Bytes,
325        zoom: bool,
326        locs: &[IndexedLoc],
327        range: std::ops::Range<usize>,
328        path: &str,
329    ) -> Result<Self> {
330        let (header, count) = if zoom {
331            // Counted from the block as it stands: `data_size` is what the
332            // block occupies on disk, which is the compressed size on every
333            // file that sets `uncompress_buffer_size`.
334            (None, block.len() / ZOOM_RECORD_SIZE)
335        } else {
336            let header = read_wig_header(&block, path)?;
337            let count = header.item_count as usize;
338            let item_size = header.encoding.item_size();
339            // Checked once here rather than per item, so a block whose declared
340            // item count does not fit it fails cleanly.
341            if WIG_HEADER_SIZE + count * item_size > block.len() {
342                return Err(Error::corrupt(
343                    path,
344                    0,
345                    format!(
346                        "wig section declares {count} items of type {:?}, which do not fit \
347                         its {} byte block",
348                        header.encoding,
349                        block.len()
350                    ),
351                ));
352            }
353            (Some(header), count)
354        };
355        Ok(Self {
356            block,
357            path: path.to_string(),
358            bounds: LocBounds::of(locs, range),
359            header,
360            count,
361            index: 0,
362        })
363    }
364}
365
366impl Iterator for DataIntervals {
367    type Item = Result<DataInterval>;
368
369    fn next(&mut self) -> Option<Self::Item> {
370        while self.index < self.count {
371            let index = self.index;
372            self.index += 1;
373
374            let data = match self.header {
375                Some(header) => match read_wig_item(&self.block, &header, index, &self.path) {
376                    Ok(d) => d,
377                    Err(e) => return Some(Err(e)),
378                },
379                None => {
380                    let record =
381                        match read_zoom_record(&self.block, index * ZOOM_RECORD_SIZE, &self.path) {
382                            Ok(r) => r,
383                            Err(e) => return Some(Err(e)),
384                        };
385                    // A window the file recorded but valued from nothing has no
386                    // mean to give.
387                    if record.valid_count == 0 {
388                        continue;
389                    }
390                    DataInterval {
391                        chr_index: record.chr_index,
392                        start: record.chr_start,
393                        end: record.chr_end,
394                        value: record.sum_data / record.valid_count as f32,
395                        valid_count: record.valid_count,
396                        min_value: record.min_value,
397                        max_value: record.max_value,
398                        sum_squared: record.sum_squared as f64,
399                    }
400                }
401            };
402
403            let b = &self.bounds;
404            if data.chr_index < b.min_chr {
405                continue;
406            }
407            if data.chr_index == b.min_chr && data.end <= b.min_start {
408                continue;
409            }
410            if data.chr_index > b.max_chr {
411                break;
412            }
413            if data.chr_index == b.max_chr && data.start >= b.max_end {
414                break;
415            }
416            return Some(Ok(data));
417        }
418        None
419    }
420}
421
422/// Walk a bigBed block's records (Supp. Table 12), keeping the loci-relevant
423/// ones and parsing only the columns asked for.
424///
425/// A record is `u32 chromIx, u32 start, u32 end` followed by a NUL-terminated,
426/// tab-separated tail. The fields past `kept` are counted but never built,
427/// which is what makes a `col_count`-limited read cheaper than a full one
428/// rather than merely narrower.
429#[derive(Debug)]
430pub struct BedRecords<'a> {
431    block: Bytes,
432    path: &'a str,
433    /// The file's column names, coordinates included. The tail's fields are
434    /// named by the entries past the first three.
435    auto_sql: &'a indexmap::IndexMap<String, String>,
436    /// Tail fields the file declares, i.e. `auto_sql.len() - 3`.
437    field_count: usize,
438    /// Tail fields to build.
439    kept: usize,
440    bounds: LocBounds,
441    offset: usize,
442}
443
444impl<'a> BedRecords<'a> {
445    /// `col_count` counts the three coordinate columns; 0 keeps everything.
446    /// Clamped rather than checked — the readers validate what they were handed
447    /// against the file's own column count before they get here.
448    pub fn new(
449        block: Bytes,
450        auto_sql: &'a indexmap::IndexMap<String, String>,
451        col_count: usize,
452        locs: &[IndexedLoc],
453        range: std::ops::Range<usize>,
454        path: &'a str,
455    ) -> Result<Self> {
456        if auto_sql.len() < 3 {
457            return Err(Error::format(
458                path,
459                format!(
460                    "bed entries need the 3 standard fields, autosql describes {}",
461                    auto_sql.len()
462                ),
463            ));
464        }
465        let field_count = auto_sql.len() - 3;
466        let kept = if col_count == 0 {
467            field_count
468        } else {
469            field_count.min(col_count.saturating_sub(3))
470        };
471        Ok(Self {
472            block,
473            path,
474            auto_sql,
475            field_count,
476            kept,
477            bounds: LocBounds::of(locs, range),
478            offset: 0,
479        })
480    }
481
482    /// Read the tab-separated tail into named fields.
483    ///
484    /// A trailing delimiter closes a field of its own, so `"name\t0\t"` is three
485    /// fields, the last blank — dropping it would fail the record for a field
486    /// count mismatch it does not have. By the same count an empty tail is one
487    /// blank field, since n fields carry n-1 delimiters; that is also what a
488    /// bed3 record looks like, and the two cannot be told apart from the bytes.
489    /// So the autoSql decides, and only a file declaring no tail fields reads an
490    /// empty tail as none.
491    fn read_fields(&self, tail: &[u8]) -> Result<Vec<(String, String)>> {
492        let mut fields = Vec::with_capacity(self.kept);
493        let mut found = 0usize;
494        if !tail.is_empty() || self.field_count > 0 {
495            for part in tail.split(|b| *b == b'\t') {
496                if found < self.kept {
497                    let name = self
498                        .auto_sql
499                        .get_index(3 + found)
500                        .map(|(k, _)| k.clone())
501                        .unwrap_or_else(|| format!("field{}", 4 + found));
502                    fields.push((name, String::from_utf8_lossy(part).into_owned()));
503                }
504                found += 1;
505            }
506        }
507        if found != self.field_count {
508            return Err(Error::corrupt(
509                self.path,
510                self.offset as u64,
511                format!(
512                    "invalid bed entry (found {found} fields past the first 3, \
513                     autosql declares {})",
514                    self.field_count
515                ),
516            ));
517        }
518        Ok(fields)
519    }
520}
521
522impl Iterator for BedRecords<'_> {
523    type Item = Result<(u32, i64, i64, Vec<(String, String)>)>;
524
525    fn next(&mut self) -> Option<Self::Item> {
526        while self.offset < self.block.len() {
527            if self.offset + BED_RECORD_HEADER_SIZE > self.block.len() {
528                return Some(Err(Error::corrupt(
529                    self.path,
530                    self.offset as u64,
531                    format!(
532                        "truncated bed record at {} ({} bytes left in its block)",
533                        self.offset,
534                        self.block.len() - self.offset
535                    ),
536                )));
537            }
538            let head = &self.block[self.offset..self.offset + BED_RECORD_HEADER_SIZE];
539            let chr_index = u32::from_le_bytes([head[0], head[1], head[2], head[3]]);
540            let start = u32::from_le_bytes([head[4], head[5], head[6], head[7]]) as i64;
541            let end = u32::from_le_bytes([head[8], head[9], head[10], head[11]]) as i64;
542
543            let tail_start = self.offset + BED_RECORD_HEADER_SIZE;
544            let Some(nul) = memchr::memchr(0, &self.block[tail_start..]) else {
545                return Some(Err(Error::corrupt(
546                    self.path,
547                    tail_start as u64,
548                    "invalid bed entry (null terminator not found)",
549                )));
550            };
551            let tail_end = tail_start + nul;
552
553            // Bounds first: a record no locus of this batch can reach costs one
554            // comparison rather than a field split.
555            //
556            // `reach` rather than `end`: BED allows `start == end` — an
557            // insertion — and a half-open interval of no width overlaps
558            // nothing at all, so such an entry would be dropped by every test
559            // it met. It occupies the single base it names instead, which is
560            // the rule `advance_cursor` and `walk_bed_batch` also apply.
561            let b = &self.bounds;
562            let reach = end.max(start + 1);
563            let reachable =
564                !(chr_index < b.min_chr || (chr_index == b.min_chr && reach <= b.min_start));
565            let past = chr_index > b.max_chr || (chr_index == b.max_chr && start >= b.max_end);
566
567            let fields = if reachable && !past {
568                match self.read_fields(&self.block[tail_start..tail_end]) {
569                    Ok(f) => Some(f),
570                    Err(e) => return Some(Err(e)),
571                }
572            } else {
573                None
574            };
575            self.offset = tail_end + 1;
576
577            if past {
578                break;
579            }
580            if let Some(fields) = fields {
581                return Some(Ok((chr_index, start, end, fields)));
582            }
583        }
584        None
585    }
586}
587
588/// Call `visit(chr_index, start, end)` for every record of a bed block,
589/// reading the coordinates alone.
590///
591/// What the zoom pass wants: the tail is skipped rather than split, which is all
592/// the coverage a bigBed's summaries need.
593pub fn visit_bed_records(
594    block: &[u8],
595    path: &str,
596    mut visit: impl FnMut(u32, i64, i64),
597) -> Result<()> {
598    let mut offset = 0usize;
599    while offset < block.len() {
600        if offset + BED_RECORD_HEADER_SIZE > block.len() {
601            return Err(Error::corrupt(
602                path,
603                offset as u64,
604                format!(
605                    "truncated bed record at {offset} ({} bytes left in its block)",
606                    block.len() - offset
607                ),
608            ));
609        }
610        let head = &block[offset..offset + BED_RECORD_HEADER_SIZE];
611        let chr_index = u32::from_le_bytes([head[0], head[1], head[2], head[3]]);
612        let start = u32::from_le_bytes([head[4], head[5], head[6], head[7]]) as i64;
613        let end = u32::from_le_bytes([head[8], head[9], head[10], head[11]]) as i64;
614        let tail_start = offset + BED_RECORD_HEADER_SIZE;
615        let Some(nul) = memchr::memchr(0, &block[tail_start..]) else {
616            return Err(Error::corrupt(
617                path,
618                tail_start as u64,
619                "invalid bed entry (null terminator not found)",
620            ));
621        };
622        offset = tail_start + nul + 1;
623        visit(chr_index, start, end);
624    }
625    Ok(())
626}
627
628#[cfg(test)]
629mod tests {
630    use super::*;
631
632    /// Deflate `bytes` the way a bbi writer would, so `decompress` can read it.
633    fn deflated(bytes: &[u8]) -> Vec<u8> {
634        use std::io::Write as _;
635        let mut e = flate2::write::ZlibEncoder::new(Vec::new(), flate2::Compression::best());
636        e.write_all(bytes).unwrap();
637        e.finish().unwrap()
638    }
639
640    #[test]
641    fn a_declared_buffer_size_is_a_hint_and_not_an_allocation() {
642        // Four gigabytes, which is what a mutated header can name and what
643        // reserving would abort the process for. The block is 32 bytes.
644        let block = bytes::Bytes::from(deflated(&[7u8; 32]));
645        let out = decompress(block, u32::MAX, "corrupt.bigwig").unwrap();
646        assert_eq!(out.len(), 32);
647        assert!(out.iter().all(|b| *b == 7));
648    }
649
650    #[test]
651    fn a_block_that_inflates_past_the_limit_is_refused_rather_than_read() {
652        // The cheap half of a decompression bomb: a run of zeros deflates to
653        // almost nothing, so refusing it has to happen while inflating rather
654        // than after. The real cap is a gibibyte; proving the rule does not
655        // need one.
656        let block = bytes::Bytes::from(deflated(&vec![0u8; 4097]));
657        let err = decompress_limited(block, 4096, "bomb.bigwig", 4096)
658            .unwrap_err()
659            .to_string();
660        assert!(err.contains("exceeds limit (4096)"), "{err}");
661    }
662
663    #[test]
664    fn a_block_exactly_at_the_limit_is_still_read() {
665        let block = bytes::Bytes::from(deflated(&vec![0u8; 4096]));
666        assert_eq!(
667            decompress_limited(block, 4096, "big.bigwig", 4096)
668                .unwrap()
669                .len(),
670            4096
671        );
672    }
673
674    fn loc(chr: usize, start: i64, end: i64) -> IndexedLoc {
675        IndexedLoc {
676            chr_index: chr,
677            start,
678            end,
679            binned_start: start,
680            binned_end: end,
681            bin_size: 1.0,
682            reverse: false,
683            output_start: 0,
684            output_end: (end - start) as usize,
685        }
686    }
687
688    fn wig_block(
689        encoding: WigEncoding,
690        items: &[(u32, u32, f32)],
691        span: u32,
692        step: u32,
693    ) -> Vec<u8> {
694        let mut b = Vec::new();
695        b.extend_from_slice(&7u32.to_le_bytes()); // chromId
696        b.extend_from_slice(&items[0].0.to_le_bytes()); // chromStart
697        b.extend_from_slice(&items[items.len() - 1].1.to_le_bytes()); // chromEnd
698        b.extend_from_slice(&step.to_le_bytes());
699        b.extend_from_slice(&span.to_le_bytes());
700        b.push(encoding as u8);
701        b.push(0);
702        b.extend_from_slice(&(items.len() as u16).to_le_bytes());
703        assert_eq!(b.len(), WIG_HEADER_SIZE);
704        for (start, end, value) in items {
705            match encoding {
706                WigEncoding::BedGraph => {
707                    b.extend_from_slice(&start.to_le_bytes());
708                    b.extend_from_slice(&end.to_le_bytes());
709                    b.extend_from_slice(&value.to_le_bytes());
710                }
711                WigEncoding::VarStep => {
712                    b.extend_from_slice(&start.to_le_bytes());
713                    b.extend_from_slice(&value.to_le_bytes());
714                }
715                WigEncoding::FixedStep => b.extend_from_slice(&value.to_le_bytes()),
716            }
717        }
718        b
719    }
720
721    fn collect(block: Vec<u8>, zoom: bool, locs: &[IndexedLoc]) -> Vec<DataInterval> {
722        DataIntervals::new(Bytes::from(block), zoom, locs, 0..locs.len(), "test")
723            .unwrap()
724            .map(|r| r.unwrap())
725            .collect()
726    }
727
728    #[test]
729    fn the_three_encodings_decode_to_the_same_intervals() {
730        let items = [(100u32, 110u32, 1.5f32), (110, 120, 2.5), (120, 130, 3.5)];
731        let locs = [loc(7, 0, 1000)];
732
733        let bg = collect(
734            wig_block(WigEncoding::BedGraph, &items, 10, 10),
735            false,
736            &locs,
737        );
738        let vs = collect(
739            wig_block(WigEncoding::VarStep, &items, 10, 10),
740            false,
741            &locs,
742        );
743        let fs = collect(
744            wig_block(WigEncoding::FixedStep, &items, 10, 10),
745            false,
746            &locs,
747        );
748
749        assert_eq!(bg.len(), 3);
750        assert_eq!(bg, vs);
751        assert_eq!(bg, fs);
752        assert_eq!(bg[0].start, 100);
753        assert_eq!(bg[0].end, 110);
754        assert_eq!(bg[0].value, 1.5);
755        assert_eq!(bg[0].valid_count, 10);
756        assert_eq!(bg[0].min_value, 1.5);
757        assert_eq!(bg[0].sum_squared, 1.5f64 * 1.5 * 10.0);
758        assert_eq!(bg[0].chr_index, 7);
759    }
760
761    #[test]
762    fn items_outside_the_batchs_bounds_are_skipped_and_the_walk_stops() {
763        let items = [
764            (0u32, 10u32, 1.0f32),
765            (10, 20, 2.0),
766            (100, 110, 3.0),
767            (200, 210, 4.0),
768            (300, 310, 5.0),
769        ];
770        // Only [100, 210) is wanted.
771        let locs = [loc(7, 100, 210)];
772        let got = collect(
773            wig_block(WigEncoding::BedGraph, &items, 10, 10),
774            false,
775            &locs,
776        );
777        assert_eq!(got.iter().map(|d| d.start).collect::<Vec<_>>(), [100, 200]);
778    }
779
780    #[test]
781    fn a_block_straddling_two_chromosomes_keeps_only_the_reachable_side() {
782        // Items on chr 7 then chr 8; the batch only wants chr 8.
783        let mut b = wig_block(WigEncoding::BedGraph, &[(0, 10, 1.0)], 10, 10);
784        b[0..4].copy_from_slice(&8u32.to_le_bytes());
785        let locs = [loc(8, 0, 10)];
786        assert_eq!(collect(b.clone(), false, &locs).len(), 1);
787        let locs = [loc(9, 0, 10)];
788        assert!(collect(b, false, &locs).is_empty());
789    }
790
791    #[test]
792    fn zoom_records_become_intervals_and_zero_count_ones_are_dropped() {
793        let mut b = Vec::new();
794        for (start, end, valid, min, max, sum, sq) in [
795            (0u32, 100u32, 100u32, 1.0f32, 5.0f32, 300.0f32, 1000.0f32),
796            (100, 200, 0, 0.0, 0.0, 0.0, 0.0), // recorded but valued from nothing
797            (200, 300, 50, 2.0, 4.0, 150.0, 500.0),
798        ] {
799            b.extend_from_slice(&7u32.to_le_bytes());
800            b.extend_from_slice(&start.to_le_bytes());
801            b.extend_from_slice(&end.to_le_bytes());
802            b.extend_from_slice(&valid.to_le_bytes());
803            b.extend_from_slice(&min.to_le_bytes());
804            b.extend_from_slice(&max.to_le_bytes());
805            b.extend_from_slice(&sum.to_le_bytes());
806            b.extend_from_slice(&sq.to_le_bytes());
807        }
808        let locs = [loc(7, 0, 1000)];
809        let got = collect(b, true, &locs);
810        assert_eq!(got.len(), 2);
811        // value is the window mean, not the stored sum.
812        assert_eq!(got[0].value, 3.0);
813        assert_eq!(got[0].valid_count, 100);
814        assert_eq!(got[0].min_value, 1.0);
815        assert_eq!(got[0].max_value, 5.0);
816        assert_eq!(got[1].value, 3.0);
817        assert_eq!(got[1].valid_count, 50);
818    }
819
820    #[test]
821    fn a_bad_wig_type_is_corrupt_not_a_panic() {
822        let mut b = wig_block(WigEncoding::BedGraph, &[(0, 10, 1.0)], 10, 10);
823        b[20] = 9;
824        let locs = [loc(7, 0, 10)];
825        let err = DataIntervals::new(Bytes::from(b), false, &locs, 0..1, "test").unwrap_err();
826        assert!(err.to_string().contains("wig data type 9 invalid"), "{err}");
827    }
828
829    #[test]
830    fn a_block_declaring_more_items_than_it_holds_is_refused_up_front() {
831        let mut b = wig_block(WigEncoding::BedGraph, &[(0, 10, 1.0)], 10, 10);
832        b[22..24].copy_from_slice(&500u16.to_le_bytes());
833        let locs = [loc(7, 0, 10)];
834        let err = DataIntervals::new(Bytes::from(b), false, &locs, 0..1, "test").unwrap_err();
835        assert!(err.to_string().contains("do not fit"), "{err}");
836    }
837
838    #[test]
839    fn a_block_too_short_for_a_header_is_refused() {
840        let locs = [loc(7, 0, 10)];
841        let err =
842            DataIntervals::new(Bytes::from(vec![0u8; 8]), false, &locs, 0..1, "test").unwrap_err();
843        assert!(err.to_string().contains("header needs 24 bytes"), "{err}");
844    }
845
846    #[test]
847    fn an_uncompressed_file_hands_its_block_back_untouched() {
848        let block = Bytes::from(vec![1u8, 2, 3]);
849        assert_eq!(decompress(block.clone(), 0, "test").unwrap(), block);
850    }
851
852    #[test]
853    fn a_compressed_block_round_trips() {
854        use std::io::Write;
855        let raw: Vec<u8> = (0..5000).map(|i| (i % 251) as u8).collect();
856        let mut encoder = flate2::write::ZlibEncoder::new(Vec::new(), flate2::Compression::new(6));
857        encoder.write_all(&raw).unwrap();
858        let compressed = Bytes::from(encoder.finish().unwrap());
859        assert_eq!(&decompress(compressed, 8192, "test").unwrap()[..], &raw[..]);
860    }
861
862    #[test]
863    fn garbage_where_a_deflate_stream_should_be_is_corrupt() {
864        let err = decompress(Bytes::from(vec![9u8; 64]), 4096, "test").unwrap_err();
865        assert!(err.to_string().contains("could not inflate"), "{err}");
866    }
867
868    #[test]
869    fn bounds_take_the_widest_end_on_the_last_chromosome() {
870        // Sorted by (chr, start): the last locus is not the one reaching
871        // furthest, which is the case a naive `last.binned_end` gets wrong.
872        let locs = [loc(1, 0, 10), loc(2, 0, 500), loc(2, 100, 200)];
873        let bounds = LocBounds::of(&locs, 0..3);
874        assert_eq!(bounds.min_chr, 1);
875        assert_eq!(bounds.min_start, 0);
876        assert_eq!(bounds.max_chr, 2);
877        assert_eq!(bounds.max_end, 500);
878    }
879}