use crate::bbi::block::WigEncoding;
use crate::bbi::header::to_bbi_u32;
use crate::error::Result;
pub const MAX_ITEMS_PER_SECTION: usize = 65535;
pub const SECTION_SPLIT_COST: i64 = 64;
pub const DOWNGRADE_COMPRESSION_PERCENT: i64 = 25;
pub const MIN_SPLIT_ITEMS: i64 = 32;
pub const RUNT_PATIENCE: u32 = 2;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum SectionPolicy {
Cost,
Split,
Widen,
PrefixOnly,
Room,
Runt,
#[default]
Adaptive,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct CostModel {
pub split_cost: i64,
pub compression_percent: i64,
pub min_split_items: i64,
pub runt_patience: u32,
}
impl Default for CostModel {
fn default() -> Self {
Self {
split_cost: SECTION_SPLIT_COST,
compression_percent: DOWNGRADE_COMPRESSION_PERCENT,
min_split_items: MIN_SPLIT_ITEMS,
runt_patience: RUNT_PATIENCE,
}
}
}
pub const WIG_HEADER_SIZE: usize = 24;
#[derive(Debug, Default)]
pub struct WigSection {
pub chr_ix: u32,
pub encoding: Option<WigEncoding>,
pub first_start: i64,
pub uniform_step: Option<i64>,
pub uniform_span: i64,
pub last_start: i64,
pub last_end: i64,
pub starts: Vec<i64>,
pub ends: Vec<i64>,
pub values: Vec<f32>,
items_per_slot: usize,
policy: SectionPolicy,
cost: CostModel,
runt_streak: u32,
}
impl WigSection {
pub fn with_policy(items_per_slot: usize, policy: SectionPolicy, cost: CostModel) -> Self {
Self {
items_per_slot,
policy,
cost,
..Default::default()
}
}
pub fn len(&self) -> usize {
self.values.len()
}
pub fn is_empty(&self) -> bool {
self.values.is_empty()
}
pub fn is_full(&self) -> bool {
self.values.len() >= self.items_per_slot
}
pub fn encoding(&self) -> WigEncoding {
self.encoding.unwrap_or(WigEncoding::FixedStep)
}
pub fn clear(&mut self) {
if !self.is_empty() {
if (self.len() as i64) < self.cost.min_split_items {
self.runt_streak = self.runt_streak.saturating_add(1);
} else if self.encoding() == WigEncoding::FixedStep {
self.runt_streak = 0;
}
}
self.chr_ix = 0;
self.encoding = None;
self.first_start = -1;
self.uniform_step = None;
self.uniform_span = -1;
self.last_start = -1;
self.last_end = -1;
self.starts.clear();
self.ends.clear();
self.values.clear();
}
fn should_widen(&self, old: WigEncoding, new: WigEncoding, batch_remaining: usize) -> bool {
match self.policy {
SectionPolicy::Split => return false,
SectionPolicy::Widen => return true,
SectionPolicy::Cost
| SectionPolicy::PrefixOnly
| SectionPolicy::Room
| SectionPolicy::Runt
| SectionPolicy::Adaptive => {}
}
let count = self.len() as i64;
if self.policy == SectionPolicy::Runt && count < self.cost.min_split_items {
return true;
}
if self.policy == SectionPolicy::Adaptive
&& count < self.cost.min_split_items
&& self.runt_streak >= self.cost.runt_patience
{
return true;
}
let (old_size, new_size) = (old.item_size() as i64, new.item_size() as i64);
let prefix_cost = (new_size - old_size) * count;
let room = self.items_per_slot as i64 - count;
let widened_items = match self.policy {
SectionPolicy::PrefixOnly => 0,
SectionPolicy::Cost => (batch_remaining as i64).min(room),
_ => room,
};
let tail_cost = (new_size - WigEncoding::FixedStep.item_size() as i64) * widened_items;
(prefix_cost + tail_cost) * self.cost.compression_percent / 100 < self.cost.split_cost
}
fn widen(&mut self, target: WigEncoding) {
let count = self.len();
if self.encoding() == WigEncoding::FixedStep && target != WigEncoding::FixedStep {
let step = if count >= 2 {
self.uniform_step.unwrap_or(self.uniform_span)
} else {
self.uniform_span
};
self.starts.clear();
self.starts
.extend((0..count).map(|i| self.first_start + i as i64 * step));
self.encoding = Some(WigEncoding::VarStep);
}
if self.encoding() == WigEncoding::VarStep && target == WigEncoding::BedGraph {
self.ends.clear();
let span = self.uniform_span;
self.ends.extend(self.starts.iter().map(|s| s + span));
self.encoding = Some(WigEncoding::BedGraph);
}
}
fn push(&mut self, start: i64, end: i64, value: f32) {
self.values.push(value);
if self.encoding() != WigEncoding::FixedStep {
self.starts.push(start);
}
if self.encoding() == WigEncoding::BedGraph {
self.ends.push(end);
}
self.last_start = start;
self.last_end = end;
}
fn open(&mut self, chr_ix: u32, start: i64, span: i64) {
self.chr_ix = chr_ix;
self.first_start = start;
self.uniform_span = span;
self.uniform_step = None;
self.encoding = Some(WigEncoding::FixedStep);
}
pub fn offer(
&mut self,
chr_ix: u32,
start: i64,
span: i64,
value: f32,
batch_remaining: usize,
) -> Accept {
if !self.is_empty() && chr_ix != self.chr_ix {
return Accept::Flush;
}
if self.is_empty() {
self.open(chr_ix, start, span);
self.push(start, start + span, value);
return Accept::Buffered;
}
let count = self.len();
let mut needed = WigEncoding::FixedStep;
if span != self.uniform_span {
needed = WigEncoding::BedGraph;
} else if count >= 2 && Some(start - self.last_start) != self.uniform_step {
needed = WigEncoding::VarStep;
}
if needed.item_size() > self.encoding().item_size() {
if !self.should_widen(self.encoding(), needed, batch_remaining) {
return Accept::Flush;
}
self.widen(needed);
}
if count == 1 && self.encoding() == WigEncoding::FixedStep {
self.uniform_step = Some(start - self.first_start);
}
self.push(start, start + span, value);
Accept::Buffered
}
pub fn should_flush_for_run(&self, chr_ix: u32, start: i64, span: i64, run_len: usize) -> bool {
if !matches!(self.policy, SectionPolicy::Adaptive) {
return false;
}
if self.is_empty() || self.extends_run(chr_ix, start, span) {
return false;
}
if self.chr_ix != chr_ix {
return false;
}
let poured = WigEncoding::BedGraph.item_size() as i64;
let fresh = WigEncoding::FixedStep.item_size() as i64;
let room = (self.items_per_slot - self.len()) as i64;
let saved = (poured - fresh) * (run_len as i64).min(room);
saved * self.cost.compression_percent / 100 > self.cost.split_cost
}
pub fn extends_run(&self, chr_ix: u32, start: i64, span: i64) -> bool {
self.encoding() == WigEncoding::FixedStep
&& !self.is_empty()
&& self.chr_ix == chr_ix
&& self.uniform_span == span
&& (self.len() == 1 || self.uniform_step == Some(span))
&& self.last_end == start
}
pub fn extend_run(&mut self, chr_ix: u32, start: i64, span: i64, values: &[f32]) {
debug_assert!(!values.is_empty());
if self.is_empty() {
self.open(chr_ix, start, span);
}
self.uniform_step = Some(span);
self.values.extend_from_slice(values);
self.last_start = start + span * (values.len() as i64 - 1);
self.last_end = start + span * values.len() as i64;
}
pub fn encode(&self) -> Result<Vec<u8>> {
let count = self.len();
let encoding = self.encoding();
let mut out = Vec::with_capacity(WIG_HEADER_SIZE + count * encoding.item_size());
out.extend_from_slice(&to_bbi_u32(self.chr_ix as i64, "chromId")?.to_le_bytes());
out.extend_from_slice(&to_bbi_u32(self.first_start, "chromStart")?.to_le_bytes());
out.extend_from_slice(&to_bbi_u32(self.last_end, "chromEnd")?.to_le_bytes());
let step = if encoding == WigEncoding::FixedStep {
if count >= 2 {
self.uniform_step.unwrap_or(self.uniform_span)
} else {
self.uniform_span
}
} else {
0
};
out.extend_from_slice(&to_bbi_u32(step, "itemStep")?.to_le_bytes());
let span = if encoding == WigEncoding::BedGraph {
0
} else {
self.uniform_span
};
out.extend_from_slice(&to_bbi_u32(span, "itemSpan")?.to_le_bytes());
out.push(encoding as u8);
out.push(0); if count > MAX_ITEMS_PER_SECTION {
return Err(crate::error::Error::invalid(format!(
"wig section of {count} items exceeds the {MAX_ITEMS_PER_SECTION} its item \
count is stored on"
)));
}
out.extend_from_slice(&(count as u16).to_le_bytes());
debug_assert_eq!(out.len(), WIG_HEADER_SIZE);
match encoding {
WigEncoding::FixedStep => {
for value in &self.values {
out.extend_from_slice(&value.to_le_bytes());
}
}
WigEncoding::VarStep => {
for i in 0..count {
out.extend_from_slice(&to_bbi_u32(self.starts[i], "chromStart")?.to_le_bytes());
out.extend_from_slice(&self.values[i].to_le_bytes());
}
}
WigEncoding::BedGraph => {
for i in 0..count {
out.extend_from_slice(&to_bbi_u32(self.starts[i], "chromStart")?.to_le_bytes());
out.extend_from_slice(&to_bbi_u32(self.ends[i], "chromEnd")?.to_le_bytes());
out.extend_from_slice(&self.values[i].to_le_bytes());
}
}
}
Ok(out)
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Accept {
Buffered,
Flush,
}
#[cfg(test)]
mod tests {
use super::*;
fn section() -> WigSection {
WigSection::with_policy(1024, SectionPolicy::default(), CostModel::default())
}
fn cost_section() -> WigSection {
WigSection::with_policy(1024, SectionPolicy::Cost, CostModel::default())
}
fn widening_section() -> WigSection {
WigSection::with_policy(1024, SectionPolicy::Widen, CostModel::default())
}
fn offer(s: &mut WigSection, start: i64, span: i64, value: f32, remaining: usize) -> bool {
match s.offer(0, start, span, value, remaining) {
Accept::Buffered => false,
Accept::Flush => {
s.clear();
assert_eq!(s.offer(0, start, span, value, remaining), Accept::Buffered);
true
}
}
}
#[test]
fn a_regular_run_stays_fixed_step() {
let mut s = section();
for i in 0..10 {
assert!(!offer(&mut s, i * 10, 10, i as f32, 0));
}
assert_eq!(s.encoding(), WigEncoding::FixedStep);
assert_eq!(s.len(), 10);
assert!(s.starts.is_empty() && s.ends.is_empty());
assert_eq!(s.uniform_step, Some(10));
}
#[test]
fn a_gap_widens_to_var_step_when_the_widening_is_cheap() {
let mut s = cost_section();
for i in 0..4 {
offer(&mut s, i * 10, 10, 1.0, 0);
}
assert!(!offer(&mut s, 100, 10, 1.0, 0));
assert_eq!(s.encoding(), WigEncoding::VarStep);
assert_eq!(s.starts, [0, 10, 20, 30, 100]);
}
#[test]
fn a_different_span_widens_all_the_way_to_bedgraph() {
let mut s = cost_section();
for i in 0..3 {
offer(&mut s, i * 10, 10, 1.0, 0);
}
assert!(!offer(&mut s, 30, 25, 1.0, 0));
assert_eq!(s.encoding(), WigEncoding::BedGraph);
assert_eq!(s.starts, [0, 10, 20, 30]);
assert_eq!(s.ends, [10, 20, 30, 55]);
}
#[test]
fn a_long_tail_of_the_new_shape_splits_the_section_instead() {
let mut s = cost_section();
for i in 0..4 {
offer(&mut s, i * 10, 10, 1.0, 0);
}
assert!(offer(&mut s, 30, 25, 1.0, 500));
assert_eq!(s.encoding(), WigEncoding::FixedStep);
assert_eq!(s.len(), 1);
assert_eq!(s.uniform_span, 25);
}
#[test]
fn the_tail_is_costed_against_fixed_step_not_against_the_current_encoding() {
let mut s = cost_section();
offer(&mut s, 0, 10, 1.0, 0);
offer(&mut s, 10, 10, 1.0, 0);
offer(&mut s, 40, 10, 1.0, 0); assert_eq!(s.encoding(), WigEncoding::VarStep);
let mut narrow = cost_section();
for (start, span) in [(0, 10), (10, 10), (40, 10)] {
offer(&mut narrow, start, span, 1.0, 0);
}
assert!(!offer(&mut narrow, 50, 25, 1.0, 30));
assert_eq!(narrow.encoding(), WigEncoding::BedGraph);
let mut wide = cost_section();
for (start, span) in [(0, 10), (10, 10), (40, 10)] {
offer(&mut wide, start, span, 1.0, 0);
}
assert!(offer(&mut wide, 50, 25, 1.0, 31));
}
#[test]
fn a_cascade_of_short_sections_stops_itself() {
let mut s = section();
let mut flushes = 0;
let mut at = 0i64;
for i in 0..200 {
let span = if i % 2 == 0 { 10 } else { 11 };
if offer(&mut s, at, span, 1.0, 0) {
flushes += 1;
}
at += span;
}
assert!(
flushes <= RUNT_PATIENCE as usize + 1,
"{flushes} sections for 200 values"
);
assert_eq!(s.encoding(), WigEncoding::BedGraph);
assert!(s.len() > 190, "{} values buffered", s.len());
}
#[test]
fn a_long_section_splits_however_bad_the_recent_history() {
let mut s = section();
let mut at = 0i64;
for i in 0..12 {
let span = 7 + (i % 5);
offer(&mut s, at, span, 1.0, 0);
at += span + 1;
}
assert!(s.encoding() != WigEncoding::FixedStep);
s.clear();
let mut at = 10_000i64;
for _ in 0..500 {
offer(&mut s, at, 10, 1.0, 0);
at += 10;
}
assert_eq!(s.encoding(), WigEncoding::FixedStep);
assert_eq!(s.len(), 500);
assert!(offer(&mut s, at, 25, 1.0, 0), "a 500-value section widened");
assert_eq!(s.len(), 1);
}
#[test]
fn a_long_run_flushes_a_section_it_cannot_extend() {
let mut s = section();
let mut at = 0i64;
for i in 0..6 {
let span = 7 + (i % 3);
offer(&mut s, at, span, 1.0, 0);
at += span + 2;
}
assert_eq!(s.encoding(), WigEncoding::BedGraph);
assert!(s.should_flush_for_run(0, at, 10, 1000));
assert!(!s.should_flush_for_run(0, at, 10, 4));
assert!(!s.should_flush_for_run(1, at, 10, 1000));
let mut fixed = section();
for i in 0..50 {
offer(&mut fixed, i * 10, 10, 1.0, 0);
}
assert!(!fixed.should_flush_for_run(0, 500, 10, 1000));
}
#[test]
fn the_runt_floor_is_read_from_the_cost_model() {
let cost = CostModel {
min_split_items: 4,
..Default::default()
};
let mut s = WigSection::with_policy(1024, SectionPolicy::Adaptive, cost);
for i in 0..3 {
offer(&mut s, i * 10, 10, 1.0, 0);
}
s.clear();
for i in 0..3 {
offer(&mut s, i * 10, 10, 1.0, 0);
}
s.clear();
let mut at = 0i64;
offer(&mut s, at, 10, 1.0, 0);
at += 10;
assert!(
!offer(&mut s, at, 25, 1.0, 0),
"still splitting after 2 runts"
);
assert_eq!(s.encoding(), WigEncoding::BedGraph);
}
#[test]
fn the_tail_is_capped_by_the_room_left_in_the_section() {
let mut s = WigSection::with_policy(8, SectionPolicy::Cost, CostModel::default());
for i in 0..4 {
offer(&mut s, i * 10, 10, 1.0, 0);
}
assert!(!offer(&mut s, 30, 25, 1.0, 1_000_000));
assert_eq!(s.encoding(), WigEncoding::BedGraph);
}
#[test]
fn a_change_of_chromosome_always_flushes() {
let mut s = section();
s.offer(0, 0, 10, 1.0, 0);
assert_eq!(s.offer(1, 0, 10, 1.0, 0), Accept::Flush);
}
#[test]
fn a_second_item_of_the_same_span_stays_fixed_step_whatever_its_start() {
let mut s = section();
offer(&mut s, 0, 10, 1.0, 0);
assert!(!offer(&mut s, 900, 10, 1.0, 1_000_000));
assert_eq!(s.encoding(), WigEncoding::FixedStep);
assert_eq!(s.uniform_step, Some(900));
}
#[test]
fn widening_a_single_item_section_uses_its_span_as_the_step() {
let mut s = widening_section();
offer(&mut s, 40, 10, 1.0, 0);
offer(&mut s, 60, 25, 1.0, 0); assert_eq!(s.encoding(), WigEncoding::BedGraph);
assert_eq!(s.starts, [40, 60]);
assert_eq!(s.ends, [50, 85]);
}
#[test]
fn a_run_extends_the_open_section_only_when_it_lines_up() {
let mut s = section();
s.extend_run(0, 0, 10, &[1.0, 2.0, 3.0]);
assert_eq!(s.len(), 3);
assert_eq!(s.last_end, 30);
assert!(s.extends_run(0, 30, 10));
assert!(!s.extends_run(0, 40, 10), "a gap does not extend");
assert!(!s.extends_run(0, 30, 20), "a different span does not");
assert!(!s.extends_run(1, 30, 10), "another chromosome does not");
s.extend_run(0, 30, 10, &[4.0]);
assert_eq!(s.len(), 4);
assert_eq!(s.encoding(), WigEncoding::FixedStep);
assert!(s.starts.is_empty(), "the fast path materialises no starts");
}
#[test]
fn a_single_item_section_is_extended_by_a_run_that_meets_it() {
let mut s = section();
offer(&mut s, 0, 10, 1.0, 0);
assert!(s.extends_run(0, 10, 10));
s.extend_run(0, 10, 10, &[2.0, 3.0]);
assert_eq!(s.uniform_step, Some(10));
assert_eq!(s.len(), 3);
}
fn header_of(bytes: &[u8]) -> (u32, u32, u32, u32, u32, u8, u16) {
let u32_at = |o: usize| u32::from_le_bytes(bytes[o..o + 4].try_into().unwrap());
(
u32_at(0),
u32_at(4),
u32_at(8),
u32_at(12),
u32_at(16),
bytes[20],
u16::from_le_bytes(bytes[22..24].try_into().unwrap()),
)
}
#[test]
fn a_fixed_step_section_encodes_to_a_header_and_a_column_of_values() {
let mut s = section();
for i in 0..3 {
offer(&mut s, 100 + i * 10, 10, i as f32, 0);
}
let bytes = s.encode().unwrap();
assert_eq!(bytes.len(), WIG_HEADER_SIZE + 3 * 4);
let (chr, start, end, step, span, kind, count) = header_of(&bytes);
assert_eq!((chr, start, end), (0, 100, 130));
assert_eq!((step, span), (10, 10));
assert_eq!(kind, WigEncoding::FixedStep as u8);
assert_eq!(count, 3);
}
#[test]
fn a_sparse_fixed_step_section_ends_at_its_last_item_not_its_last_step() {
let mut s = section();
for i in 0..3 {
offer(&mut s, i * 100, 10, 1.0, 0);
}
let (_, start, end, step, span, _, _) = header_of(&s.encode().unwrap());
assert_eq!((start, end), (0, 210));
assert_eq!((step, span), (100, 10));
}
#[test]
fn a_var_step_section_writes_starts_and_a_span_but_no_step() {
let mut s = widening_section();
offer(&mut s, 0, 10, 1.0, 0);
offer(&mut s, 10, 10, 2.0, 0);
offer(&mut s, 40, 10, 3.0, 0);
let bytes = s.encode().unwrap();
assert_eq!(bytes.len(), WIG_HEADER_SIZE + 3 * 8);
let (_, _, end, step, span, kind, count) = header_of(&bytes);
assert_eq!(kind, WigEncoding::VarStep as u8);
assert_eq!((step, span, end, count), (0, 10, 50, 3));
assert_eq!(u32::from_le_bytes(bytes[24..28].try_into().unwrap()), 0);
assert_eq!(u32::from_le_bytes(bytes[32..36].try_into().unwrap()), 10);
assert_eq!(u32::from_le_bytes(bytes[40..44].try_into().unwrap()), 40);
}
#[test]
fn a_bedgraph_section_writes_both_coordinates_and_neither_step_nor_span() {
let mut s = widening_section();
offer(&mut s, 0, 10, 1.0, 0);
offer(&mut s, 10, 25, 2.0, 0);
let bytes = s.encode().unwrap();
assert_eq!(bytes.len(), WIG_HEADER_SIZE + 2 * 12);
let (_, _, end, step, span, kind, count) = header_of(&bytes);
assert_eq!(kind, WigEncoding::BedGraph as u8);
assert_eq!((step, span, end, count), (0, 0, 35, 2));
}
#[test]
fn a_single_item_section_writes_its_span_as_its_step() {
let mut s = section();
offer(&mut s, 60, 15, 1.0, 0);
let (_, start, end, step, span, _, count) = header_of(&s.encode().unwrap());
assert_eq!((start, end, step, span, count), (60, 75, 15, 15, 1));
}
#[test]
fn an_encoded_section_reads_back_through_the_reader() {
for shape in 0..3 {
let mut s = section();
match shape {
0 => {
for i in 0..5 {
offer(&mut s, i * 10, 10, i as f32 * 1.5, 0);
}
}
1 => {
offer(&mut s, 0, 10, 1.0, 0);
offer(&mut s, 10, 10, 2.0, 0);
offer(&mut s, 45, 10, 3.0, 0);
}
_ => {
offer(&mut s, 0, 10, 1.0, 0);
offer(&mut s, 10, 25, 2.0, 0);
offer(&mut s, 60, 5, 3.0, 0);
}
}
let want: Vec<(i64, i64, f32)> = (0..s.len())
.map(|i| {
let start = if s.encoding() == WigEncoding::FixedStep {
s.first_start + i as i64 * s.uniform_step.unwrap_or(s.uniform_span)
} else {
s.starts[i]
};
let end = if s.encoding() == WigEncoding::BedGraph {
s.ends[i]
} else {
start + s.uniform_span
};
(start, end, s.values[i])
})
.collect();
let bytes = bytes::Bytes::from(s.encode().unwrap());
let header = crate::bbi::block::read_wig_header(&bytes, "test.bigwig").unwrap();
assert_eq!(header.item_count as usize, s.len(), "shape {shape}");
let got: Vec<(i64, i64, f32)> = (0..header.item_count as usize)
.map(|i| {
let item = crate::bbi::block::read_wig_item(&bytes, &header, i, "test.bigwig")
.unwrap();
(item.start, item.end, item.value)
})
.collect();
assert_eq!(got, want, "shape {shape}");
}
}
#[test]
fn a_coordinate_past_32_bits_is_refused_rather_than_truncated() {
let mut s = section();
offer(&mut s, 5_000_000_000, 10, 1.0, 0);
let err = s.encode().unwrap_err().to_string();
assert!(err.contains("chromStart 5000000000"), "{err}");
}
}