use crate::container::Container;
use crate::precompressed::SkipPolicy;
use crate::split::{self, Segment};
use crate::unit::{DeflateError, Encoder};
#[derive(Debug, Clone, Copy)]
pub struct Config {
pub level: u32,
pub coalesce_below: usize,
pub batch_bytes: usize,
pub batch_count: usize,
pub split_above: usize,
pub segment: usize,
pub workers: usize,
pub claim: usize,
pub window_bits: u8,
pub container: Container,
pub skip: SkipPolicy,
}
impl Default for Config {
fn default() -> Self {
Config {
level: 6,
coalesce_below: 64 * 1024,
batch_bytes: 512 * 1024,
batch_count: 256,
split_above: 4 << 20,
segment: 1 << 20,
workers: 0,
claim: 4,
window_bits: crate::unit::WINDOW_BITS_DEFAULT,
container: Container::Zlib,
skip: SkipPolicy::resolve(),
}
}
}
impl Config {
pub fn level(mut self, level: u32) -> Self {
self.level = level;
self
}
pub fn workers(mut self, w: usize) -> Self {
self.workers = w;
self
}
pub fn without_coalescing(mut self) -> Self {
self.coalesce_below = 0;
self
}
pub fn without_splitting(mut self) -> Self {
self.split_above = usize::MAX;
self
}
pub fn container(mut self, container: Container) -> Self {
self.container = container;
self
}
pub fn skip(mut self, skip: SkipPolicy) -> Self {
self.skip = skip;
self
}
pub fn compress_everything(mut self) -> Self {
self.skip = SkipPolicy::compress_everything();
self
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Unit {
Batch { lo: usize, hi: usize },
Whole { idx: usize },
Segment { idx: usize, seg: usize, start: usize, end: usize, is_last: bool },
}
pub fn plan(sizes: &[usize], cfg: &Config) -> Result<Vec<Unit>, DeflateError> {
let mut units = Vec::new();
let mut batch_lo: Option<usize> = None;
let mut batch_bytes = 0usize;
let mut batch_count = 0usize;
let flush = |units: &mut Vec<Unit>,
batch_lo: &mut Option<usize>,
batch_bytes: &mut usize,
batch_count: &mut usize,
hi: usize| {
if let Some(lo) = batch_lo.take() {
if hi > lo {
units.push(Unit::Batch { lo, hi });
}
}
*batch_bytes = 0;
*batch_count = 0;
};
for (idx, &size) in sizes.iter().enumerate() {
if size < cfg.coalesce_below {
if batch_lo.is_none() {
batch_lo = Some(idx);
}
batch_bytes += size;
batch_count += 1;
if batch_bytes >= cfg.batch_bytes || batch_count >= cfg.batch_count {
flush(&mut units, &mut batch_lo, &mut batch_bytes, &mut batch_count, idx + 1);
}
continue;
}
flush(&mut units, &mut batch_lo, &mut batch_bytes, &mut batch_count, idx);
if size >= cfg.split_above {
for (seg, (start, end, is_last)) in
split::plan_segments(size, cfg.segment)?.into_iter().enumerate()
{
units.push(Unit::Segment { idx, seg, start, end, is_last });
}
} else {
units.push(Unit::Whole { idx });
}
}
flush(&mut units, &mut batch_lo, &mut batch_bytes, &mut batch_count, sizes.len());
Ok(units)
}
#[inline]
pub fn weight(u: &Unit, sizes: &[usize]) -> u64 {
match *u {
Unit::Batch { lo, hi } => sizes[lo..hi].iter().map(|&s| s as u64).sum(),
Unit::Whole { idx } => sizes[idx] as u64,
Unit::Segment { start, end, .. } => (end - start) as u64,
}
}
enum UnitOut {
Batch(Vec<(Vec<u8>, u32)>),
Whole((Vec<u8>, u32)),
Segment(Segment),
}
fn run_unit(
u: &Unit,
objects: &[&[u8]],
enc_wrapped: &mut Encoder,
enc_raw: &mut Encoder,
level: u32,
container: Container,
want_sum: bool,
skip: SkipPolicy,
) -> Result<UnitOut, DeflateError> {
let one = |src: &[u8],
ew: &mut Encoder,
er: &mut Encoder|
-> Result<(Vec<u8>, u32), DeflateError> {
let sum = if want_sum { container.checksum(src) } else { 0 };
if skip.skip(src) {
let body = crate::precompressed::stored_deflate(src, true);
return Ok((container.wrap(level, &body, src), sum));
}
if container.encoder_is_zlib_wrapped() {
let mut o = Vec::new();
ew.compress_into(src, &mut o)?;
Ok((o, sum))
} else {
let mut body = Vec::new();
er.compress_into(src, &mut body)?;
Ok((container.wrap(level, &body, src), sum))
}
};
match *u {
Unit::Batch { lo, hi } => {
let mut outs = Vec::with_capacity(hi - lo);
for obj in &objects[lo..hi] {
outs.push(one(obj, enc_wrapped, enc_raw)?);
}
Ok(UnitOut::Batch(outs))
}
Unit::Whole { idx } => Ok(UnitOut::Whole(one(objects[idx], enc_wrapped, enc_raw)?)),
Unit::Segment { idx, start, end, is_last, .. } => {
Ok(UnitOut::Segment(split::compress_segment_skipping(
enc_raw,
objects[idx],
start,
end,
is_last,
container,
skip.skip(objects[idx]),
)?))
}
}
}
pub fn compress_objects(
objects: &[&[u8]],
cfg: &Config,
) -> Result<Vec<Vec<u8>>, DeflateError> {
Ok(compress_objects_inner(objects, cfg, false)?.into_iter().map(|(b, _)| b).collect())
}
pub fn compress_objects_checksummed(
objects: &[&[u8]],
cfg: &Config,
) -> Result<Vec<(Vec<u8>, u32)>, DeflateError> {
compress_objects_inner(objects, cfg, true)
}
fn compress_objects_inner(
objects: &[&[u8]],
cfg: &Config,
want_sum: bool,
) -> Result<Vec<(Vec<u8>, u32)>, DeflateError> {
let sizes: Vec<usize> = objects.iter().map(|o| o.len()).collect();
let units = plan(&sizes, cfg)?;
if units.is_empty() {
return Ok(Vec::new());
}
let level = cfg.level;
let wb = cfg.window_bits;
let ct = cfg.container;
let results: Vec<Result<UnitOut, DeflateError>> =
gatling::gatling_forkjoin::gatling_map_balanced(
&units,
cfg.workers,
cfg.claim,
|u| weight(u, &sizes),
|_i, u| {
let mut ez = Encoder::new_with_window(level, ct.encoder_is_zlib_wrapped(), wb)?;
let mut er = Encoder::new_with_window(level, false, wb)?;
run_unit(u, objects, &mut ez, &mut er, level, ct, want_sum, cfg.skip)
},
);
let mut out: Vec<(Vec<u8>, u32)> = Vec::with_capacity(objects.len());
let mut pending: Vec<Segment> = Vec::new();
let mut pending_idx: Option<usize> = None;
for (u, r) in units.iter().zip(results) {
match (u, r?) {
(Unit::Batch { .. }, UnitOut::Batch(v)) => out.extend(v),
(Unit::Whole { .. }, UnitOut::Whole(v)) => out.push(v),
(Unit::Segment { idx, is_last, .. }, UnitOut::Segment(s)) => {
debug_assert!(pending_idx.is_none_or(|p| p == *idx));
pending_idx = Some(*idx);
pending.push(s);
if *is_last {
out.push(split::assemble_checksummed(level, &pending, ct));
pending.clear();
pending_idx = None;
}
}
_ => unreachable!("unit kind and output kind cannot disagree"),
}
}
debug_assert!(pending.is_empty(), "a split object was left unassembled");
debug_assert_eq!(out.len(), objects.len(), "one output stream per input object");
gatling::functional_status(
"ldeflate",
"compress_objects",
true,
&"objects fanned out over one no-barrier gatling pass",
);
Ok(out)
}
pub fn compress_objects_serial(
objects: &[&[u8]],
cfg: &Config,
) -> Result<Vec<Vec<u8>>, DeflateError> {
let sizes: Vec<usize> = objects.iter().map(|o| o.len()).collect();
let units = plan(&sizes, cfg)?;
let mut ez =
Encoder::new_with_window(cfg.level, cfg.container.encoder_is_zlib_wrapped(), cfg.window_bits)?;
let mut er = Encoder::new_with_window(cfg.level, false, cfg.window_bits)?;
let mut out: Vec<Vec<u8>> = Vec::with_capacity(objects.len());
let mut pending: Vec<Segment> = Vec::new();
for u in &units {
match run_unit(u, objects, &mut ez, &mut er, cfg.level, cfg.container, false, cfg.skip)? {
UnitOut::Batch(v) => out.extend(v.into_iter().map(|(b, _)| b)),
UnitOut::Whole((b, _)) => out.push(b),
UnitOut::Segment(s) => {
let last = matches!(u, Unit::Segment { is_last: true, .. });
pending.push(s);
if last {
out.push(split::assemble(cfg.level, &pending, cfg.container));
pending.clear();
}
}
}
}
Ok(out)
}
#[cfg(test)]
mod tests {
use super::*;
fn objs(sizes: &[usize]) -> Vec<Vec<u8>> {
sizes
.iter()
.enumerate()
.map(|(i, &n)| {
let tag = format!("object-{i}-");
let mut v = Vec::with_capacity(n);
while v.len() < n {
v.extend_from_slice(tag.as_bytes());
v.extend_from_slice(b"payload payload payload ");
}
v.truncate(n);
v
})
.collect()
}
fn refs(v: &[Vec<u8>]) -> Vec<&[u8]> {
v.iter().map(|x| x.as_slice()).collect()
}
#[test]
fn plan_covers_every_object_exactly_once() {
let cfg = Config::default();
let sizes: Vec<usize> = (0..500)
.map(|i| match i % 7 {
0 => 120,
1 => 900,
2 => 4_000,
3 => 70_000,
4 => 300_000,
5 => 3_000_000,
_ => 40,
})
.collect();
let units = plan(&sizes, &cfg).unwrap();
let mut seen = vec![0usize; sizes.len()];
let mut order = Vec::new();
for u in &units {
match *u {
Unit::Batch { lo, hi } => {
for i in lo..hi {
seen[i] += 1;
order.push(i);
}
}
Unit::Whole { idx } => {
seen[idx] += 1;
order.push(idx);
}
Unit::Segment { idx, is_last, .. } => {
if is_last {
seen[idx] += 1;
order.push(idx);
}
}
}
}
assert!(seen.iter().all(|&c| c == 1), "every object exactly once: {seen:?}");
assert!(order.windows(2).all(|w| w[0] < w[1]), "objects must stay in order");
}
#[test]
fn parallel_equals_serial_at_every_worker_count() {
let data = objs(&[80, 300, 900, 64_000, 200_000, 12 << 20, 40, 5, 1_500_000]);
let r = refs(&data);
let base = compress_objects_serial(&r, &Config::default()).unwrap();
for w in [1usize, 2, 3, 8, 32] {
let got = compress_objects(&r, &Config::default().workers(w)).unwrap();
assert_eq!(got, base, "workers={w}: output differs from the serial oracle");
}
}
#[test]
fn segment_size_actually_changes_the_bytes() {
let data = objs(&[24 << 20]);
let r = refs(&data);
assert!(data[0].len() >= Config::default().split_above, "fixture must reach the split queue");
let a = compress_objects(&r, &Config { segment: 256 << 10, ..Config::default() }).unwrap();
let b = compress_objects(&r, &Config { segment: 1 << 20, ..Config::default() }).unwrap();
assert_ne!(
a, b,
"two different segment sizes produced identical bytes — the split path \
is not actually segmenting, so the determinism test above proves nothing"
);
}
#[test]
fn every_queue_round_trips() {
use std::io::Read;
let data = objs(&[10, 200, 65_000, 300_000, 9 << 20, 7, 1_100_000]);
let r = refs(&data);
let out = compress_objects(&r, &Config::default()).unwrap();
assert_eq!(out.len(), data.len());
for (i, (z, src)) in out.iter().zip(&data).enumerate() {
let mut got = Vec::new();
flate2::read::ZlibDecoder::new(&z[..])
.read_to_end(&mut got)
.unwrap_or_else(|e| panic!("object {i} ({} B) failed to inflate: {e}", src.len()));
assert_eq!(&got, src, "object {i} ({} B) round-trip differs", src.len());
}
}
#[test]
fn coalescing_actually_changes_the_unit_count() {
let sizes = vec![200usize; 10_000];
let batched = plan(&sizes, &Config::default()).unwrap();
let unbatched = plan(&sizes, &Config::default().without_coalescing()).unwrap();
assert_eq!(unbatched.len(), 10_000, "no-coalesce must be one unit per object");
assert!(
batched.len() * 8 < unbatched.len(),
"coalescing produced {} units vs {} — barely fewer, so the small queue \
is not doing anything and the A/B arm is measuring noise",
batched.len(),
unbatched.len()
);
let data = objs(&vec![200usize; 400]);
let r = refs(&data);
assert_eq!(
compress_objects(&r, &Config::default()).unwrap(),
compress_objects(&r, &Config::default().without_coalescing()).unwrap(),
"coalescing must not change a single output byte"
);
}
#[test]
fn degenerate_inputs() {
let cfg = Config::default();
assert!(compress_objects(&[], &cfg).unwrap().is_empty());
let empties: Vec<Vec<u8>> = vec![Vec::new(); 5];
let r = refs(&empties);
let out = compress_objects(&r, &cfg).unwrap();
assert_eq!(out.len(), 5);
use std::io::Read;
for z in &out {
let mut got = Vec::new();
flate2::read::ZlibDecoder::new(&z[..]).read_to_end(&mut got).unwrap();
assert!(got.is_empty());
}
}
}