Make huffman decoder/encoder building function return io::Result<_>

sile · sile · commit dc77163c1cdd · 2018-07-02T03:02:19.000+09:00
This resolve RazrFalcon/usvg#22
diff --git 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
@@ -232,7 +232,7 @@ impl Decoder {
 }
 
 pub trait HuffmanCodec {
-    fn build(&self, symbols: &[Symbol]) -> Encoder;
+    fn build(&self, symbols: &[Symbol]) -> io::Result<Encoder>;
     fn save<W>(&self, writer: &mut bit::BitWriter<W>, codec: &Encoder) -> io::Result<()>
     where
         W: io::Write;
@@ -245,27 +245,27 @@ 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<Encoder> {
         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
                 .clone()
                 .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<W>(&self, writer: &mut bit::BitWriter<W>, 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<Encoder> {
         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<W>(&self, writer: &mut bit::BitWriter<W>, 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
@@ -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<Self::Instance> {
         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<u16>) -> Decoder {
+    pub fn from_bitwidthes(bitwidthes: &[u8], eob_symbol: Option<u16>) -> io::Result<Decoder> {
         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<Encoder> {
         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<Encoder> {
         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 }

Original file line number	Diff line number	Diff line change
`@@ -400,7 +400,7 @@ where`
`400`	`400`	`{`
`401`	`401`	`self.lz77.flush(&mut self.buf);`
`402`	`402`	`self.buf.push(symbol::Symbol::EndOfBlock);`
`403`		`- let symbol_encoder = self.huffman.build(&self.buf);`
	`403`	`+ let symbol_encoder = self.huffman.build(&self.buf)?;`
`404`	`404`	`self.huffman.save(writer, &symbol_encoder)?;`
`405`	`405`	`for s in self.buf.drain(..) {`
`406`	`406`	`symbol_encoder.encode(writer, &s)?;`