use alloc::boxed::Box;
use alloc::vec::Vec;
use core::mem::MaybeUninit;
use core::ops::Range;
use super::config::Chunking;
use super::config::{ConfigError, EncoderConfig, SizeOverflow};
use super::core::graphs::SERIAL_CHUNK_BYTE_SIZE;
use super::core::params::LocalParams as CoreParams;
use super::core::registry::{RuntimeParams, gid};
use super::core::sink::{SliceSink, UninitSink, VecSink};
use super::core::stream::Stream;
use super::core::{Engine, GraphPlan, Params};
use super::error::EncodeError;
use super::session::{EncoderSession, EncoderSessionOwned, StreamConfig};
use crate::graph::{Comment, CompressorGraph, GraphError};
use crate::input::Input;
use crate::shared::wire::COMMENT_VERSION_MIN;
use crate::{Backend, RetentionPolicy};
use alloc::sync::Arc;
const FRAME_OVERHEAD_MAX: usize = 64;
const CHUNK_OVERHEAD_MAX: usize = 16;
const BOUND_CHUNK: usize = 32 << 10;
#[derive(Debug)]
pub struct Compressor {
config: EncoderConfig,
params: Params,
graph: Option<CompressorGraph>,
comment: Option<Comment>,
plan_error: Option<EncodeError>,
backend: Backend,
retention: RetentionPolicy,
pub(crate) engine: Engine,
pub(crate) stream: Stream,
}
impl Default for Compressor {
fn default() -> Self {
Self::with_parts(
EncoderConfig::DEFAULT,
RetentionPolicy::default(),
Backend::default(),
None,
)
}
}
impl Compressor {
fn with_parts(
config: EncoderConfig,
retention: RetentionPolicy,
backend: Backend,
graph: Option<CompressorGraph>,
) -> Self {
let params = Params::from(&config);
Self {
config,
params,
graph,
comment: None,
plan_error: None,
backend,
retention,
engine: Engine::new(¶ms, backend),
stream: Stream::default(),
}
}
pub fn new(config: EncoderConfig) -> Result<Self, ConfigError> {
Self::builder(config).build()
}
#[must_use]
pub const fn builder(config: EncoderConfig) -> CompressorBuilder {
CompressorBuilder {
config,
retention: RetentionPolicy::Aggressive,
backend: None,
graph: None,
}
}
#[must_use]
pub const fn config(&self) -> &EncoderConfig {
&self.config
}
#[must_use]
pub const fn backend(&self) -> Backend {
self.backend
}
#[must_use]
pub const fn retention(&self) -> RetentionPolicy {
self.retention
}
pub fn reconfigure(&mut self, config: EncoderConfig) -> Result<(), ConfigError> {
if self.comment.is_some() && u32::from(config.format_version()) < COMMENT_VERSION_MIN {
return Err(ConfigError::CommentUnsupported {
version: u32::from(config.format_version()),
});
}
self.config = config;
self.params = Params::from(&config);
self.engine.reconfigure(&self.params);
self.stream.reset();
Ok(())
}
pub const fn max_compressed_size(input_size: usize) -> Result<usize, SizeOverflow> {
let chunks = (input_size / BOUND_CHUNK) + 1;
let Some(overhead) = chunks.checked_mul(CHUNK_OVERHEAD_MAX) else {
return Err(SizeOverflow);
};
let Some(with_frame) = input_size.checked_add(FRAME_OVERHEAD_MAX) else {
return Err(SizeOverflow);
};
match with_frame.checked_add(overhead) {
Some(bound) => Ok(bound),
None => Err(SizeOverflow),
}
}
pub fn compress(&mut self, src: &[u8]) -> Result<Vec<u8>, EncodeError> {
let mut dst = Vec::new();
self.compress_into(src, &mut dst)?;
Ok(dst)
}
pub fn compress_into(
&mut self,
src: &[u8],
dst: &mut Vec<u8>,
) -> Result<Range<usize>, EncodeError> {
self.sync_plan()?;
let start = dst.len();
let result = self.engine.frame(&self.params, src, &mut VecSink::new(dst));
self.apply_retention();
match result {
Ok(()) => Ok(start..dst.len()),
Err(error) => {
dst.truncate(start);
Err(EncodeError::from_core(error, usize::MAX))
}
}
}
pub fn compress_to_slice(&mut self, src: &[u8], dst: &mut [u8]) -> Result<usize, EncodeError> {
self.sync_plan()?;
let provided = dst.len();
let mut sink = SliceSink::new(dst);
let result = self.engine.frame(&self.params, src, &mut sink);
let written = super::core::sink::Sink::len(&sink);
self.apply_retention();
result
.map(|()| written)
.map_err(|error| EncodeError::from_core(error, provided))
}
pub fn compress_to_uninit(
&mut self,
src: &[u8],
dst: &mut [MaybeUninit<u8>],
) -> Result<usize, EncodeError> {
self.sync_plan()?;
let provided = dst.len();
let mut sink = UninitSink::from(dst);
let result = self.engine.frame(&self.params, src, &mut sink);
let written = super::core::sink::Sink::len(&sink);
self.apply_retention();
result
.map(|()| written)
.map_err(|error| EncodeError::from_core(error, provided))
}
pub fn start(&mut self, stream: StreamConfig) -> EncoderSession<'_> {
self.begin_stream();
EncoderSession::new(self, stream)
}
#[must_use]
pub fn into_session(mut self, stream: StreamConfig) -> EncoderSessionOwned {
self.begin_stream();
EncoderSessionOwned::new(Box::new(self), stream)
}
#[must_use]
pub fn retained_bytes(&self) -> usize {
self.engine.retained_bytes() + self.stream.retained_bytes()
}
pub fn trim(&mut self, policy: RetentionPolicy) {
if !policy.allows(self.stream.retained_bytes()) {
self.stream.release();
}
self.engine.trim(policy);
}
pub(crate) fn begin_stream(&mut self) {
self.stream.reset();
self.plan_error = self.sync_plan().err();
}
pub(crate) fn take_plan_error(&mut self) -> Option<EncodeError> {
self.plan_error.take()
}
fn sync_plan(&mut self) -> Result<(), EncodeError> {
self.params = Params::from(&self.config);
let custom = self.graph.is_some() || self.comment.is_some() || !self.params.fast_path();
if !custom {
self.engine.set_plan(None);
return Ok(());
}
let graph = self.graph.clone().unwrap_or_default();
let registry = Arc::clone(&graph.0);
let mut start = registry.start;
let mut start_params = registry.start_params.clone();
if let Chunking::Fixed(size) = self.config.chunking() {
if matches!(
registry.graph(start).map(|def| &def.kind),
Ok(super::core::registry::GraphKind::Segmenter(_))
) {
return Err(GraphError::Invalid(
"fixed chunking cannot wrap a graph that is already a segmenter",
)
.into());
}
let chunk = i32::try_from(usize::from(size)).unwrap_or(i32::MAX);
let mut params = CoreParams::default();
params.ints.to_mut().push((SERIAL_CHUNK_BYTE_SIZE, chunk));
start_params = Some(RuntimeParams {
params: Some(params),
custom_graphs: Some(alloc::vec![start]),
custom_nodes: None,
});
start = gid::SEGMENT_SERIAL;
}
let segmented = registry.serial_plan(start, start_params.as_ref());
if let Some((chunk, _)) = segmented {
self.params.chunk_limit = chunk;
}
self.engine.set_plan(Some(GraphPlan {
registry,
start,
start_params,
comment: self.comment.as_ref().map(Comment::shared),
segmented,
}));
Ok(())
}
pub fn set_comment(&mut self, comment: Option<Comment>) -> Result<(), ConfigError> {
let version = u32::from(self.config.format_version());
if comment.is_some() && version < COMMENT_VERSION_MIN {
return Err(ConfigError::CommentUnsupported { version });
}
self.comment = comment;
Ok(())
}
pub fn compress_inputs(&mut self, inputs: &[Input<'_>]) -> Result<Vec<u8>, EncodeError> {
let mut dst = Vec::new();
self.compress_inputs_into(inputs, &mut dst)?;
Ok(dst)
}
pub fn compress_inputs_into(
&mut self,
inputs: &[Input<'_>],
dst: &mut Vec<u8>,
) -> Result<Range<usize>, EncodeError> {
self.sync_plan()?;
let streams = inputs.iter().map(|input| input.to_stream()).collect();
let start = dst.len();
let result = self
.engine
.frame_inputs(&self.params, streams, &mut VecSink::new(dst));
self.apply_retention();
match result {
Ok(()) => Ok(start..dst.len()),
Err(error) => {
dst.truncate(start);
Err(EncodeError::from_core(error, usize::MAX))
}
}
}
pub(crate) fn apply_retention(&mut self) {
if !self.retention.allows(self.retained_bytes()) {
self.trim(self.retention);
}
}
pub(crate) const fn params(&self) -> &Params {
&self.params
}
}
#[derive(Clone, Debug)]
pub struct CompressorBuilder {
config: EncoderConfig,
retention: RetentionPolicy,
backend: Option<Backend>,
graph: Option<CompressorGraph>,
}
impl CompressorBuilder {
#[must_use]
pub const fn with_retention(mut self, retention: RetentionPolicy) -> Self {
self.retention = retention;
self
}
#[must_use]
pub fn with_graph(mut self, graph: CompressorGraph) -> Self {
self.graph = Some(graph);
self
}
#[must_use]
pub const fn with_backend(mut self, backend: Backend) -> Self {
self.backend = Some(backend);
self
}
pub fn build(self) -> Result<Compressor, ConfigError> {
Ok(Compressor::with_parts(
self.config,
self.retention,
self.backend.unwrap_or_default(),
self.graph,
))
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn bound_matches_the_reference_formula() {
for size in [0usize, 1, 32_767, 32_768, 1 << 20] {
assert_eq!(
Compressor::max_compressed_size(size),
Ok(openzl_ffi::compress_bound(size))
);
}
assert_eq!(
Compressor::max_compressed_size(usize::MAX),
Err(SizeOverflow)
);
assert_eq!(
Compressor::max_compressed_size(usize::MAX - 64),
Err(SizeOverflow)
);
}
#[test]
fn slices_too_small_fail_and_vec_appends_roll_back() {
let mut encoder = Compressor::default();
let mut tiny = [0u8; 10];
assert_eq!(
encoder.compress_to_slice(b"payload", &mut tiny),
Err(EncodeError::OutputTooSmall { provided: 10 })
);
let mut out = b"keep".to_vec();
let range = encoder.compress_into(b"payload", &mut out).unwrap();
assert_eq!(&out[..4], b"keep");
assert_eq!(range.end, out.len());
}
#[cfg(feature = "zstd")]
#[test]
fn reconfiguring_switches_the_next_frame() {
let mut encoder = Compressor::default();
assert_eq!(encoder.retention(), RetentionPolicy::Aggressive);
let input: Vec<u8> = (0..30_000u32)
.map(|i| b"abcdefgh"[(i.wrapping_mul(2_654_435_761) >> 29) as usize])
.collect();
let default = encoder.compress(&input).unwrap();
let config = EncoderConfig::from(crate::Level::try_from(1).unwrap());
encoder.reconfigure(config).unwrap();
assert_eq!(encoder.config(), &config);
let fast = encoder.compress(&input).unwrap();
assert_ne!(fast, default);
assert_eq!(
fast,
Compressor::new(config).unwrap().compress(&input).unwrap()
);
}
#[test]
fn release_all_keeps_nothing_after_an_operation() {
let mut encoder = Compressor::builder(EncoderConfig::default())
.with_retention(RetentionPolicy::ReleaseAll)
.build()
.unwrap();
encoder.compress(&[7u8; 10_000]).unwrap();
assert_eq!(encoder.retained_bytes(), 0);
assert_eq!(encoder.params().level, 6);
}
}