diff --git a/src/deflate/encode.rs b/src/deflate/encode.rs index 031d851..d944c13 100644 --- a/src/deflate/encode.rs +++ b/src/deflate/encode.rs @@ -400,7 +400,7 @@ where { self.lz77.flush(&mut self.buf); self.buf.push(symbol::Symbol::EndOfBlock); - let symbol_encoder = self.huffman.build(&self.buf); + let symbol_encoder = self.huffman.build(&self.buf)?; self.huffman.save(writer, &symbol_encoder)?; for s in self.buf.drain(..) { symbol_encoder.encode(writer, &s)?; diff --git a/src/deflate/symbol.rs b/src/deflate/symbol.rs index 1c670a5..66a5a8c 100644 --- a/src/deflate/symbol.rs +++ b/src/deflate/symbol.rs @@ -232,7 +232,7 @@ impl Decoder { } pub trait HuffmanCodec { - fn build(&self, symbols: &[Symbol]) -> Encoder; + fn build(&self, symbols: &[Symbol]) -> io::Result; fn save(&self, writer: &mut bit::BitWriter, codec: &Encoder) -> io::Result<()> where W: io::Write; @@ -245,7 +245,7 @@ pub trait HuffmanCodec { pub struct FixedHuffmanCodec; impl HuffmanCodec for FixedHuffmanCodec { #[allow(unused_variables)] - fn build(&self, symbols: &[Symbol]) -> Encoder { + fn build(&self, symbols: &[Symbol]) -> io::Result { let mut literal_builder = huffman::EncoderBuilder::new(288); for &(bitwidth, ref symbols, code_base) in &FIXED_LITERAL_OR_LENGTH_CODE_TABLE { for (code, symbol) in symbols @@ -253,19 +253,19 @@ impl HuffmanCodec for FixedHuffmanCodec { .enumerate() .map(|(i, s)| (code_base + i as u16, s)) { - literal_builder.set_mapping(symbol, huffman::Code::new(bitwidth, code)); + literal_builder.set_mapping(symbol, huffman::Code::new(bitwidth, code))?; } } let mut distance_builder = huffman::EncoderBuilder::new(30); for i in 0..30 { - distance_builder.set_mapping(i, huffman::Code::new(5, i)); + distance_builder.set_mapping(i, huffman::Code::new(5, i))?; } - Encoder { + Ok(Encoder { literal: literal_builder.finish(), distance: distance_builder.finish(), - } + }) } #[allow(unused_variables)] fn save(&self, writer: &mut bit::BitWriter, codec: &Encoder) -> io::Result<()> @@ -286,13 +286,13 @@ impl HuffmanCodec for FixedHuffmanCodec { .enumerate() .map(|(i, s)| (code_base + i as u16, s)) { - literal_builder.set_mapping(symbol, huffman::Code::new(bitwidth, code)); + literal_builder.set_mapping(symbol, huffman::Code::new(bitwidth, code))?; } } let mut distance_builder = huffman::DecoderBuilder::new(5, None); for i in 0..30 { - distance_builder.set_mapping(i, huffman::Code::new(5, i)); + distance_builder.set_mapping(i, huffman::Code::new(5, i))?; } Ok(Decoder { @@ -305,7 +305,7 @@ impl HuffmanCodec for FixedHuffmanCodec { #[derive(Debug)] pub struct DynamicHuffmanCodec; impl HuffmanCodec for DynamicHuffmanCodec { - fn build(&self, symbols: &[Symbol]) -> Encoder { + fn build(&self, symbols: &[Symbol]) -> io::Result { let mut literal_counts = [0; 286]; let mut distance_counts = [0; 30]; for s in symbols { @@ -314,10 +314,10 @@ impl HuffmanCodec for DynamicHuffmanCodec { distance_counts[d as usize] += 1; } } - Encoder { - literal: huffman::EncoderBuilder::from_frequencies(&literal_counts, 15), - distance: huffman::EncoderBuilder::from_frequencies(&distance_counts, 15), - } + Ok(Encoder { + literal: huffman::EncoderBuilder::from_frequencies(&literal_counts, 15)?, + distance: huffman::EncoderBuilder::from_frequencies(&distance_counts, 15)?, + }) } fn save(&self, writer: &mut bit::BitWriter, codec: &Encoder) -> io::Result<()> where @@ -331,7 +331,7 @@ impl HuffmanCodec for DynamicHuffmanCodec { for x in &codes { code_counts[x.0 as usize] += 1; } - let bitwidth_encoder = huffman::EncoderBuilder::from_frequencies(&code_counts, 7); + let bitwidth_encoder = huffman::EncoderBuilder::from_frequencies(&code_counts, 7)?; let bitwidth_code_count = cmp::max( 4, @@ -379,7 +379,7 @@ impl HuffmanCodec for DynamicHuffmanCodec { bitwidth_code_bitwidthes[i] = reader.read_bits(3)? as u8; } let bitwidth_decoder = - huffman::DecoderBuilder::from_bitwidthes(&bitwidth_code_bitwidthes, None); + huffman::DecoderBuilder::from_bitwidthes(&bitwidth_code_bitwidthes, None)?; let mut literal_code_bitwidthes = Vec::with_capacity(literal_code_count as usize); while literal_code_bitwidthes.len() < literal_code_count as usize { @@ -412,8 +412,8 @@ impl HuffmanCodec for DynamicHuffmanCodec { literal: huffman::DecoderBuilder::from_bitwidthes( &literal_code_bitwidthes, Some(END_OF_BLOCK), - ), - distance: huffman::DecoderBuilder::from_bitwidthes(&distance_code_bitwidthes, None), + )?, + distance: huffman::DecoderBuilder::from_bitwidthes(&distance_code_bitwidthes, None)?, }) } } diff --git a/src/huffman.rs b/src/huffman.rs index 7a68bf1..b7a1a6b 100644 --- a/src/huffman.rs +++ b/src/huffman.rs @@ -34,9 +34,9 @@ impl Code { pub trait Builder: Sized { type Instance; - fn set_mapping(&mut self, symbol: u16, code: Code); + fn set_mapping(&mut self, symbol: u16, code: Code) -> io::Result<()>; fn finish(self) -> Self::Instance; - fn restore_canonical_huffman_codes(mut self, bitwidthes: &[u8]) -> Self::Instance { + fn restore_canonical_huffman_codes(mut self, bitwidthes: &[u8]) -> io::Result { debug_assert!(!bitwidthes.is_empty()); let mut symbols = bitwidthes @@ -51,11 +51,11 @@ pub trait Builder: Sized { let mut prev_width = 0; for (symbol, bitwidth) in symbols { code <<= bitwidth - prev_width; - self.set_mapping(symbol, Code::new(bitwidth, code)); + self.set_mapping(symbol, Code::new(bitwidth, code))?; code += 1; prev_width = bitwidth; } - self.finish() + Ok(self.finish()) } } @@ -75,14 +75,14 @@ impl DecoderBuilder { max_bitwidth: max_bitwidth, } } - pub fn from_bitwidthes(bitwidthes: &[u8], eob_symbol: Option) -> Decoder { + pub fn from_bitwidthes(bitwidthes: &[u8], eob_symbol: Option) -> io::Result { let builder = Self::new(bitwidthes.iter().cloned().max().unwrap_or(0), eob_symbol); builder.restore_canonical_huffman_codes(bitwidthes) } } impl Builder for DecoderBuilder { type Instance = Decoder; - fn set_mapping(&mut self, symbol: u16, code: Code) { + fn set_mapping(&mut self, symbol: u16, code: Code) -> io::Result<()> { debug_assert!(code.width <= self.max_bitwidth); if Some(symbol) == self.eob_symbol { self.eob_bitwidth = code.width; @@ -95,11 +95,18 @@ impl Builder for DecoderBuilder { let code_be = code.inverse_endian(); for padding in 0..(1 << (self.max_bitwidth - code.width)) { let i = ((padding << code.width) | code_be.bits) as usize; - debug_assert_eq!(self.table[i], u16::from(MAX_BITWIDTH) + 1); + if self.table[i] != u16::from(MAX_BITWIDTH) + 1 { + let message = format!( + "Bit region conflict: i={}, old_value={}, new_value={}, symbol={}, code={:?}", + i, self.table[i], value, symbol, code + ); + return Err(io::Error::new(io::ErrorKind::InvalidData, message)); + } unsafe { *self.table.get_unchecked_mut(i) = value; } } + Ok(()) } fn finish(self) -> Self::Instance { Decoder { @@ -158,7 +165,7 @@ impl EncoderBuilder { table: vec![Code::new(0, 0); symbol_count], } } - pub fn from_bitwidthes(bitwidthes: &[u8]) -> Encoder { + pub fn from_bitwidthes(bitwidthes: &[u8]) -> io::Result { let symbol_count = bitwidthes .iter() .enumerate() @@ -168,7 +175,7 @@ impl EncoderBuilder { let builder = Self::new(symbol_count); builder.restore_canonical_huffman_codes(bitwidthes) } - pub fn from_frequencies(symbol_frequencies: &[usize], max_bitwidth: u8) -> Encoder { + pub fn from_frequencies(symbol_frequencies: &[usize], max_bitwidth: u8) -> io::Result { let max_bitwidth = cmp::min( max_bitwidth, ordinary_huffman_codes::calc_optimal_max_bitwidth(symbol_frequencies), @@ -179,9 +186,10 @@ impl EncoderBuilder { } impl Builder for EncoderBuilder { type Instance = Encoder; - fn set_mapping(&mut self, symbol: u16, code: Code) { + fn set_mapping(&mut self, symbol: u16, code: Code) -> io::Result<()> { debug_assert_eq!(self.table[symbol as usize], Code::new(0, 0)); self.table[symbol as usize] = code.inverse_endian(); + Ok(()) } fn finish(self) -> Self::Instance { Encoder { table: self.table }