[refactor] consolidate RLE compression

src/compression.rs

@@ -0,0 +1,86 @@
+use crate::sections::PsdCursor;
+
+pub(crate) struct RLECompressed<'a> {
+    cursor: PsdCursor<'a>,
+    repeat: usize,       // 0 signifies end of current section
+    literal: Option<u8>, // None - next repeat bytes are literal, Some(byte) - repeat this byte
+}
+
+impl<'a> RLECompressed<'a> {
+    pub(crate) fn new(bytes: &'a [u8]) -> RLECompressed<'a> {
+        RLECompressed {
+            cursor: PsdCursor::new(bytes),
+            literal: None,
+            repeat: 0,
+        }
+    }
+}
+
+/// https://en.wikipedia.org/wiki/PackBits - algorithm used for decompression
+impl<'a> Iterator for RLECompressed<'a> {
+    type Item = u8;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        if self.repeat > 0 {
+            self.repeat -= 1;
+            return match self.literal {
+                Some(_) => self.literal,
+                None => Some(self.cursor.read_u8()),
+            };
+        }
+
+        if self.cursor.position() >= self.cursor.get_ref().len() as u64 {
+            return None;
+        }
+
+        if self.repeat == 0 {
+            let header = self.cursor.read_i8() as i16;
+            if header == -128 {
+                return self.next();
+            }
+            // TODO: do we need this?
+            // if self.cursor.position() == self.cursor.get_ref().len() as u64 {
+            //     return None;
+            // }
+
+            if header >= 0 {
+                self.literal = None;
+                self.repeat = 1 + header as usize
+            } else {
+                self.literal = Some(self.cursor.read_u8());
+                self.repeat = (1 - header) as usize
+            }
+        }
+
+        self.next()
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_empty() {
+        let empty = vec![];
+        assert_eq!(RLECompressed::new(&empty).collect::<Vec<u8>>(), empty);
+    }
+
+    #[test]
+    fn test_literal() {
+        let value = vec![0, 1, 0, 2, 0, 3, 0, 4];
+        assert_eq!(
+            RLECompressed::new(&value).collect::<Vec<u8>>(),
+            vec![1, 2, 3, 4]
+        );
+    }
+
+    #[test]
+    fn test_repeat() {
+        let value = vec![253, 1];
+        assert_eq!(
+            RLECompressed::new(&value).collect::<Vec<u8>>(),
+            vec![1, 1, 1, 1]
+        );
+    }
+}

src/lib.rs

@@ -33,6 +33,7 @@ use crate::sections::MajorSections;
 use self::sections::file_header_section::FileHeaderSection;
 
 mod blend;
+mod compression;
 mod psd_channel;
 mod render;
 mod sections;

src/psd_channel.rs

@@ -1,3 +1,4 @@
+use crate::compression::RLECompressed;
 use crate::sections::image_data_section::ChannelBytes;
 use crate::sections::PsdCursor;
 use thiserror::Error;
@@ -137,72 +138,27 @@ pub trait IntoRgba {
     ///
     /// We use the channels offset to know where to put it.. So red would go in 0, 4, 8..
     /// blue would go in 1, 5, 9.. etc
-    ///
-    /// https://en.wikipedia.org/wiki/PackBits - algorithm used for decompression
     fn insert_rle_channel(
         &self,
         rgba: &mut Vec<u8>,
         channel_kind: PsdChannelKind,
         channel_bytes: &[u8],
     ) {
-        let mut cursor = PsdCursor::new(&channel_bytes[..]);
-
-        let mut idx = 0;
+        let compressed = RLECompressed::new(channel_bytes);
         let offset = channel_kind.rgba_offset().unwrap();
 
-        while cursor.position() != cursor.get_ref().len() as u64 {
-            let header = cursor.read_i8() as i16;
-
-            if header == -128 {
-                continue;
-            } else if header >= 0 {
-                let bytes_to_read = 1 + header;
-                for byte in cursor.read(bytes_to_read as u32) {
-                    let rgba_idx = self.rgba_idx(idx);
-                    rgba[rgba_idx * 4 + offset] = *byte;
-
-                    idx += 1;
-                }
-            } else {
-                let repeat = 1 - header;
-                let byte = cursor.read_1()[0];
-                for _ in 0..repeat as usize {
-                    let rgba_idx = self.rgba_idx(idx);
-                    rgba[rgba_idx * 4 + offset] = byte;
-
-                    idx += 1;
-                }
-            };
+        for (idx, byte) in compressed.enumerate() {
+            let rgba_idx = self.rgba_idx(idx);
+            let target = rgba_idx * 4 + offset;
+            rgba[target] = byte;
         }
     }
 }
 
 /// Rle decompress a channel
 fn rle_decompress(bytes: &[u8]) -> Vec<u8> {
-    let mut cursor = PsdCursor::new(&bytes[..]);
-
-    let mut decompressed = vec![];
-
-    while cursor.position() != cursor.get_ref().len() as u64 {
-        let header = cursor.read_i8() as i16;
-
-        if header == -128 {
-            continue;
-        } else if header >= 0 {
-            let bytes_to_read = 1 + header;
-            for byte in cursor.read(bytes_to_read as u32) {
-                decompressed.push(*byte);
-            }
-        } else {
-            let repeat = 1 - header;
-            let byte = cursor.read_1()[0];
-            for _ in 0..repeat as usize {
-                decompressed.push(byte);
-            }
-        };
-    }
-
-    decompressed
+    let compressed = RLECompressed::new(bytes);
+    compressed.collect()
 }
 
 /// Take two 8 bit channels that together represent a 16 bit channel and convert them down