qoi-lib/src/decoder.rs

use imageproc::image::{ RgbImage, RgbaImage };
use itertools::Itertools;
use crate::common::*;

impl Header {
    pub fn read (header: &[u8]) -> Option<Self> {
        let magic = header[..4].iter().eq(MAGIC.iter());
        if !magic { return None; }

        let width  = u32::from_be_bytes( header[4..8].try_into().ok()?);
        let height = u32::from_be_bytes(header[8..12].try_into().ok()?);

        let channels = Self::channels(u8::from_be_bytes(header[12..13].try_into().ok()?))?;
        let colorspace = Self::colorspace(u8::from_be_bytes(header[13..14].try_into().ok()?))?;

        Some(Self { width, height, channels, colorspace })
    }

    fn colorspace (space: u8) -> Option<Colorspace> {
        match space {
            0 => Some(Colorspace::Srgb),
            1 => Some(Colorspace::Linear),
            _ => None
        }
    }

    fn channels (ch: u8) -> Option<Channels> {
        match ch {
            3 => Some(Channels::RGB),
            4 => Some(Channels::RGBA),
            _ => None
        }
    }
}

pub trait Decoder {
    fn decode (header: Header, bytes: &[u8]) -> Option<Self> where Self: Sized;
}

impl Decoder for Vec<[u8; 4]> {
    /// decode a complete slice into array of pixel color values
    fn decode (header: Header, bytes: &[u8]) -> Option<Self> {
        if bytes.len() < 8 { return None; }
        let end = [0, 0, 0, 0, 0, 0, 0, 1];
        if bytes[(bytes.len() - 8)..] != end { return None; }
        let bytes = &bytes[..bytes.len() - 8];

        let mut known = [[0; 4]; 64];
        let mut prev = [0, 0, 0, 255];
        let mut out = vec![];
        let mut iter = bytes.iter();

        loop {
            let Some(&byte) = iter.next() else { break; };
            let (color, len) = match byte {
                254 => { // QOI_OP_RGB
                    let (r, g, b) = iter.next_tuple().unwrap();
                    ([*r, *g, *b, prev[3]], 1)
                },
                255 => { // QOI_OP_RGBA
                    let (r, g, b, a) = iter.next_tuple().unwrap();
                    ([*r, *g, *b, *a], 1)
                },
                q if q >> 6 == 0 => { // QOI_OP_INDEX
                    (known[q as usize & 63], 1)
                },
                q if q >> 6 == 1 => { // QOI_OP_DIFF
                    let dr = ((q & 48) >> 4) as i8 - 2;
                    let dg = ((q & 12) >> 2) as i8 - 2;
                    let db = (q & 3) as i8 - 2;
                    ([ prev[0].wrapping_add_signed(dr), prev[1].wrapping_add_signed(dg), prev[2].wrapping_add_signed(db), prev[3]], 1)
                },
                q if q >> 6 == 2 => { // QOI_OP_LUMA
                    let extra = *iter.next()?;
                    let dg = (q & 63) as i8 - 32;
                    let dr = ((extra & 240) >> 4) as i8 - 8;
                    let db = (extra & 15) as i8 - 8;
                    ([ prev[0].wrapping_add_signed(dr + dg), prev[1].wrapping_add_signed(dg), prev[2].wrapping_add_signed(db + dg), prev[3]], 1)
                },
                q if q >> 6 == 3 => { // QOI_OP_RUN
                    let len = ((q & 63) + 1) as usize;
                    (prev.clone(), len)
                },
                _ => return None
            };

            known[hash(color)] = color;
            prev = color.clone();
            (0..len).for_each(|_| out.push(color));
        }

        if header.width as u64 * header.height as u64 != out.len() as u64 { return None; }

        Some(out)
    }
}

impl Decoder for Vec<[u8; 3]> {
    fn decode (header: Header, bytes: &[u8]) -> Option<Self> where Self: Sized {
        let pixels: Vec<[u8; 4]> = Decoder::decode(header, bytes)?;
        let pixels: Vec<[u8; 3]> = pixels.iter().filter_map(|e| (&e[0..3]).try_into().ok()).collect();
        if pixels.len() as u64 != header.width as u64 * header.height as u64 { return None; }
        Some(pixels)
    }
}

impl Decoder for Vec<u8> {
    fn decode (header: Header, bytes: &[u8]) -> Option<Self> where Self: Sized {
        let pixels: Vec<[u8; 4]> = Decoder::decode(header, bytes)?;
        if header.channels == Channels::RGBA {
            bytemuck::try_cast_vec(pixels).ok()
        } else {
            let pixels: Vec<[u8; 3]> = pixels.iter().filter_map(|e| (&e[0..3]).try_into().ok()).collect();
            if pixels.len() as u64 != header.width as u64 * header.height as u64 { return None; }
            bytemuck::try_cast_vec(pixels).ok()
        }
    }
}

impl Decoder for RgbaImage {
    fn decode (mut header: Header, bytes: &[u8]) -> Option<Self> where Self: Sized {
        header.channels = Channels::RGBA;
        let pixels = Decoder::decode(header, bytes)?;
        RgbaImage::from_raw(header.width, header.height, pixels)
    }
}

impl Decoder for RgbImage {
    fn decode (mut header: Header, bytes: &[u8]) -> Option<Self> where Self: Sized {
        header.channels = Channels::RGB;
        let pixels = Decoder::decode(header, bytes)?;
        RgbImage::from_raw(header.width, header.height, pixels)
    }
}

/// a generic decoder over a type which implements Decoder trait
pub fn decode <T> (bytes: &[u8]) -> Option<T> where T: Decoder {
    let header = Header::read(&bytes[..14])?;
    let rest = &bytes[14..];
    Decoder::decode(header, rest)
}


#[cfg(test)]
mod test {
    use std::time::Instant;

    use imageproc::image::RgbaImage;
    use super::*;

    #[test]
    fn basic () {
        let file = std::fs::read("in/dice.qoi").unwrap();
        let now = Instant::now();

        let header = Header::read(&file[..14]).unwrap();
        assert_eq!(Header { width: 800, height: 600, channels: Channels::RGBA, colorspace: Colorspace::Srgb }, header);

        let rest = &file[14..];
        let pixels: Vec<[u8; 4]> = Decoder::decode(header, rest).unwrap();
        println!("dice decoding done in {}ms", now.elapsed().as_millis());

        let now = Instant::now();

        let pixels = bytemuck::try_cast_vec(pixels).unwrap();
        println!("dice cast done in {}ns", now.elapsed().as_nanos());
        let rgba = RgbaImage::from_raw(header.width, header.height, pixels).unwrap();
        rgba.save("out/dice.png").unwrap();
    }

    #[test]
    fn edge () {
        let file = std::fs::read("in/edgecase.qoi").unwrap();

        let now = Instant::now();

        let header = Header::read(&file[..14]).unwrap();
        assert_eq!(Header { width: 256, height: 64, channels: Channels::RGBA, colorspace: Colorspace::Srgb }, header);

        let rest = &file[14..];
        let pixels: Vec<[u8; 4]> = Decoder::decode(header, rest).unwrap();
        println!("edgecase decoding done in {}ms", now.elapsed().as_millis());

        let pixels = bytemuck::try_cast_vec(pixels).unwrap();
        let rgba = RgbaImage::from_raw(header.width, header.height, pixels).unwrap();
        rgba.save("out/edgecase.png").unwrap();
    }


    #[test]
    fn edge_as_jpeg () {
        let file = std::fs::read("in/edgecase.qoi").unwrap();

        let header = Header::read(&file[..14]).unwrap();
        assert_eq!(Header { width: 256, height: 64, channels: Channels::RGBA, colorspace: Colorspace::Srgb }, header);

        let rest = &file[14..];
        let rgb = RgbImage::decode(header, rest).unwrap();

        rgb.save_with_format("out/edgecase.jpg", imageproc::image::ImageFormat::Jpeg).unwrap();
    }
}