Skip to content

Commit

Permalink
Use a single Vec between Decoder and ZlibStream
Browse files Browse the repository at this point in the history
  • Loading branch information
@fintelia
fintelia committed Jan 20, 2024
1 parent ed54082 commit f3f819e
Show file tree
Hide file tree
Showing 4 changed files with 82 additions and 142 deletions.
2 changes: 1 addition & 1 deletion examples/pngcheck.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -174,7 +174,7 @@ fn check_image<P: AsRef<Path>>(c: Config, fname: P) -> io::Result<()> {
}
buf = &data[..n];
}
match decoder.update(buf, &mut Vec::new()) {
match decoder.update(buf, &mut Vec::new(), &mut 0) {
Ok((_, ImageEnd)) => {
if !have_idat {
// This isn't beautiful. But it works.
Expand Down
83 changes: 56 additions & 27 deletions src/decoder/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -204,7 +204,7 @@ impl<R: Read> Decoder<R> {
let mut buf = Vec::new();
while self.read_decoder.info().is_none() {
buf.clear();
if self.read_decoder.decode_next(&mut buf)?.is_none() {
if self.read_decoder.decode_next(&mut buf, &mut 0)?.is_none() {
return Err(DecodingError::Format(
FormatErrorInner::UnexpectedEof.into(),
));
Expand All @@ -226,6 +226,8 @@ impl<R: Read> Decoder<R> {
data_stream: Vec::new(),
prev_start: 0,
current_start: 0,
lookback_start: 0,
write_start: 0,
transform: self.transform,
scratch_buffer: Vec::new(),
};
Expand Down Expand Up @@ -295,7 +297,7 @@ struct ReadDecoder<R: Read> {
impl<R: Read> ReadDecoder<R> {
/// Returns the next decoded chunk. If the chunk is an ImageData chunk, its contents are written
/// into image_data.
fn decode_next(&mut self, image_data: &mut Vec<u8>) -> Result<Option<Decoded>, DecodingError> {
fn decode_next(&mut self, image_data: &mut Vec<u8>, out_pos: &mut usize) -> Result<Option<Decoded>, DecodingError> {
while !self.at_eof {
let (consumed, result) = {
let buf = self.reader.fill_buf()?;
Expand All @@ -304,7 +306,7 @@ impl<R: Read> ReadDecoder<R> {
FormatErrorInner::UnexpectedEof.into(),
));
}
self.decoder.update(buf, image_data)?
self.decoder.update(buf, image_data, out_pos)?
};
self.reader.consume(consumed);
match result {
Expand All @@ -324,7 +326,7 @@ impl<R: Read> ReadDecoder<R> {
FormatErrorInner::UnexpectedEof.into(),
));
}
let (consumed, event) = self.decoder.update(buf, &mut vec![])?;
let (consumed, event) = self.decoder.update(buf, &mut vec![], &mut 0)?;
self.reader.consume(consumed);
match event {
Decoded::Nothing => (),
Expand Down Expand Up @@ -365,6 +367,10 @@ pub struct Reader<R: Read> {
prev_start: usize,
/// Index in `data_stream` where the current row starts.
current_start: usize,
/// Index in `data_stream` of the first byte still needed by deflate lookback.
lookback_start: usize,
/// Index in `data_stream` where new data can be written.
write_start: usize,
/// Output transformations
transform: Transformations,
/// This buffer is only used so that `next_row` and `next_interlaced_row` can return reference
Expand Down Expand Up @@ -415,7 +421,7 @@ impl<R: Read> Reader<R> {
// know that we will stop before reading any image data from the stream. Thus pass an
// empty buffer and assert that remains empty.
let mut buf = Vec::new();
let state = self.decoder.decode_next(&mut buf)?;
let state = self.decoder.decode_next(&mut buf, &mut 0)?;
assert!(buf.is_empty());

match state {
Expand Down Expand Up @@ -514,6 +520,8 @@ impl<R: Read> Reader<R> {
self.data_stream.clear();
self.current_start = 0;
self.prev_start = 0;
self.lookback_start = 0;
self.write_start = 0;
let width = self.info().width;
if self.info().interlaced {
while let Some(InterlacedRow {
Expand Down Expand Up @@ -607,7 +615,7 @@ impl<R: Read> Reader<R> {
self.prev_start = 0;
loop {
let mut buf = Vec::new();
let state = self.decoder.decode_next(&mut buf)?;
let state = self.decoder.decode_next(&mut buf, &mut 0)?;

if state.is_none() {
break;
Expand Down Expand Up @@ -749,44 +757,65 @@ impl<R: Read> Reader<R> {
///
/// The scanline is filtered against the previous scanline according to the specification.
fn next_raw_interlaced_row(&mut self, rowlen: usize) -> Result<(), DecodingError> {
// Read image data until we have at least one full row (but possibly more than one).
while self.data_stream.len() - self.current_start < rowlen {
if self.subframe.consumed_and_flushed {
return Err(DecodingError::Format(
FormatErrorInner::NoMoreImageData.into(),
));
}

// Read image data if we don't yet have enough data for the next row.
if self.lookback_start - self.current_start < rowlen {
// Clear the current buffer before appending more data.
if self.prev_start > 0 {
self.data_stream.copy_within(self.prev_start.., 0);
self.data_stream
.truncate(self.data_stream.len() - self.prev_start);
self.write_start -= self.prev_start;
self.lookback_start -= self.prev_start;
self.current_start -= self.prev_start;
self.prev_start = 0;
}

match self.decoder.decode_next(&mut self.data_stream)? {
Some(Decoded::ImageData) => {}
Some(Decoded::ImageDataFlushed) => {
self.subframe.consumed_and_flushed = true;
}
None => {
let image_bytes = self
.subframe
.rowlen
.saturating_mul(self.subframe.height as usize);
let target_bytes = (256 << 10).min(image_bytes);

self.data_stream.resize(self.data_stream.len().max(target_bytes), 0);

while self.lookback_start - self.current_start < target_bytes {
if self.subframe.consumed_and_flushed {
return Err(DecodingError::Format(
if self.data_stream.is_empty() {
FormatErrorInner::NoMoreImageData
FormatErrorInner::NoMoreImageData.into(),
));
}

match self.decoder.decode_next(&mut self.data_stream, &mut self.write_start)? {
Some(Decoded::ImageData) => {
// TODO: interlaced
if self.data_stream.len() >= image_bytes {
self.lookback_start = self.write_start;
} else {
FormatErrorInner::UnexpectedEndOfChunk
self.lookback_start = self.write_start.saturating_sub(32 << 10);
}
.into(),
));
}
Some(Decoded::ImageDataFlushed) => {
self.lookback_start = self.write_start;
self.subframe.consumed_and_flushed = true;
break;
}
None => {
return Err(DecodingError::Format(
if self.data_stream.is_empty() {
FormatErrorInner::NoMoreImageData
} else {
FormatErrorInner::UnexpectedEndOfChunk
}
.into(),
));
}
_ => (),
}
_ => (),
}
}

// Get a reference to the current row and point scan_start to the next one.
let (prev, row) = self.data_stream.split_at_mut(self.current_start);
let (prev, row) = self.data_stream[..self.lookback_start].split_at_mut(self.current_start);

// Unfilter the row.
let filter = FilterType::from_u8(row[0]).ok_or(DecodingError::Format(
Expand Down
13 changes: 8 additions & 5 deletions src/decoder/stream.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -581,10 +581,11 @@ impl StreamingDecoder {
&mut self,
mut buf: &[u8],
image_data: &mut Vec<u8>,
out_pos: &mut usize,
) -> Result<(usize, Decoded), DecodingError> {
let len = buf.len();
while !buf.is_empty() && self.state.is_some() {
match self.next_state(buf, image_data) {
match self.next_state(buf, image_data, out_pos) {
Ok((bytes, Decoded::Nothing)) => buf = &buf[bytes..],
Ok((bytes, result)) => {
buf = &buf[bytes..];
Expand All @@ -600,6 +601,7 @@ impl StreamingDecoder {
&mut self,
buf: &[u8],
image_data: &mut Vec<u8>,
out_pos: &mut usize,
) -> Result<(usize, Decoded), DecodingError> {
use self::State::*;

Expand All @@ -620,7 +622,7 @@ impl StreamingDecoder {
// values is that they occur fairly frequently and special-casing them results
// in performance gains.
const CONSUMED_BYTES: usize = 4;
self.parse_u32(kind, &buf[0..4], image_data)
self.parse_u32(kind, &buf[0..4], image_data, out_pos)
.map(|decoded| (CONSUMED_BYTES, decoded))
} else {
let remaining_count = 4 - accumulated_count;
Expand All @@ -641,7 +643,7 @@ impl StreamingDecoder {
Ok((consumed_bytes, Decoded::Nothing))
} else {
debug_assert_eq!(accumulated_count, 4);
self.parse_u32(kind, &bytes, image_data)
self.parse_u32(kind, &bytes, image_data, out_pos)
.map(|decoded| (consumed_bytes, decoded))
}
}
Expand Down Expand Up @@ -700,7 +702,7 @@ impl StreamingDecoder {
debug_assert!(type_str == IDAT || type_str == chunk::fdAT);
let len = std::cmp::min(buf.len(), self.current_chunk.remaining as usize);
let buf = &buf[..len];
let consumed = self.inflater.decompress(buf, image_data)?;
let consumed = self.inflater.decompress(buf, image_data, out_pos)?;
self.current_chunk.crc.update(&buf[..consumed]);
self.current_chunk.remaining -= consumed as u32;
if self.current_chunk.remaining == 0 {
Expand All @@ -718,6 +720,7 @@ impl StreamingDecoder {
kind: U32ValueKind,
u32_be_bytes: &[u8],
image_data: &mut Vec<u8>,
out_pos: &mut usize,
) -> Result<Decoded, DecodingError> {
debug_assert_eq!(u32_be_bytes.len(), 4);
let bytes = u32_be_bytes.try_into().unwrap();
Expand Down Expand Up @@ -759,7 +762,7 @@ impl StreamingDecoder {
&& (self.current_chunk.type_ == IDAT || self.current_chunk.type_ == chunk::fdAT)
{
self.current_chunk.type_ = type_str;
self.inflater.finish_compressed_chunks(image_data)?;
self.inflater.finish_compressed_chunks(image_data, out_pos)?;
self.inflater.reset();
self.state = Some(State::U32 {
kind,
Expand Down

0 comments on commit f3f819e

Please sign in to comment.