use crate::container::Container;
use crate::unit::{DeflateError, Encoder};
#[derive(Debug, Clone)]
pub struct Segment {
pub src_start: usize,
pub src_end: usize,
pub bytes: Vec<u8>,
pub sum: u32,
}
pub fn plan_segments(
len: usize,
segment: usize,
) -> Result<Vec<(usize, usize, bool)>, DeflateError> {
let n = crate::limits::check_pieces("split segments", len, segment)?;
if len == 0 {
return Ok(vec![(0, 0, true)]);
}
let segment = segment.max(1);
Ok((0..n)
.map(|i| {
let start = i * segment;
let end = ((i + 1) * segment).min(len);
(start, end, i + 1 == n)
})
.collect())
}
pub fn compress_segment(
enc: &mut Encoder,
src: &[u8],
start: usize,
end: usize,
is_last: bool,
container: Container,
) -> Result<Segment, DeflateError> {
compress_segment_skipping(enc, src, start, end, is_last, container, false)
}
pub fn compress_segment_skipping(
enc: &mut Encoder,
src: &[u8],
start: usize,
end: usize,
is_last: bool,
container: Container,
store: bool,
) -> Result<Segment, DeflateError> {
debug_assert!(!enc.is_zlib_wrapped(), "segments must be bare DEFLATE");
let chunk = &src[start..end];
let mut bytes = Vec::new();
if store {
bytes = crate::precompressed::stored_deflate(chunk, is_last);
} else if is_last {
enc.compress_into(chunk, &mut bytes)?;
} else {
enc.compress_segment_into(chunk, &mut bytes)?;
}
let sum = container.checksum(chunk);
Ok(Segment {
src_start: start,
src_end: end,
bytes,
sum,
})
}
pub fn assemble(level: u32, segments: &[Segment], container: Container) -> Vec<u8> {
assemble_checksummed(level, segments, container).0
}
pub fn assemble_checksummed(
level: u32,
segments: &[Segment],
container: Container,
) -> (Vec<u8>, u32) {
debug_assert!(!segments.is_empty(), "assemble needs at least one segment");
debug_assert!(
segments.windows(2).all(|w| w[0].src_end == w[1].src_start),
"segments must be contiguous and in stream order"
);
let body: usize = segments.iter().map(|s| s.bytes.len()).sum();
let total: u64 = segments
.iter()
.map(|s| (s.src_end - s.src_start) as u64)
.sum();
let mut out = Vec::with_capacity(12 + body + 8);
out.extend_from_slice(&container.header(level));
let mut sum = container.checksum_init();
for s in segments {
out.extend_from_slice(&s.bytes);
let n = (s.src_end - s.src_start) as u64;
sum = container.checksum_combine(sum, s.sum, n);
}
out.extend_from_slice(&container.trailer(sum, total));
(out, sum)
}
pub fn split_compress_serial(
src: &[u8],
level: u32,
segment: usize,
) -> Result<Vec<u8>, DeflateError> {
split_compress_serial_in(src, level, segment, Container::Zlib)
}
pub fn split_compress_serial_in(
src: &[u8],
level: u32,
segment: usize,
container: Container,
) -> Result<Vec<u8>, DeflateError> {
let mut enc = Encoder::new(level, false)?;
let plan = plan_segments(src.len(), segment)?;
let mut segs = Vec::with_capacity(plan.len());
for (start, end, last) in plan {
segs.push(compress_segment(
&mut enc, src, start, end, last, container,
)?);
}
Ok(assemble(level, &segs, container))
}
#[cfg(test)]
mod tests {
use super::*;
use crate::plan::Config;
use std::io::Read;
fn corpus(n: usize) -> Vec<u8> {
let words: &[&[u8]] = &[
b"fn main() { println!(\"hello\"); }\n",
b"use std::collections::HashMap;\n",
b" let mut total = 0usize;\n",
b"// TODO: revisit this once the planner lands\n",
];
let mut out = Vec::with_capacity(n + 64);
let mut s: u64 = 0x9E3779B97F4A7C15;
while out.len() < n {
s = s
.wrapping_mul(6364136223846793005)
.wrapping_add(1442695040888963407);
out.extend_from_slice(words[(s >> 33) as usize % words.len()]);
}
out.truncate(n);
out
}
fn round_trips(len: usize, seg: usize) {
let src = corpus(len);
let z = split_compress_serial(&src, 6, seg).unwrap();
let mut got = Vec::new();
flate2::read::ZlibDecoder::new(&z[..])
.read_to_end(&mut got)
.unwrap_or_else(|e| panic!("len={len} seg={seg}: inflate failed: {e}"));
assert_eq!(got, src, "len={len} seg={seg}: bytes differ");
}
#[test]
fn degenerate_segment_sizes_round_trip_on_small_inputs() {
for &len in &[0usize, 1, 2, 3, 7, 8, 4095, 4096, 4097, 8192] {
for &seg in &[1usize, 2, 3, 7, 4096] {
round_trips(len, seg);
}
}
}
#[test]
fn large_inputs_round_trip_at_realistic_segment_sizes() {
for &len in &[100_000usize, 1_000_000] {
for &seg in &[4096usize, 65_536, 1 << 20] {
round_trips(len, seg);
}
}
}
#[test]
fn a_degenerate_segment_size_is_refused_not_attempted() {
use crate::limits::MAX_PIECES;
let just_over = MAX_PIECES + 1;
let e = plan_segments(just_over, 1).expect_err("one segment per byte must be refused");
assert!(
matches!(e, DeflateError::TooManyPieces { pieces, max, .. }
if pieces == just_over && max == MAX_PIECES),
"wrong error: {e}"
);
assert!(
e.to_string().contains("split segments"),
"message does not name the axis: {e}"
);
assert!(
plan_segments(1_000_000, 1).is_err(),
"a million one-byte segments were planned"
);
assert!(
plan_segments(1_000_000, 0).is_err(),
"segment=0 escaped the ceiling"
);
assert!(
split_compress_serial(&vec![b'x'; 1_000_000], 6, 1).is_err(),
"the serial split path attempted a million one-byte segments"
);
assert!(
crate::compress_split(
&vec![b'x'; 1_000_000],
&Config {
segment: 1,
..Config::default()
}
)
.is_err(),
"the public compress_split entry point bypassed the ceiling"
);
assert!(
crate::compress_objects(
&[&vec![b'x'; 1_000_000][..]],
&Config {
split_above: 1,
segment: 1,
..Config::default()
},
)
.is_err(),
"compress_objects bypassed the ceiling"
);
assert_eq!(plan_segments(MAX_PIECES, 1).unwrap().len(), MAX_PIECES);
assert!(
plan_segments(256 << 30, 1 << 20).is_ok(),
"256 GiB at 1 MiB segments"
);
assert!(split_compress_serial(&vec![b'x'; 1_000_000], 6, 4096).is_ok());
}
#[test]
fn a_wrong_adler_trailer_is_actually_rejected() {
let src = corpus(200_000);
let mut z = split_compress_serial(&src, 6, 16_384).unwrap();
let mut ok = Vec::new();
flate2::read::ZlibDecoder::new(&z[..])
.read_to_end(&mut ok)
.unwrap();
assert_eq!(ok, src);
let n = z.len();
z[n - 1] ^= 0xff;
let mut bad = Vec::new();
let r = flate2::read::ZlibDecoder::new(&z[..]).read_to_end(&mut bad);
assert!(
r.is_err(),
"the decoder did NOT verify the Adler-32 trailer — the round-trip \
tests in this module prove nothing about adler32_combine"
);
}
#[test]
fn a_stream_with_no_final_block_is_rejected() {
let src = corpus(50_000);
let mut enc = Encoder::new(6, false).unwrap();
let plan = plan_segments(src.len(), 8192).unwrap();
let mut segs = Vec::new();
for (start, end, _last) in plan {
segs.push(
compress_segment(&mut enc, &src, start, end, false, Container::Zlib).unwrap(),
);
}
let z = assemble(6, &segs, Container::Zlib);
let mut got = Vec::new();
let r = flate2::read::ZlibDecoder::new(&z[..]).read_to_end(&mut got);
assert!(
r.is_err() || got != src,
"a stream whose last block is not BFINAL was accepted as complete — \
the `is_last` branch of compress_segment is not being tested"
);
}
#[test]
fn segment_plan_is_contiguous_and_covering() {
let check = |len: usize, seg: usize| {
let p = plan_segments(len, seg)
.unwrap_or_else(|e| panic!("len={len} seg={seg}: refused: {e}"));
assert!(!p.is_empty());
assert_eq!(p[0].0, 0, "len={len} seg={seg}: must start at 0");
assert_eq!(
p.last().unwrap().1,
len,
"len={len} seg={seg}: must cover to len"
);
assert_eq!(
p.iter().filter(|(_, _, l)| *l).count(),
1,
"len={len} seg={seg}: exactly one final segment"
);
for w in p.windows(2) {
assert_eq!(w[0].1, w[1].0, "len={len} seg={seg}: gap or overlap");
assert!(
!w[0].2,
"len={len} seg={seg}: is_last set on a non-final segment"
);
}
};
for &len in &[0usize, 1, 999, 1000, 1001, 9999] {
for &seg in &[1usize, 7, 1000] {
check(len, seg);
}
}
for &len in &[1usize << 20, (1 << 20) + 1, 64 << 20] {
for &seg in &[1000usize, 1 << 16, 1 << 20] {
check(len, seg);
}
}
}
#[test]
fn splitting_costs_ratio_and_the_cost_is_bounded() {
let src = corpus(4 << 20);
let whole = crate::unit::zlib_compress(&src, 6).unwrap();
for &seg in &[1 << 20, 256 << 10, 64 << 10] {
let split = split_compress_serial(&src, 6, seg).unwrap();
let overhead = split.len() as f64 / whole.len() as f64;
assert!(
overhead >= 1.0,
"seg={seg}: split output ({}) SMALLER than unsplit ({}) — impossible \
unless the window is not actually being reset",
split.len(),
whole.len()
);
assert!(
overhead < 1.25,
"seg={seg}: split costs {:.1}% more than unsplit — far beyond the \
expected few percent; a boundary is being emitted per block, not per segment",
(overhead - 1.0) * 100.0
);
}
}
}