use core::borrow::BorrowMut;
use crate::common::MAX_BLOCK_SIZE;
use crate::decoding::errors::{FrameDecoderError, ReadFrameHeaderError};
use crate::decoding::{BlockDecodingStrategy, DictionaryHandle, FrameDecoder};
use crate::io::{Error, ErrorKind, Read};
pub struct StreamingDecoder<READ: Read, DEC: BorrowMut<FrameDecoder>> {
pub decoder: DEC,
source: READ,
forced_dictionary: bool,
position: Position,
input: Input,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct FinishedFrame {
pub compressed_size: u64,
pub decompressed_size: u64,
}
pub struct DictionaryRefused<READ, DEC = ()> {
source: READ,
decoder: DEC,
error: FrameDecoderError,
}
impl<READ, DEC> DictionaryRefused<READ, DEC> {
pub fn error(&self) -> &FrameDecoderError {
&self.error
}
pub fn into_source(self) -> READ {
self.source
}
pub fn into_parts(self) -> (READ, DEC, FrameDecoderError) {
(self.source, self.decoder, self.error)
}
}
impl<READ, DEC> core::fmt::Debug for DictionaryRefused<READ, DEC> {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
core::fmt::Debug::fmt(&self.error, f)
}
}
impl<READ, DEC> core::fmt::Display for DictionaryRefused<READ, DEC> {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
core::fmt::Display::fmt(&self.error, f)
}
}
#[cfg(feature = "std")]
impl<READ, DEC> std::error::Error for DictionaryRefused<READ, DEC> {
fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
std::error::Error::source(&self.error)
}
}
#[derive(Clone, Copy, PartialEq, Eq)]
enum Position {
BeforeFirstFrame,
BetweenFrames,
InFrame,
End,
}
const SKIP_CHUNK: u32 = 8 * 1024;
struct Input {
buf: alloc::vec::Vec<u8>,
start: usize,
end: usize,
skip_left: u32,
raw: Option<RawBlock>,
}
struct RawBlock {
size: u32,
left: u32,
last: bool,
#[cfg(all(feature = "lsm", feature = "hash"))]
digest: Option<twox_hash::XxHash64>,
}
const RAW_CHUNK: usize = 32 * 1024;
impl Input {
const fn new() -> Self {
Self {
buf: alloc::vec::Vec::new(),
start: 0,
end: 0,
skip_left: 0,
raw: None,
}
}
fn pending(&self) -> &[u8] {
&self.buf[self.start..self.end]
}
fn read_rest<R: Read>(&mut self, source: &mut R) -> Result<(), Error> {
if self.start > 0 {
self.buf.copy_within(self.start..self.end, 0);
self.end -= self.start;
self.start = 0;
}
self.buf.truncate(self.end);
let result = source.read_to_end(&mut self.buf);
self.end = self.buf.len();
result.map(drop)
}
fn read_raw<R: Read>(&mut self, source: &mut R, left: u32) -> Result<usize, Error> {
debug_assert!(
self.pending().is_empty(),
"a Raw block is read with nothing pending"
);
let chunk = left.min(RAW_CHUNK as u32) as usize;
if self.buf.len() < chunk {
self.buf.resize(chunk, 0);
}
loop {
match source.read(&mut self.buf[..chunk]) {
Err(e) if e.kind() == ErrorKind::Interrupted => {}
other => return other,
}
}
}
fn release(&mut self) {
self.buf = alloc::vec::Vec::new();
self.start = 0;
self.end = 0;
}
fn consume(&mut self, n: usize) {
debug_assert!(n <= self.end - self.start);
self.start += n;
if self.start == self.end {
self.start = 0;
self.end = 0;
}
}
fn fill<R: Read>(&mut self, source: &mut R, want: usize) -> Result<usize, Error> {
if self.start > 0 {
self.buf.copy_within(self.start..self.end, 0);
self.end -= self.start;
self.start = 0;
}
debug_assert!(self.end < want, "a fill must ask for more than is pending");
if self.buf.len() < want {
let len = want.max(2 * self.buf.len());
self.buf.resize(len, 0);
}
loop {
match source.read(&mut self.buf[self.end..want]) {
Ok(n) => {
self.end += n;
return Ok(n);
}
Err(e) if e.kind() == ErrorKind::Interrupted => {}
Err(e) => return Err(e),
}
}
}
}
fn header_len(bytes: &[u8], magicless: bool) -> usize {
let descriptor_at = if magicless {
0
} else {
let Some(magic) = bytes.first_chunk::<4>() else {
return 4;
};
let magic = u32::from_le_bytes(*magic);
if magic & 0xFFFF_FFF0 == 0x184D_2A50 {
return 8;
}
if magic != crate::common::MAGIC_NUM {
#[cfg(feature = "legacy")]
if let Some(version) = crate::decoding::LegacyVersion::of(bytes) {
return crate::decoding::legacy::header_len(version, bytes).unwrap_or(5);
}
return 4;
}
4
};
let Some(&descriptor) = bytes.get(descriptor_at) else {
return descriptor_at + 1;
};
descriptor_at + crate::decoding::frame::FrameDescriptor(descriptor).header_len()
}
impl<READ: Read, DEC: BorrowMut<FrameDecoder>> StreamingDecoder<READ, DEC> {
pub fn new_with_decoder(source: READ, decoder: DEC) -> StreamingDecoder<READ, DEC> {
StreamingDecoder {
decoder,
source,
forced_dictionary: false,
position: Position::BeforeFirstFrame,
input: Input::new(),
}
}
pub fn new_with_decoder_and_dictionary_handle(
source: READ,
mut decoder: DEC,
dict: &DictionaryHandle,
) -> Result<StreamingDecoder<READ, DEC>, DictionaryRefused<READ, DEC>> {
if let Err(error) = decoder.borrow_mut().arm_dict(dict) {
return Err(DictionaryRefused {
source,
decoder,
error,
});
}
Ok(StreamingDecoder {
decoder,
source,
forced_dictionary: true,
position: Position::BeforeFirstFrame,
input: Input::new(),
})
}
}
impl<READ: Read> StreamingDecoder<READ, FrameDecoder> {
pub fn new(source: READ) -> StreamingDecoder<READ, FrameDecoder> {
StreamingDecoder::new_with_decoder(source, FrameDecoder::new())
}
pub fn new_with_dictionary_handle(
source: READ,
dict: &DictionaryHandle,
) -> Result<StreamingDecoder<READ, FrameDecoder>, DictionaryRefused<READ>> {
Self::new_with_decoder_and_dictionary_handle(source, FrameDecoder::new(), dict).map_err(
|refused| DictionaryRefused {
source: refused.source,
decoder: (),
error: refused.error,
},
)
}
pub fn new_with_dictionary_bytes(
source: READ,
raw_dictionary: &[u8],
) -> Result<StreamingDecoder<READ, FrameDecoder>, DictionaryRefused<READ>> {
match DictionaryHandle::decode_dict(raw_dictionary) {
Ok(dict) => Self::new_with_dictionary_handle(source, &dict),
Err(error) => Err(DictionaryRefused {
source,
decoder: (),
error: FrameDecoderError::from(error),
}),
}
}
}
impl<READ: Read, DEC: BorrowMut<FrameDecoder>> StreamingDecoder<READ, DEC> {
pub fn get_ref(&self) -> &READ {
&self.source
}
pub fn get_mut(&mut self) -> &mut READ {
&mut self.source
}
pub fn decoder_mut(&mut self) -> &mut FrameDecoder {
self.decoder.borrow_mut()
}
pub fn into_inner(self) -> READ
where
READ: Sized,
{
self.source
}
pub fn into_parts(self) -> (READ, DEC)
where
READ: Sized,
{
(self.source, self.decoder)
}
pub fn into_frame_decoder(self) -> DEC {
self.decoder
}
fn cut_short(&mut self) -> Error {
let decoder = self.decoder.borrow_mut();
if decoder.awaits_checksum() {
return frame_error(FrameDecoderError::FailedToReadChecksum(Error::from(
ErrorKind::UnexpectedEof,
)));
}
let mut rest = self.input.pending();
decoder
.decode_blocks(&mut rest, BlockDecodingStrategy::All)
.map_or_else(frame_error, |_| Error::from(ErrorKind::UnexpectedEof))
}
fn start_frame(&mut self) -> Result<bool, Error> {
let how = if self.forced_dictionary {
FrameStart::HeldDictionary
} else {
FrameStart::Plain
};
let magicless = self.decoder.borrow_mut().is_magicless();
loop {
while self.input.skip_left > 0 {
let step = self.input.skip_left.min(SKIP_CHUNK);
if self.input.pending().is_empty()
&& self.input.fill(&mut self.source, step as usize)? == 0
{
return Err(frame_error(FrameDecoderError::FailedToSkipFrame));
}
let pending = u32::try_from(self.input.pending().len()).unwrap_or(u32::MAX);
let take = self.input.skip_left.min(pending);
self.input.consume(take as usize);
self.input.skip_left -= take;
}
loop {
let have = self.input.pending().len();
let want = header_len(self.input.pending(), magicless);
if have >= want {
break;
}
if self.input.fill(&mut self.source, want)? == 0 {
if have == 0 {
if self.position == Position::BeforeFirstFrame {
return Err(no_frame_error());
}
return Ok(false);
}
break;
}
}
let mut header = self.input.pending();
let before = header.len();
let started = how.start(self.decoder.borrow_mut(), &mut header);
let used = before - header.len();
self.input.consume(used);
match started {
Ok(()) => {
self.position = Position::InFrame;
return Ok(true);
}
Err(e) => match skippable_frame_length(&e) {
Some(length) => {
self.position = Position::BetweenFrames;
self.input.skip_left = length;
}
None => return Err(frame_error(e)),
},
}
}
}
}
fn skippable_frame_length(e: &FrameDecoderError) -> Option<u32> {
match e {
FrameDecoderError::ReadFrameHeaderError(ReadFrameHeaderError::SkipFrame {
length, ..
}) => Some(*length),
_ => None,
}
}
enum FrameStart {
Plain,
HeldDictionary,
}
impl FrameStart {
fn start(
&self,
decoder: &mut FrameDecoder,
source: impl Read,
) -> Result<(), FrameDecoderError> {
match self {
Self::Plain => decoder.init(source),
Self::HeldDictionary => decoder.reset_with_active_dict(source),
}
}
}
fn finish_raw_block(decoder: &mut FrameDecoder, raw: RawBlock) {
#[cfg(all(feature = "lsm", feature = "hash"))]
if let Some(digest) = raw.digest {
decoder.record_block_checksum(core::hash::Hasher::finish(&digest) as u32);
}
decoder.finish_raw_block(raw.size, raw.last);
}
fn frame_error(e: FrameDecoderError) -> Error {
#[cfg(feature = "std")]
return Error::other(e);
#[cfg(not(feature = "std"))]
return Error::new(ErrorKind::Other, alloc::boxed::Box::new(e));
}
fn no_frame_error() -> Error {
frame_error(FrameDecoderError::ReadFrameHeaderError(
ReadFrameHeaderError::MagicNumberReadError(Error::from(ErrorKind::UnexpectedEof)),
))
}
fn verify_finished_frame(decoder: &FrameDecoder) -> Result<(), Error> {
decoder.verify_content_size().map_err(frame_error)?;
#[cfg(feature = "hash")]
decoder.verify_content_checksum().map_err(frame_error)?;
Ok(())
}
impl<READ: Read, DEC: BorrowMut<FrameDecoder>> StreamingDecoder<READ, DEC> {
pub fn read_reporting_frames(
&mut self,
buf: &mut [u8],
mut finished: impl FnMut(FinishedFrame),
) -> Result<usize, Error> {
if buf.is_empty() || self.position == Position::End {
return Ok(0);
}
let mut written = 0;
loop {
if self.position != Position::InFrame {
debug_assert_eq!(written, 0, "a frame boundary is reached empty-handed");
if !self.start_frame()? {
self.position = Position::End;
return Ok(0);
}
continue;
}
let decoder = self.decoder.borrow_mut();
if decoder.is_finished() && decoder.can_collect() == 0 {
if written > 0 {
return Ok(written);
}
verify_finished_frame(decoder)?;
finished(FinishedFrame {
compressed_size: decoder.bytes_read_from_source(),
decompressed_size: decoder.decoded_size(),
});
self.position = Position::BetweenFrames;
continue;
}
if written == buf.len() {
if decoder.awaits_checksum() && decoder.can_collect() == 0 {
while self.input.pending().len() < 4 {
match self.input.fill(&mut self.source, 4) {
Ok(n) if n > 0 => {}
_ => break,
}
}
if self.input.pending().len() >= 4
&& let Ok((consumed, _)) =
decoder.decode_available(&self.input.pending()[..4], &mut [], true)
{
self.input.consume(consumed);
}
}
return Ok(written);
}
if self.input.raw.is_some() && decoder.can_collect() == 0 {
if written > 0 {
return Ok(written);
}
decoder.check_window_still_admitted().map_err(frame_error)?;
let left = self.input.raw.as_ref().map_or(0, |raw| raw.left);
let n = self.input.read_raw(&mut self.source, left)?;
let mut raw = self.input.raw.take().expect("checked just above");
if n == 0 {
return Err(frame_error(decoder.raw_block_cut_short(raw.size, raw.last)));
}
let content = &self.input.buf[..n];
decoder.raw_block_push(content);
#[cfg(all(feature = "lsm", feature = "hash"))]
if let Some(digest) = raw.digest.as_mut() {
core::hash::Hasher::write(digest, content);
}
raw.left -= n as u32;
if raw.left == 0 {
finish_raw_block(decoder, raw);
} else {
self.input.raw = Some(raw);
}
continue;
}
let want = decoder.input_needed(self.input.pending());
if self.input.raw.is_none()
&& want > self.input.pending().len()
&& decoder.can_collect() == 0
{
if written > 0 {
return Ok(written);
}
if let Some(&header) = self.input.pending().first_chunk::<3>()
&& self.input.pending().len() == 3
&& let Some((size, last)) = decoder
.start_raw_block(&header, true)
.map_err(frame_error)?
{
self.input.consume(3);
let raw = RawBlock {
size,
left: size,
last,
#[cfg(all(feature = "lsm", feature = "hash"))]
digest: decoder
.per_block_checksums_enabled()
.then(|| twox_hash::XxHash64::with_seed(0)),
};
if size == 0 {
finish_raw_block(decoder, raw);
} else {
self.input.raw = Some(raw);
}
continue;
}
if self.input.fill(&mut self.source, want)? == 0 {
return Err(self.cut_short());
}
continue;
}
let (consumed, produced) = decoder
.decode_available(self.input.pending(), &mut buf[written..], true)
.map_err(frame_error)?;
self.input.consume(consumed);
written += produced;
if consumed > 0 || produced > 0 {
continue;
}
if written > 0 {
return Ok(written);
}
return Err(self.cut_short());
}
}
}
impl<READ: Read, DEC: BorrowMut<FrameDecoder>> Read for StreamingDecoder<READ, DEC> {
fn read(&mut self, buf: &mut [u8]) -> Result<usize, Error> {
self.read_reporting_frames(buf, |_| {})
}
#[cfg(feature = "std")]
fn read_to_end(&mut self, output: &mut alloc::vec::Vec<u8>) -> Result<usize, Error> {
let start_total = output.len();
self.decode_rest_to_vec(output)?;
Ok(output.len() - start_total)
}
#[cfg(not(feature = "std"))]
fn read_to_end(&mut self, output: &mut alloc::vec::Vec<u8>) -> Result<(), Error> {
self.decode_rest_to_vec(output)
}
}
impl<READ: Read, DEC: BorrowMut<FrameDecoder>> StreamingDecoder<READ, DEC> {
fn decode_rest_to_vec(&mut self, output: &mut alloc::vec::Vec<u8>) -> Result<(), Error> {
let keep_dictionary = self.forced_dictionary;
if self.position == Position::BeforeFirstFrame && !self.start_frame()? {
self.position = Position::End;
}
match self.position {
Position::End => return Ok(()),
Position::InFrame
if self.input.raw.is_none() && {
let d = self.decoder.borrow_mut();
d.is_at_frame_start() && d.can_collect() == 0
} =>
{
self.input.read_rest(&mut self.source)?;
let decoded = self.decoder.borrow_mut().decode_current_frame_to_vec(
self.input.pending(),
output,
keep_dictionary,
);
self.input.release();
decoded.map_err(frame_error)?;
self.position = Position::End;
return Ok(());
}
Position::BeforeFirstFrame | Position::InFrame | Position::BetweenFrames => {}
}
loop {
let start = output.len();
output.resize(start + MAX_BLOCK_SIZE as usize, 0);
let n = match self.read(&mut output[start..]) {
Ok(n) => n,
Err(e) => {
output.truncate(start);
return Err(e);
}
};
output.truncate(start + n);
if n == 0 {
return Ok(());
}
}
}
}
#[cfg(test)]
mod tests;