Made decompression work across row boundaries (#5)

Author: jorgecarleitao

src/error.rs

@@ -0,0 +1,25 @@
+use crate::proto::stream::Kind;
+
+#[derive(Debug, Clone)]
+pub enum Error {
+    OutOfSpec,
+    RleLiteralTooLarge,
+    InvalidUtf8,
+    InvalidColumn(u32),
+    InvalidKind(u32, Kind),
+    DecodeFloat,
+    Decompression,
+    InvalidProto,
+}
+
+impl From<prost::DecodeError> for Error {
+    fn from(_: prost::DecodeError) -> Self {
+        Self::InvalidProto
+    }
+}
+
+impl From<std::io::Error> for Error {
+    fn from(_: std::io::Error) -> Self {
+        Self::OutOfSpec
+    }
+}

src/lib.rs

@@ -1,26 +1,5 @@
-use proto::stream::Kind;
-
+pub mod error;
 pub mod proto;
 pub mod read;
 
 pub use fallible_streaming_iterator;
-
-#[derive(Debug, Clone)]
-pub enum Error {
-    OutOfSpec,
-    RleLiteralTooLarge,
-    InvalidUtf8,
-    InvalidColumn(u32, Kind),
-}
-
-impl From<prost::DecodeError> for Error {
-    fn from(_: prost::DecodeError) -> Self {
-        Self::OutOfSpec
-    }
-}
-
-impl From<std::io::Error> for Error {
-    fn from(_: std::io::Error) -> Self {
-        Self::OutOfSpec
-    }
-}

src/read/decode/boolean_rle.rs

@@ -1,65 +1,84 @@
-use crate::Error;
+use std::io::Read;
+
+use crate::error::Error;
+
+use super::read_u8;
 
 #[derive(Debug, Copy, Clone, PartialEq)]
-pub enum BooleanRun<'a> {
+#[allow(clippy::large_enum_variant)]
+pub enum BooleanRun {
     Run(u8, u16),
-    Literals(&'a [u8]),
+    Literals([u8; 255]),
 }
 
-pub struct BooleanRleRunIter<'a> {
-    stream: &'a [u8],
+pub struct BooleanRleRunIter<'a, R: Read> {
+    reader: &'a mut R,
 }
 
-impl<'a> BooleanRleRunIter<'a> {
-    pub fn new(stream: &'a [u8]) -> Self {
-        Self { stream }
+impl<'a, R: Read> BooleanRleRunIter<'a, R> {
+    pub fn new(reader: &'a mut R) -> Self {
+        Self { reader }
     }
 }
 
-impl<'a> Iterator for BooleanRleRunIter<'a> {
-    type Item = Result<BooleanRun<'a>, Error>;
+fn read_literals<R: Read>(reader: &mut R, header: i8) -> Result<[u8; 255], Error> {
+    let length = (-header) as usize;
+
+    let mut literals = [0u8; 255];
+
+    reader
+        .take(length as u64)
+        .read_exact(&mut literals[..length])?;
+
+    Ok(literals)
+}
+
+impl<'a, R: Read> Iterator for BooleanRleRunIter<'a, R> {
+    type Item = Result<BooleanRun, Error>;
 
     #[inline]
     fn next(&mut self) -> Option<Self::Item> {
-        let header = *self.stream.first()?;
-        self.stream = &self.stream[1..];
-        let header = i8::from_le_bytes([header]);
+        let header = read_u8(self.reader);
+        let header = match header {
+            Ok(header) => header as i8,
+            Err(e) => return Some(Err(e.into())),
+        };
         if header < 0 {
-            let length = (-header) as usize;
-            if length > self.stream.len() {
-                return Some(Err(Error::RleLiteralTooLarge));
-            }
-            let (literals, remaining) = self.stream.split_at(length);
-            self.stream = remaining;
-            Some(Ok(BooleanRun::Literals(literals)))
+            Some(read_literals(self.reader, header).map(BooleanRun::Literals))
         } else {
             let length = header as u16 + 3;
-            let value = self.stream[0];
-            self.stream = &self.stream[1..];
+            // this is not ok - it may require more than one byte
+            let value = read_u8(self.reader);
+            let value = match value {
+                Ok(value) => value,
+                Err(e) => return Some(Err(e.into())),
+            };
             Some(Ok(BooleanRun::Run(value, length)))
         }
     }
 }
 
-pub struct BooleanIter<'a> {
-    iter: BooleanRleRunIter<'a>,
-    current: Option<BooleanRun<'a>>,
+pub struct BooleanIter<'a, R: Read> {
+    iter: BooleanRleRunIter<'a, R>,
+    current: Option<BooleanRun>,
     position: u8,
+    byte_position: usize,
     remaining: usize,
 }
 
-impl<'a> BooleanIter<'a> {
-    pub fn new(stream: &'a [u8], length: usize) -> Self {
+impl<'a, R: Read> BooleanIter<'a, R> {
+    pub fn new(reader: &'a mut R, length: usize) -> Self {
         Self {
-            iter: BooleanRleRunIter::new(stream),
+            iter: BooleanRleRunIter::new(reader),
             current: None,
             position: 0,
+            byte_position: 0,
             remaining: length,
         }
     }
 }
 
-impl<'a> Iterator for BooleanIter<'a> {
+impl<'a, R: Read> Iterator for BooleanIter<'a, R> {
     type Item = Result<bool, Error>;
 
     #[inline]
@@ -89,7 +108,7 @@ impl<'a> Iterator for BooleanIter<'a> {
                 }
                 BooleanRun::Literals(bytes) => {
                     let mask = 128u8 >> self.position;
-                    let result = bytes[0] & mask == mask;
+                    let result = bytes[self.byte_position] & mask == mask;
                     self.position += 1;
                     if self.remaining == 0 {
                         self.current = None;
@@ -100,8 +119,9 @@ impl<'a> Iterator for BooleanIter<'a> {
                     if self.position == 8 {
                         if bytes.len() == 1 {
                             self.current = None;
+                            self.byte_position = 0;
                         } else {
-                            self.current = Some(BooleanRun::Literals(&bytes[1..]));
+                            self.byte_position += 1;
                         }
                         self.position = 0;
                     }
@@ -128,3 +148,54 @@ impl<'a> Iterator for BooleanIter<'a> {
         (self.remaining, Some(self.remaining))
     }
 }
+
+#[cfg(test)]
+mod test {
+    use super::*;
+
+    #[test]
+    fn basic() {
+        let data = [0x61u8, 0x00];
+
+        let data = &mut data.as_ref();
+
+        let iter = BooleanIter::new(data, 100)
+            .collect::<Result<Vec<_>, Error>>()
+            .unwrap();
+        assert_eq!(iter, vec![false; 100])
+    }
+
+    #[test]
+    fn literals() {
+        let data = [0xfeu8, 0b01000100, 0b01000101];
+
+        let data = &mut data.as_ref();
+
+        let iter = BooleanIter::new(data, 16)
+            .collect::<Result<Vec<_>, Error>>()
+            .unwrap();
+        assert_eq!(
+            iter,
+            vec![
+                false, true, false, false, false, true, false, false, // 0b01000100
+                false, true, false, false, false, true, false, true, // 0b01000101
+            ]
+        )
+    }
+
+    #[test]
+    fn another() {
+        // "For example, the byte sequence [0xff, 0x80] would be one true followed by seven false values."
+        let data = [0xff, 0x80];
+
+        let data = &mut data.as_ref();
+
+        let iter = BooleanIter::new(data, 8)
+            .collect::<Result<Vec<_>, Error>>()
+            .unwrap();
+        assert_eq!(
+            iter,
+            vec![true, false, false, false, false, false, false, false,]
+        )
+    }
+}

src/read/decode/float.rs

@@ -0,0 +1,83 @@
+use crate::error::Error;
+
+/// Sealead trait to generically represent f32 and f64.
+pub trait Float: Default + Copy + private::Sealed {
+    type Bytes: AsRef<[u8]> + AsMut<[u8]> + Default;
+    fn from_le_bytes(bytes: Self::Bytes) -> Self;
+}
+
+mod private {
+    pub trait Sealed {} // Users in other crates cannot name this trait.
+    impl Sealed for f32 {}
+    impl Sealed for f64 {}
+}
+
+impl Float for f32 {
+    type Bytes = [u8; 4];
+
+    #[inline]
+    fn from_le_bytes(bytes: Self::Bytes) -> Self {
+        Self::from_le_bytes(bytes)
+    }
+}
+
+impl Float for f64 {
+    type Bytes = [u8; 8];
+
+    #[inline]
+    fn from_le_bytes(bytes: Self::Bytes) -> Self {
+        Self::from_le_bytes(bytes)
+    }
+}
+
+/// An iterator
+pub struct FloatIter<'a, T: Float, R: std::io::Read> {
+    reader: &'a mut R,
+    remaining: usize,
+    phantom: std::marker::PhantomData<T>,
+}
+
+impl<'a, T: Float, R: std::io::Read> FloatIter<'a, T, R> {
+    #[inline]
+    pub fn new(reader: &'a mut R, length: usize) -> Self {
+        Self {
+            reader,
+            remaining: length,
+            phantom: Default::default(),
+        }
+    }
+
+    #[inline]
+    pub fn len(&self) -> usize {
+        self.remaining
+    }
+
+    #[must_use]
+    pub fn is_empty(&self) -> bool {
+        self.len() == 0
+    }
+}
+
+impl<'a, T: Float, R: std::io::Read> Iterator for FloatIter<'a, T, R> {
+    type Item = Result<T, Error>;
+
+    #[inline]
+    fn next(&mut self) -> Option<Self::Item> {
+        if self.remaining == 0 {
+            return None;
+        }
+        let mut chunk: T::Bytes = Default::default();
+        let error = self.reader.read_exact(chunk.as_mut());
+        if error.is_err() {
+            return Some(Err(Error::DecodeFloat));
+        };
+        self.remaining -= 1;
+        Some(Ok(T::from_le_bytes(chunk)))
+    }
+
+    #[inline]
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        let remaining = self.len();
+        (remaining, Some(remaining))
+    }
+}

src/read/decode/mod.rs

@@ -1,18 +1,17 @@
 mod boolean_rle;
+mod float;
 mod rle_v2;
+mod variable_length;
 
 pub use boolean_rle::{BooleanIter, BooleanRleRunIter, BooleanRun};
+pub use float::FloatIter;
 pub use rle_v2::IteratorEnum;
 pub use rle_v2::{SignedRleV2Iter, SignedRleV2Run, UnsignedRleV2Iter, UnsignedRleV2Run};
+pub use variable_length::Values;
 
-pub fn deserialize_f32(stream: &[u8]) -> impl Iterator<Item = f32> + '_ {
-    stream
-        .chunks_exact(4)
-        .map(|chunk| f32::from_le_bytes(chunk.try_into().unwrap()))
-}
-
-pub fn deserialize_f64(stream: &[u8]) -> impl Iterator<Item = f64> + '_ {
-    stream
-        .chunks_exact(8)
-        .map(|chunk| f64::from_le_bytes(chunk.try_into().unwrap()))
+#[inline]
+fn read_u8<R: std::io::Read>(reader: &mut R) -> Result<u8, std::io::Error> {
+    let mut buf = [0; 1];
+    reader.read_exact(&mut buf)?;
+    Ok(buf[0])
 }