4 files changed, 139 insertions, 39 deletions

diff --git a/src/decode.rs b/src/decode.rs
index 59dae97..a9abd37 100644
--- a/src/decode.rs
+++ b/src/decode.rs
@@ -1,15 +1,22 @@
-use std::io::{self, Cursor, Read};
+use std::io::{Cursor, Read};
 
 use bitvec::prelude::*;
 use byteorder::{LittleEndian, ReadBytesExt};
 
+use crate::error::{OBIError, OBIResult};
 use crate::{FileHeader, Image, ImageInfoHeader};
 
-pub fn decode_image(obi_data: &mut Cursor<Vec<u8>>) -> io::Result<Image> {
+pub fn decode_image(obi_data: &mut Cursor<Vec<u8>>) -> OBIResult<Image> {
     // file header
-    let version = obi_data.read_u16::<LittleEndian>()?;
-    let file_size = obi_data.read_u32::<LittleEndian>()?;
-    let data_offset = obi_data.read_u32::<LittleEndian>()?;
+    let version = obi_data
+        .read_u16::<LittleEndian>()
+        .map_err(|_| OBIError::Decode)?;
+    let file_size = obi_data
+        .read_u32::<LittleEndian>()
+        .map_err(|_| OBIError::Decode)?;
+    let data_offset = obi_data
+        .read_u32::<LittleEndian>()
+        .map_err(|_| OBIError::Decode)?;
     let file_header = FileHeader {
         version,
         file_size,
@@ -17,10 +24,18 @@ pub fn decode_image(obi_data: &mut Cursor<Vec<u8>>) -> io::Result<Image> {
     };
 
     // image info header
-    let width = obi_data.read_u32::<LittleEndian>()?;
-    let height = obi_data.read_u32::<LittleEndian>()?;
-    let compression_type = obi_data.read_u32::<LittleEndian>()?;
-    let post_compression_size = obi_data.read_u32::<LittleEndian>()?;
+    let width = obi_data
+        .read_u32::<LittleEndian>()
+        .map_err(|_| OBIError::Decode)?;
+    let height = obi_data
+        .read_u32::<LittleEndian>()
+        .map_err(|_| OBIError::Decode)?;
+    let compression_type = obi_data
+        .read_u32::<LittleEndian>()
+        .map_err(|_| OBIError::Decode)?;
+    let post_compression_size = obi_data
+        .read_u32::<LittleEndian>()
+        .map_err(|_| OBIError::Decode)?;
     let image_info_header = ImageInfoHeader {
         width,
         height,
@@ -30,7 +45,9 @@ pub fn decode_image(obi_data: &mut Cursor<Vec<u8>>) -> io::Result<Image> {
 
     // pixmap data
     let mut data_bytes = vec![];
-    obi_data.read_to_end(&mut data_bytes)?;
+    obi_data
+        .read_to_end(&mut data_bytes)
+        .map_err(|_| OBIError::Decode)?;
     let data = data_bytes
         .iter()
         .map(|&b| {
diff --git a/src/encode.rs b/src/encode.rs
index e5760a8..a8c58e0 100644
--- a/src/encode.rs
+++ b/src/encode.rs
@@ -1,34 +1,55 @@
-use std::io;
+use std::borrow::Borrow;
 
 use bitvec::{prelude::*, vec::BitVec};
 use byteorder::{LittleEndian, WriteBytesExt};
 
+use crate::error::{OBIError, OBIResult};
 use crate::Image;
 
-pub fn encode_image(obi_image: Image) -> io::Result<Vec<u8>> {
+pub fn encode_image<I>(obi_image: I) -> OBIResult<Vec<u8>>
+where
+    I: Borrow<Image>,
+{
+    let obi_image = obi_image.borrow();
     let mut obi_data = Vec::with_capacity(obi_image.file_header.file_size as usize);
 
     // file header
-    let file_header = obi_image.file_header;
-    obi_data.write_u16::<LittleEndian>(file_header.version)?;
-    obi_data.write_u32::<LittleEndian>(file_header.file_size)?;
-    obi_data.write_u32::<LittleEndian>(file_header.data_offset)?;
+    let file_header = &obi_image.file_header;
+    obi_data
+        .write_u16::<LittleEndian>(file_header.version)
+        .map_err(|_| OBIError::Encode)?;
+    obi_data
+        .write_u32::<LittleEndian>(file_header.file_size)
+        .map_err(|_| OBIError::Encode)?;
+    obi_data
+        .write_u32::<LittleEndian>(file_header.data_offset)
+        .map_err(|_| OBIError::Encode)?;
 
     // image info header
-    let image_info_header = obi_image.image_info_header;
-    obi_data.write_u32::<LittleEndian>(image_info_header.width)?;
-    obi_data.write_u32::<LittleEndian>(image_info_header.height)?;
-    obi_data.write_u32::<LittleEndian>(image_info_header.compression_type)?;
-    obi_data.write_u32::<LittleEndian>(image_info_header.post_compression_size)?;
+    let image_info_header = &obi_image.image_info_header;
+    obi_data
+        .write_u32::<LittleEndian>(image_info_header.width)
+        .map_err(|_| OBIError::Encode)?;
+    obi_data
+        .write_u32::<LittleEndian>(image_info_header.height)
+        .map_err(|_| OBIError::Encode)?;
+    obi_data
+        .write_u32::<LittleEndian>(image_info_header.compression_type)
+        .map_err(|_| OBIError::Encode)?;
+    obi_data
+        .write_u32::<LittleEndian>(image_info_header.post_compression_size)
+        .map_err(|_| OBIError::Encode)?;
 
     // pixmap data
-    let pixmap = obi_image.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);
         }
-        obi_data.write_u8(bv.load::<u8>())?;
+        obi_data
+            .write_u8(bv.load::<u8>())
+            .map_err(|_| OBIError::Encode)?;
     }
 
     return Ok(obi_data);
diff --git a/src/error.rs b/src/error.rs
new file mode 100644
index 0000000..8c2f096
--- /dev/null
+++ b/src/error.rs
@@ -0,0 +1,45 @@
+use std::{error, fmt};
+
+pub type OBIResult<T> = Result<T, OBIError>;
+
+#[derive(Debug)]
+pub enum OBIError {
+    Encode,
+    Decode,
+    Image,
+}
+
+impl fmt::Display for OBIError {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        match self {
+            OBIError::Encode => writeln!(f, "Encoding error"),
+            OBIError::Decode => writeln!(f, "Decoding error"),
+            OBIError::Image => writeln!(f, "Image manipulation error"),
+        }
+    }
+}
+
+impl error::Error for OBIError {}
+
+#[derive(Debug)]
+enum Encode {
+    Metadata(MetadataError),
+    ChunkError(u32),
+}
+
+#[derive(Debug)]
+enum Decode {
+    Metadata(MetadataError),
+    ChunkError(u32),
+}
+
+#[derive(Debug)]
+enum MetadataError {
+    VersionError,
+    FileSizeError,
+    DataOffsetError,
+    WidthError,
+    HeightError,
+    CompressionTypeError,
+    PostCompressionSizeError,
+}
diff --git a/src/lib.rs b/src/lib.rs
index cd7923c..7429101 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -1,8 +1,13 @@
 #![allow(unreachable_patterns)]
 #![allow(non_snake_case)]
 
-pub mod decode;
-pub mod encode;
+use std::io;
+
+mod decode;
+mod encode;
+pub mod error;
+
+use crate::error::{OBIError, OBIResult};
 
 #[non_exhaustive]
 #[derive(Copy, Clone, Debug, PartialEq)]
@@ -59,6 +64,7 @@ impl ImageInfoHeader {
 }
 
 #[non_exhaustive]
+#[derive(Copy, Clone, Debug, PartialEq)]
 pub enum CompressionType {
     RLE,
     Kosinki,
@@ -96,23 +102,34 @@ impl Image {
             data,
         }
     }
-}
 
-#[cfg(test)]
-mod tests {
-    use super::*;
-    use std::mem::size_of;
+    pub fn width(&self) -> u32 {
+        self.image_info_header.width
+    }
+
+    pub fn height(&self) -> u32 {
+        self.image_info_header.height
+    }
+
+    fn to_index(&self, x: u32, y: u32) -> usize {
+        (y * self.width() + x) as usize
+    }
+
+    pub fn set_pixel(&mut self, x: u32, y: u32, val: bool) -> OBIResult<()> {
+        if x >= self.width() || y > self.height() {
+            Err(OBIError::Image)
+        } else {
+            let index = self.to_index(x, y);
+            self.data[index] = val;
+            Ok(())
+        }
+    }
 
-    #[test]
-    fn size_of_image_info_header() {
-        let file_header_size = size_of::<ImageInfoHeader>();
-        assert_eq!(16, file_header_size);
+    pub fn encode(&self) -> OBIResult<Vec<u8>> {
+        encode::encode_image(self)
     }
 
-    #[test]
-    fn encode_decode() {
-        let img = Image::new(100, 80);
-        let encoded = encode::encode_image(img).unwrap();
-        assert_eq!(encoded.len(), 1026);
+    pub fn decode(data: &mut io::Cursor<Vec<u8>>) -> OBIResult<Image> {
+        decode::decode_image(data)
     }
 }