Merge pull request #40 from newpavlov/inline_optim

Inline based optimization and overflow bug fixes

flif/src/components/header.rs

@@ -8,6 +8,7 @@ use numbers::chances::UpdateTable;
 use numbers::rac::RacRead;
 use numbers::symbol::UniformSymbolCoder;
 use numbers::FlifReadExt;
+use Limits;
 
 #[derive(Eq, PartialEq, Debug, Clone, Copy)]
 pub enum BytesPerChannel {
@@ -21,13 +22,39 @@ pub struct Header {
     pub interlaced: bool,
     pub channels: ColorSpace,
     pub bytes_per_channel: BytesPerChannel,
-    pub width: usize,
-    pub height: usize,
+    pub width: u32,
+    pub height: u32,
     pub num_frames: u32,
 }
 
+/// Helper function for reading width, height and num_frames
+fn read_varint<R: Read>(reader: &mut R, delta: u32) -> Result<u32> {
+    reader.read_varint::<u32>().and_then(|v| {
+        v.checked_add(delta).ok_or_else(|| {
+            Error::LimitViolation("number of pixels exceeds limit: overflow".to_string())
+        })
+    })
+}
+
+/// Check if number of pixels uphelds provided limit
+fn check_limit(width: u32, height: u32, frames: u32, limit: u32) -> Result<()> {
+    let pixels = frames
+        .checked_mul(width)
+        .and_then(|val| val.checked_mul(height));
+    match pixels {
+        Some(pix) if pix > limit => Err(Error::LimitViolation(format!(
+            "number of pixels exceeds limit: {}/{}",
+            pix, limit,
+        ))),
+        None => Err(Error::LimitViolation(
+            "number of pixels exceeds limit: overflow".to_string(),
+        )),
+        Some(_) => Ok(()),
+    }
+}
+
 impl Header {
-    pub(crate) fn from_reader<R: Read>(mut reader: R) -> Result<Self> {
+    pub(crate) fn from_reader<R: Read>(mut reader: R, limits: &Limits) -> Result<Self> {
         // first read in some magic
         let mut magic_buf = [0; 4];
         reader.read_exact(&mut magic_buf)?;
@@ -67,15 +94,17 @@ impl Header {
                 desc: "bytes per channel was not a valid value",
             })?,
         };
-        let width = 1 + reader.read_varint::<usize>()?;
-        let height = 1 + reader.read_varint::<usize>()?;
+        let width = read_varint(&mut reader, 1)?;
+        let height = read_varint(&mut reader, 1)?;
 
         let num_frames = if animated {
-            2 + reader.read_varint::<u32>()?
+            read_varint(&mut reader, 2)?
         } else {
             1
         };
 
+        check_limit(width, height, num_frames, limits.pixels)?;
+
         Ok(Header {
             interlaced,
             channels,

flif/src/components/metadata.rs

@@ -29,13 +29,13 @@ impl Metadata {
         mut reader: R,
         limits: &Limits,
     ) -> Result<(Vec<Metadata>, u8)> {
-        let mut ret = Vec::with_capacity(limits.metadata_count);
+        let mut ret = Vec::with_capacity(limits.metadata_count as usize);
         let required_type = loop {
             match Self::from_reader(&mut reader, limits)? {
                 MetadataType::Optional(metadata) => ret.push(metadata),
                 MetadataType::Required(byte) => break byte,
             }
-            if ret.len() > limits.metadata_count {
+            if ret.len() > limits.metadata_count as usize {
                 Err(Error::LimitViolation(format!(
                     "number of metadata entries exceeds limit: {}",
                     limits.metadata_count,
@@ -70,7 +70,7 @@ impl Metadata {
         };
 
         let chunk_size = reader.read_varint()?;
-        if chunk_size > limits.metadata_chunk {
+        if chunk_size > limits.metadata_chunk as usize {
             Err(Error::LimitViolation(format!(
                 "requested metadata chunk size exceeds limit: {} vs {}",
                 chunk_size, limits.metadata_chunk,

flif/src/decoder.rs

@@ -89,15 +89,7 @@ impl<R: Read> Decoder<R> {
 
 fn identify_internal<R: Read>(mut reader: R, limits: Limits) -> Result<(FlifInfo, Rac<R>)> {
     // read the first header
-    let main_header = Header::from_reader(&mut reader)?;
-    let frames = main_header.num_frames as usize;
-    let pixels = main_header.width * main_header.height * frames;
-    if pixels > limits.pixels {
-        Err(Error::LimitViolation(format!(
-            "number of pixels exceeds limit: {} vs {}",
-            pixels, limits.pixels,
-        )))?
-    }
+    let main_header = Header::from_reader(&mut reader, &limits)?;
 
     // read the metadata chunks
     let (metadata, non_optional_byte) = Metadata::all_from_reader(&mut reader, &limits)?;

flif/src/decoding_image.rs

@@ -10,8 +10,8 @@ use FlifInfo;
 pub use decoder::Decoder;
 
 pub(crate) struct DecodingImage {
-    height: usize,
-    width: usize,
+    height: u32,
+    width: u32,
     channels: ColorSpace,
     data: Vec<Pixel>,
 }
@@ -49,35 +49,36 @@ type Maniac<'a> = ChannelSet<Option<ManiacTree<'a>>>;
 // safety criterias defined by `debug_assert`s
 impl DecodingImage {
     pub fn new(info: &FlifInfo) -> DecodingImage {
+        let pixels = (info.header.height * info.header.width) as usize;
         DecodingImage {
             height: info.header.height,
             width: info.header.width,
             channels: info.header.channels,
-            data: vec![Pixel::default(); info.header.height * info.header.width],
+            data: vec![Pixel::default(); pixels],
         }
     }
 
-    fn get_idx(&self, x: usize, y: usize) -> usize {
-        (self.width * y) + x
+    fn check_data(&self) -> bool {
+        self.data.len() == (self.width * self.height) as usize
+    }
+
+    fn get_idx(&self, x: u32, y: u32) -> usize {
+        ((self.width * y) + x) as usize
     }
 
     pub fn get_data(&self) -> &[Pixel] {
         &self.data
     }
 
-    unsafe fn get_val(&self, x: usize, y: usize, chan: Channel) -> ColorValue {
-        debug_assert!(
-            x < self.width && y < self.height && self.data.len() == self.width * self.height
-        );
+    unsafe fn get_val(&self, x: u32, y: u32, chan: Channel) -> ColorValue {
+        debug_assert!(x < self.width && y < self.height && self.check_data());
         self.data.get_unchecked(self.get_idx(x, y))[chan]
     }
 
-    unsafe fn get_edge_vicinity(&self, x: usize, y: usize, chan: Channel) -> EdgePixelVicinity {
-        debug_assert!(
-            x < self.width && y < self.height && self.data.len() == self.width * self.height
-        );
+    unsafe fn get_edge_vicinity(&self, x: u32, y: u32, chan: Channel) -> EdgePixelVicinity {
+        debug_assert!(x < self.width && y < self.height && self.check_data());
         EdgePixelVicinity {
-            pixel: *self.data.get_unchecked((self.width * y) + x),
+            pixel: *self.data.get_unchecked(self.get_idx(x, y)),
             is_rgba: self.channels == ColorSpace::RGBA,
             chan,
             top: if y != 0 {
@@ -113,16 +114,10 @@ impl DecodingImage {
         }
     }
 
-    unsafe fn get_core_vicinity(&self, x: usize, y: usize, chan: Channel) -> CorePixelVicinity {
-        debug_assert!(
-            x < self.width - 1
-                && y < self.height
-                && x > 1
-                && y > 1
-                && self.data.len() == self.width * self.height
-        );
+    unsafe fn get_core_vicinity(&self, x: u32, y: u32, chan: Channel) -> CorePixelVicinity {
+        debug_assert!(x < self.width - 1 && y < self.height && x > 1 && y > 1 && self.check_data());
         CorePixelVicinity {
-            pixel: *self.data.get_unchecked((self.width * y) + x),
+            pixel: *self.data.get_unchecked(self.get_idx(x, y)),
             chan,
             is_rgba: self.channels == ColorSpace::RGBA,
             top: self.get_val(x, y - 1, chan),
@@ -136,8 +131,8 @@ impl DecodingImage {
 
     unsafe fn process_edge_pixel<E, R: Read>(
         &mut self,
-        x: usize,
-        y: usize,
+        x: u32,
+        y: u32,
         chan: Channel,
         maniac: &mut Maniac,
         rac: &mut Rac<R>,
@@ -146,9 +141,7 @@ impl DecodingImage {
     where
         E: FnMut(EdgePixelVicinity, &mut Maniac, &mut Rac<R>) -> Result<ColorValue>,
     {
-        debug_assert!(
-            x < self.width && y < self.height && self.data.len() == self.width * self.height
-        );
+        debug_assert!(x < self.width && y < self.height && self.check_data());
         let pix_vic = self.get_edge_vicinity(x, y, chan);
         let val = edge_f(pix_vic, maniac, rac)?;
         let idx = self.get_idx(x, y);
@@ -172,7 +165,7 @@ impl DecodingImage {
     {
         let width = self.width;
         let height = self.height;
-        debug_assert!(self.data.len() == width * height);
+        debug_assert!(self.check_data());
         // special case for small images
         if width <= 3 || height <= 2 {
             for y in 0..height {

flif/src/lib.rs

@@ -64,7 +64,7 @@ impl Flif {
         };
         let width = self.info.header.width;
         let height = self.info.header.height;
-        let mut data = Vec::with_capacity(n * width * height);
+        let mut data = Vec::with_capacity((n * width * height) as usize);
 
         for vals in self.image_data.get_data().iter() {
             for channel in self.info.header.channels {
@@ -80,13 +80,13 @@ impl Flif {
 #[derive(Copy, Clone, Debug)]
 pub struct Limits {
     /// max size of the compressed metadata in bytes (default: 1 MB)
-    pub metadata_chunk: usize,
+    pub metadata_chunk: u32,
     /// max number of metadata entries (default: 8)
-    pub metadata_count: usize,
-    /// max number of pixels: `width * height * frames` (default: 2<sup>26</sup>)
-    pub pixels: usize,
-    /// max number of MANIAC nodes (default: 4096)
-    pub maniac_nodes: usize,
+    pub metadata_count: u32,
+    /// max number of pixels: `width * height * frames` (default: 67M = 2<sup>26</sup>)
+    pub pixels: u32,
+    /// max number of MANIAC nodes (default: 16384 = 2<sup>14</sup>)
+    pub maniac_nodes: u32,
 }
 
 impl Default for Limits {
@@ -95,7 +95,7 @@ impl Default for Limits {
             metadata_chunk: 1 << 20,
             metadata_count: 8,
             pixels: 1 << 26,
-            maniac_nodes: 4096,
+            maniac_nodes: 1 << 14,
         }
     }
 }

flif/src/maniac/mod.rs

@@ -109,7 +109,7 @@ impl<'a> ManiacTree<'a> {
                 _ => break,
             };
 
-            if result_vec.len() > limits.maniac_nodes {
+            if result_vec.len() > limits.maniac_nodes as usize {
                 Err(Error::LimitViolation(format!(
                     "number of maniac nodes exceeds limit"
                 )))?;

flif/src/numbers/near_zero.rs

@@ -3,6 +3,8 @@ use num_traits::PrimInt;
 use numbers::chances::{ChanceTable, ChanceTableEntry};
 use numbers::rac::RacRead;
 
+use std::cmp;
+
 pub trait NearZeroCoder {
     fn read_near_zero<I: PrimInt>(
         &mut self,
@@ -19,16 +21,14 @@ impl<R: RacRead> NearZeroCoder for R {
         max: I,
         context: &mut ChanceTable,
     ) -> Result<I> {
-        if min > I::zero() {
-            Ok(read_near_zero_inner(self, I::zero(), max - min, context)? + min)
-        } else if max < I::zero() {
-            Ok(read_near_zero_inner(self, min - max, I::zero(), context)? + max)
-        } else {
-            read_near_zero_inner(self, min, max, context)
-        }
+        let delta = cmp::min(max, cmp::max(I::zero(), min));
+        let min = min - delta;
+        let max = max - delta;
+        Ok(read_near_zero_inner(self, min, max, context)? + delta)
     }
 }
 
+#[inline(always)]
 fn read_near_zero_inner<R: RacRead, I: PrimInt>(
     read: &mut R,
     min: I,

flif/src/numbers/rac.rs

@@ -57,6 +57,7 @@ impl<R: Read> RacRead for Rac<R> {
         self.get(chance)
     }
 
+    #[inline(always)]
     fn read(&mut self, context: &mut ChanceTable, entry: ChanceTableEntry) -> Result<bool> {
         let chance = context.get_chance(entry);
         let transformed_chance = Self::apply_chance(u32::from(chance), self.range);