add RLE compression functions

4 files changed, 191 insertions, 27 deletions

diff --git a/src/decode.rs b/src/decode.rs
index a9abd37..9c6a065 100644
--- a/src/decode.rs
+++ b/src/decode.rs
@@ -4,7 +4,7 @@ use bitvec::prelude::*;
 use byteorder::{LittleEndian, ReadBytesExt};
 
 use crate::error::{OBIError, OBIResult};
-use crate::{FileHeader, Image, ImageInfoHeader};
+use crate::{CompressionType, FileHeader, Image, ImageInfoHeader};
 
 pub fn decode_image(obi_data: &mut Cursor<Vec<u8>>) -> OBIResult<Image> {
     // file header
@@ -43,22 +43,59 @@ pub fn decode_image(obi_data: &mut Cursor<Vec<u8>>) -> OBIResult<Image> {
         post_compression_size,
     };
 
-    // pixmap data
-    let mut data_bytes = vec![];
-    obi_data
-        .read_to_end(&mut data_bytes)
-        .map_err(|_| OBIError::Decode)?;
-    let data = data_bytes
-        .iter()
-        .map(|&b| {
-            BitVec::<Lsb0, u8>::from_element(b)
-                .into_iter()
-                .map(|e| e as bool)
-                .collect::<Vec<bool>>()
-        })
-        .flatten()
-        .collect::<Vec<_>>();
+    let data: Vec<bool> = match CompressionType::from_u32(compression_type) {
+        CompressionType::RLE => {
+            let mut rest = vec![];
+            let mut lengths = vec![];
+            loop {
+                let l = obi_data
+                    .read_u32::<LittleEndian>()
+                    .map_err(|_| OBIError::Encode)?;
+                if l == 0 {
+                    break;
+                }
+                lengths.push(l);
+            }
+            obi_data
+                .read_to_end(&mut rest)
+                .map_err(|_| OBIError::Decode)?;
+            let data_points = rest
+                .iter()
+                .map(|&b| {
+                    BitVec::<Lsb0, u8>::from_element(b)
+                        .into_iter()
+                        .map(|e| e as bool)
+                        .collect::<Vec<bool>>()
+                })
+                .flatten()
+                .collect::<Vec<bool>>();
 
+            let data = data_points
+                .into_iter()
+                .zip(lengths)
+                .map(|(d, l)| vec![d; l as usize])
+                .flatten()
+                .collect::<Vec<bool>>();
+            data
+        }
+        _ => {
+            let mut rest = vec![];
+            obi_data
+                .read_to_end(&mut rest)
+                .map_err(|_| OBIError::Decode)?;
+            let data_points = rest
+                .iter()
+                .map(|&b| {
+                    BitVec::<Lsb0, u8>::from_element(b)
+                        .into_iter()
+                        .map(|e| e as bool)
+                        .collect::<Vec<bool>>()
+                })
+                .flatten()
+                .collect::<Vec<_>>();
+            data_points
+        }
+    };
     return Ok(Image {
         file_header,
         image_info_header,
diff --git a/src/encode.rs b/src/encode.rs
index a8c58e0..30edc49 100644
--- a/src/encode.rs
+++ b/src/encode.rs
@@ -3,8 +3,11 @@ use std::borrow::Borrow;
 use bitvec::{prelude::*, vec::BitVec};
 use byteorder::{LittleEndian, WriteBytesExt};
 
-use crate::error::{OBIError, OBIResult};
-use crate::Image;
+use crate::{
+    error::{OBIError, OBIResult},
+    rle::RLE,
+    CompressionType, Image,
+};
 
 pub fn encode_image<I>(obi_image: I) -> OBIResult<Vec<u8>>
 where
@@ -40,17 +43,38 @@ where
         .write_u32::<LittleEndian>(image_info_header.post_compression_size)
         .map_err(|_| OBIError::Encode)?;
 
+    let write_pixel_data = |pixels: &Vec<bool>, obi_data: &mut Vec<u8>| -> OBIResult<()> {
+        for byte in pixels.chunks(8) {
+            let mut bv = BitVec::<Lsb0, u8>::new();
+            for &b in byte {
+                bv.push(b)
+            }
+            obi_data
+                .write_u8(bv.load::<u8>())
+                .map_err(|_| OBIError::Encode)?;
+        }
+        Ok(())
+    };
+
     // pixmap data
     let pixmap = &obi_image.data;
-    for byte in pixmap.chunks(8) {
-        let mut bv = BitVec::<Lsb0, u8>::new();
-        for &b in byte {
-            bv.push(b);
+    match CompressionType::from_u32(obi_image.image_info_header.compression_type) {
+        CompressionType::RLE => {
+            let (data_points, lengths): (Vec<_>, Vec<_>) = pixmap.compress().into_iter().unzip();
+            for l in lengths {
+                obi_data
+                    .write_u32::<LittleEndian>(l as u32)
+                    .map_err(|_| OBIError::Encode)?;
+            }
+            // end length sequence with zero
+            obi_data
+                .write_u32::<LittleEndian>(0)
+                .map_err(|_| OBIError::Encode)?;
+            // begin data point sequence
+            write_pixel_data(&data_points, &mut obi_data)?;
         }
-        obi_data
-            .write_u8(bv.load::<u8>())
-            .map_err(|_| OBIError::Encode)?;
-    }
+        _ => write_pixel_data(pixmap, &mut obi_data)?,
+    };
 
     return Ok(obi_data);
 }
diff --git a/src/lib.rs b/src/lib.rs
index f949d06..3995c16 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -7,8 +7,10 @@ use std::io;
 mod decode;
 mod encode;
 pub mod error;
+mod rle;
 
-use crate::error::{OBIError, OBIResult};
+use error::{OBIError, OBIResult};
+use rle::RLE;
 
 #[non_exhaustive]
 #[derive(Copy, Clone, Debug, PartialEq)]
@@ -83,6 +85,14 @@ impl CompressionType {
             _ => panic!("Invalid compression algorithm"),
         }
     }
+    pub fn to_u32(self) -> u32 {
+        match self {
+            CompressionType::None => 0,
+            CompressionType::RLE => 1,
+            CompressionType::Kosinki => 10,
+            _ => panic!("Invalid compression algorithm"),
+        }
+    }
 }
 
 #[derive(Debug)]
@@ -111,6 +121,10 @@ impl Image {
         }
     }
 
+    pub fn use_compression(&mut self, comp: CompressionType) {
+        self.image_info_header.compression_type = comp.to_u32();
+    }
+
     pub fn width(&self) -> u32 {
         self.image_info_header.width
     }
diff --git a/src/rle.rs b/src/rle.rs
new file mode 100644
index 0000000..f6d1e05
--- /dev/null
+++ b/src/rle.rs
@@ -0,0 +1,89 @@
+pub trait RLE {
+    type Data;
+    fn compress(&self) -> Vec<Run<Self::Data>>;
+    fn decompress(runs: Vec<Run<Self::Data>>) -> Vec<Self::Data>;
+}
+
+pub type Run<T> = (T, usize);
+
+impl<T> RLE for Vec<T>
+where
+    T: Clone + PartialEq,
+{
+    type Data = T;
+    fn compress(&self) -> Vec<Run<Self::Data>> {
+        let mut runs = vec![];
+        if self.is_empty() {
+            return runs;
+        }
+        let mut idx = 0;
+        loop {
+            let first = &self[idx];
+            let run_length = self[idx..].iter().take_while(|&item| item == first).count();
+
+            runs.push((first.clone(), run_length));
+
+            idx += run_length;
+            if idx > self.len() - 1 {
+                break;
+            }
+        }
+        runs
+    }
+    fn decompress(runs: Vec<Run<Self::Data>>) -> Vec<Self::Data> {
+        runs.into_iter()
+            .map(|(item, size)| vec![item; size])
+            .flatten()
+            .collect()
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    #[test]
+    fn singleton() {
+        let data = "a".chars().collect::<Vec<_>>();
+        assert_eq!(Vec::<char>::compress(&data), vec![('a', 1),]);
+    }
+    #[test]
+    fn identity() {
+        let data = "aabbccddaabbaa".chars().collect::<Vec<_>>();
+        assert_eq!(Vec::<char>::decompress(data.compress()), data);
+    }
+    #[test]
+    fn repeated_singleton() {
+        let data = "aaaaaaaaaaaaa".chars().collect::<Vec<_>>();
+        assert_eq!(Vec::<char>::compress(&data), vec![('a', 13),]);
+    }
+    #[test]
+    fn empty_runs() {
+        let data = "".chars().collect::<Vec<_>>();
+        assert!(data.compress().is_empty());
+    }
+    #[test]
+    fn empty_decompress() {
+        assert!(Vec::<char>::decompress(vec![]).is_empty());
+    }
+    #[test]
+    fn check_runs1() {
+        let data = "aaaabbbbcccc".chars().collect::<Vec<_>>();
+        assert_eq!(data.compress(), vec![('a', 4), ('b', 4), ('c', 4)]);
+    }
+    #[test]
+    fn check_runs2() {
+        let data = "aabbccddaabbaa".chars().collect::<Vec<_>>();
+        assert_eq!(
+            data.compress(),
+            vec![
+                ('a', 2),
+                ('b', 2),
+                ('c', 2),
+                ('d', 2),
+                ('a', 2),
+                ('b', 2),
+                ('a', 2)
+            ]
+        );
+    }
+}