use alloc::boxed::Box;
use core::mem::MaybeUninit;
use super::core::stream::{Op, StepStatus};
use super::{DecodeError, Decompressor};
use crate::shared::destination::Destination;
#[derive(Clone, Copy, Debug, Default, Eq, PartialEq, Hash)]
pub enum DecodeOperation {
#[default]
Process,
Finish,
}
#[derive(Clone, Copy, Debug, Eq, PartialEq, Hash)]
pub enum DecoderStatus {
NeedsInput,
NeedsOutput,
Finished,
}
#[derive(Clone, Copy, Debug, Eq, PartialEq, Hash)]
pub struct DecodeProgress {
pub consumed: usize,
pub produced: usize,
pub status: DecoderStatus,
}
#[derive(Debug, Clone, Eq, PartialEq, thiserror::Error)]
#[error("{error}")]
pub struct DecodeFailure {
#[source]
pub error: DecodeError,
pub consumed: usize,
pub produced: usize,
}
impl DecodeFailure {
#[must_use]
pub fn into_error(self) -> DecodeError {
self.error
}
}
impl From<DecodeFailure> for DecodeError {
fn from(failure: DecodeFailure) -> Self {
failure.error
}
}
fn process<D: Destination + ?Sized>(
decoder: &mut Decompressor,
failed: &mut bool,
input: &[u8],
output: &mut D,
operation: DecodeOperation,
) -> Result<DecodeProgress, DecodeFailure> {
if *failed {
return Err(DecodeFailure {
error: DecodeError::InvalidState,
consumed: 0,
produced: 0,
});
}
let op = match operation {
DecodeOperation::Process => Op::Process,
DecodeOperation::Finish => Op::Finish,
};
let config = *decoder.config();
match decoder
.stream
.process(&mut decoder.engine, &config, input, output, op)
{
Ok(step) => {
let status = match step.status {
StepStatus::NeedsInput => DecoderStatus::NeedsInput,
StepStatus::NeedsOutput => DecoderStatus::NeedsOutput,
StepStatus::Finished => {
decoder.apply_retention();
DecoderStatus::Finished
}
};
Ok(DecodeProgress {
consumed: step.consumed,
produced: step.produced,
status,
})
}
Err(failure) => {
*failed = true;
decoder.apply_retention();
Err(DecodeFailure {
error: failure.error,
consumed: failure.consumed,
produced: failure.produced,
})
}
}
}
macro_rules! session_methods {
() => {
pub fn process(
&mut self,
input: &[u8],
output: &mut [u8],
operation: DecodeOperation,
) -> Result<DecodeProgress, DecodeFailure> {
process(
&mut self.decoder,
&mut self.failed,
input,
output,
operation,
)
}
pub fn process_uninit(
&mut self,
input: &[u8],
output: &mut [MaybeUninit<u8>],
operation: DecodeOperation,
) -> Result<DecodeProgress, DecodeFailure> {
process(
&mut self.decoder,
&mut self.failed,
input,
output,
operation,
)
}
pub fn flush(&mut self, output: &mut [u8]) -> Result<DecodeProgress, DecodeFailure> {
self.process(&[], output, DecodeOperation::Process)
}
pub fn finish(&mut self, output: &mut [u8]) -> Result<DecodeProgress, DecodeFailure> {
self.process(&[], output, DecodeOperation::Finish)
}
#[must_use]
pub fn is_finished(&self) -> bool {
!self.failed && self.decoder.stream.is_finished()
}
#[must_use]
pub fn total_in(&self) -> u64 {
self.decoder.stream.total_in()
}
#[must_use]
pub fn total_out(&self) -> u64 {
self.decoder.stream.total_out()
}
#[must_use]
pub fn frames_decoded(&self) -> u64 {
self.decoder.stream.frames()
}
};
}
#[derive(Debug)]
pub struct DecoderSession<'d> {
decoder: &'d mut Decompressor,
failed: bool,
}
impl<'d> DecoderSession<'d> {
pub(crate) const fn new(decoder: &'d mut Decompressor) -> Self {
Self {
decoder,
failed: false,
}
}
session_methods!();
}
#[derive(Debug)]
pub struct DecoderSessionOwned {
decoder: Box<Decompressor>,
failed: bool,
}
impl DecoderSessionOwned {
pub(crate) const fn new(decoder: Box<Decompressor>) -> Self {
Self {
decoder,
failed: false,
}
}
session_methods!();
pub fn reinit(&mut self) {
self.decoder.stream.reset();
self.failed = false;
}
#[must_use]
pub fn into_decompressor(mut self) -> Decompressor {
self.decoder.stream.reset();
*self.decoder
}
}
#[cfg(test)]
mod tests {
use super::*;
use alloc::string::ToString;
use alloc::vec::Vec;
#[test]
fn failures_stick_and_counters_report_progress() {
let frame = openzl_ffi::compress(b"counted", openzl_ffi::Params::default()).unwrap();
let mut decoder = Decompressor::default();
let mut session = decoder.start();
let mut out = [0u8; 32];
let progress = session
.process(&frame, &mut out, DecodeOperation::Process)
.unwrap();
assert_eq!(progress.status, DecoderStatus::Finished);
assert!(session.is_finished());
assert_eq!(session.total_in(), frame.len() as u64);
assert_eq!((session.total_out(), session.frames_decoded()), (7, 1));
assert_eq!(
session.finish(&mut out).unwrap().status,
DecoderStatus::Finished
);
let mut session = decoder.start();
let failure = session
.process(b"junk", &mut out, DecodeOperation::Process)
.unwrap_err();
assert!(matches!(failure.error, DecodeError::InvalidData { .. }));
assert_eq!(
session.flush(&mut out).unwrap_err().into_error(),
DecodeError::InvalidState
);
assert!(!session.is_finished());
assert_eq!(DecodeError::from(failure.clone()), failure.error);
assert!(!failure.to_string().is_empty());
}
#[test]
fn owned_sessions_can_restart_after_failure() {
let frame = openzl_ffi::compress(b"again", openzl_ffi::Params::default()).unwrap();
let mut session = Decompressor::default().into_session();
let mut out = [0u8; 8];
assert!(
session
.process(b"xx", &mut out, DecodeOperation::Finish)
.is_err()
);
session.reinit();
assert_eq!(session.total_in(), 0);
let mut content = Vec::new();
let progress = session
.process(&frame, &mut out, DecodeOperation::Finish)
.unwrap();
content.extend_from_slice(&out[..progress.produced]);
assert_eq!(content, b"again");
assert_eq!(
session.finish(&mut out).unwrap().status,
DecoderStatus::Finished
);
let mut decoder = session.into_decompressor();
assert_eq!(decoder.decompress(&frame).unwrap(), b"again");
}
}