From 83e75b1b6534d03211ebe9e6664124894027e10c Mon Sep 17 00:00:00 2001
From: Ivan Smirnov <i.s.smirnov@gmail.com>
Date: Mon, 29 Nov 2021 22:26:20 +0000
Subject: [PATCH] Add colorspace support to qoi_encode_to_vec()

---
 src/encode.rs | 11 +++++++----
 src/lib.rs    |  6 +++---
 2 files changed, 10 insertions(+), 7 deletions(-)

diff --git a/src/encode.rs b/src/encode.rs
index 3d68ebe..1ca5df6 100644
--- a/src/encode.rs
+++ b/src/encode.rs
@@ -1,5 +1,6 @@
 use std::slice;
 
+use crate::colorspace::ColorSpace;
 use crate::consts::{
     QOI_COLOR, QOI_DIFF_16, QOI_DIFF_24, QOI_DIFF_8, QOI_HEADER_SIZE, QOI_INDEX, QOI_PADDING,
     QOI_RUN_16, QOI_RUN_8,
@@ -122,7 +123,7 @@ fn encode_diff_wrapping<const N: usize>(
 }
 
 pub(crate) fn qoi_encode_impl<const N: usize>(
-    data: &[u8], width: u32, height: u32,
+    data: &[u8], width: u32, height: u32, colorspace: ColorSpace,
 ) -> Result<Vec<u8>>
 where
     Pixel<N>: SupportedChannels,
@@ -150,7 +151,7 @@ where
     header.width = width;
     header.height = height;
     header.channels = N as u8;
-    // TODO: colorspace
+    header.colorspace = colorspace.into();
     buf.write(header.to_bytes());
 
     let mut index = [Pixel::new(); 64];
@@ -228,11 +229,13 @@ where
 
 pub fn qoi_encode_to_vec(
     data: impl AsRef<[u8]>, width: u32, height: u32, channels: u8,
+    colorspace: impl Into<ColorSpace>,
 ) -> Result<Vec<u8>> {
     let data = data.as_ref();
+    let colorspace = colorspace.into();
     match channels {
-        3 => qoi_encode_impl::<3>(data, width, height),
-        4 => qoi_encode_impl::<4>(data, width, height),
+        3 => qoi_encode_impl::<3>(data, width, height, colorspace.into()),
+        4 => qoi_encode_impl::<4>(data, width, height, colorspace.into()),
         _ => Err(Error::InvalidChannels { channels }),
     }
 }
diff --git a/src/lib.rs b/src/lib.rs
index 46ab40c..69f22d6 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -33,14 +33,14 @@ mod tests {
     #[test]
     fn kodim_01() {
         let (w, h, c, v) = read_png("assets/kodim01.png");
-        let q = qoi_encode_to_vec(&v, w, h, c).unwrap();
+        let q = qoi_encode_to_vec(&v, w, h, c, 0).unwrap();
         std::fs::write("kodim01.qoi", q.as_slice()).unwrap();
     }
 
     #[test]
     fn wikipedia() {
         let (w, h, c, v) = read_png("assets/en.wikipedia.org.png");
-        let q = qoi_encode_to_vec(&v, w, h, c).unwrap();
+        let q = qoi_encode_to_vec(&v, w, h, c, 0).unwrap();
         std::fs::write("wikipedia.qoi", q.as_slice()).unwrap();
     }
 
@@ -48,7 +48,7 @@ mod tests {
     fn roundtrip_3() {
         let three_raw = include_bytes!("../assets/three.raw").to_vec();
         let (w, h, c) = (572, 354, 3);
-        let three_qoi = qoi_encode_to_vec(&three_raw, w, h, c).unwrap();
+        let three_qoi = qoi_encode_to_vec(&three_raw, w, h, c, 0).unwrap();
         let (header, three_rtp) = qoi_decode_to_vec(&three_qoi, c).unwrap();
         assert_eq!(header.magic, QOI_MAGIC);
         assert_eq!(header.width, w);