qoi-lib/src/encoder.rs

use std::mem::discriminant;

use itertools::Itertools;
use crate::common::*;

// @TODO replace Options with custom Errors
// @TODO encoder wrappers for RgbaImage/RgbImage

impl Ops {
    fn write_encoded (self, buf: &mut Vec<u8>) {
        // cast_diff
        let cd = |diff: i8| cast_with_bias(diff, 2);
        // cast_luma_green
        let clg = |diff: i8| cast_with_bias(diff, 32);
        // cast_luma_red_blue
        let clrb = |diff: i8| cast_with_bias(diff, 8);

        match self {
            Ops::Nop => (),
            Ops::Rgb([r, g, b]) => buf.extend_from_slice(&[254, r, g, b]),
            Ops::Rgba([r, g, b, a]) => buf.extend_from_slice(&[255, r, g, b, a]),
            Ops::Index(idx) => buf.push(idx as u8),
            Ops::Run(len) => buf.push((3 << 6) | (len - 1)),
            Ops::Diff([dr, dg, db]) => buf.push((1 << 6) | ((cd(dr) << 4) | (cd(dg) << 2) | cd(db))),
            Ops::Luma([dr, dg, db]) => buf.extend_from_slice(&[(2 << 6) | clg(dg), (clrb(dr) << 4) | clrb(db)]),
        }
    }
}

fn encode_header (header: Header) -> Option<Vec<u8>> {
    let mut out = vec![];
    out.extend_from_slice(&MAGIC);
    out.extend_from_slice(&header.width.to_be_bytes());
    out.extend_from_slice(&header.height.to_be_bytes());
    out.extend_from_slice(&(header.channels.num() as u8).to_be_bytes());
    out.extend_from_slice(&header.colorspace.to_repr().to_be_bytes());

    if out.len() != 14 { return None; }

    Some(out)
}

fn encode_body (header: Header, data: &[u8]) -> Option<Vec<u8>> {
    let mut known = [[0u8; 4]; 64];
    let mut last = [0u8, 0, 0, 255];
    let mut out = Vec::with_capacity(header.width as usize * header.height as usize);
    let mut cur = Ops::Nop;

    for mut chunk in &data.iter().chunks(header.channels.num()) {
        let (&r, &g, &b) = chunk.next_tuple()?;
        // if we need only 3 channels, we use a placeholder 255 for alpha and ignore it during write
        let a = *chunk.next().unwrap_or(&255);
        let rgba = [r, g, b, a];

        // If current pixel value is the same as of the previous pixel, we need to start/continue a run
        if rgba == last {
            match cur {
                // if run is in bounds 0..62
                Ops::Run(run) if run < 62 => cur = Ops::Run(run + 1),
                // if run is full, commit it to out vec, and start a new one
                Ops::Run(run) if run >= 62 => {
                    cur = Ops::Run(62); // clamp for safety;
                    cur.write_encoded(&mut out);
                    cur = Ops::Run(1);
                },
                // otherwise (default case, when encountering a match for the first time) start a run
                _ => cur = Ops::Run(1),
            }
        } else {
            // if current Op is Run (meaning we need to and it since there's no match)
            if discriminant(&cur) == discriminant(&Ops::Run(0)) {
                cur.write_encoded(&mut out);
            }

            let hash = hash(rgba);
            // INDEX -> DIFF -> LUMA -> RGB(A)
            if known[hash] == rgba {
                cur = Ops::Index(hash);
            } else if let Some(ops) = diff_ops([r, g, b, a], last) {
                cur = ops;
            } else {
                cur = header.channels.ops_color(&rgba[..header.channels.num()])?;
            }
            cur.write_encoded(&mut out);
            last = rgba;
            known[hash] = rgba;
        }
    }

    // if an image buffer ends with a run, flush it as well
    if discriminant(&cur) == discriminant(&Ops::Run(0)) {
        cur.write_encoded(&mut out);
    }

    Some(out)
}


pub fn encode_full (header: Header, bytes: &[u8]) -> Option<Vec<u8>> {
    let mut out = Vec::with_capacity(header.width as usize * header.height as usize * 4);
    out.extend_from_slice(&encode_header(header)?);
    out.extend_from_slice(&encode_body(header, bytes)?);
    out.extend_from_slice(&[0, 0, 0, 0, 0, 0, 0, 1]);
    Some(out)
}



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

    #[test]
    fn basic () {
        let file = imageproc::image::io::Reader::open("in/dice.png").unwrap().decode().unwrap();
        let file = file.into_rgba8().into_raw();
        let header = Header { width: 800, height: 600, channels: Channels::RGBA, colorspace: Colorspace::Srgb };

        let inst = Instant::now();
        let encoded = encode_full(header, &file).unwrap();
        println!("dice encoding done in {} ms", inst.elapsed().as_millis());
        let _write = std::fs::write("out/dice.qoi", &encoded);
    }

    #[test]
    fn _1kb () {
        let file = imageproc::image::io::Reader::open("in/1kb.png").unwrap().decode().unwrap();
        let file = file.into_rgba8().into_raw();
        let header = Header { width: 17, height: 14, channels: Channels::RGBA, colorspace: Colorspace::Srgb };

        let inst = Instant::now();
        let encoded = encode_full(header, &file).unwrap();
        println!("1kb encoding done in {} ms", inst.elapsed().as_millis());

        let _write = std::fs::write("out/1kb.qoi", &encoded);
    }

    #[test]
    fn kodim23 () {
        let file = imageproc::image::io::Reader::open("in/kodim23.png").unwrap().decode().unwrap();
        let file = file.into_rgb8().into_raw();
        let header = Header { width: 768, height: 512, channels: Channels::RGB, colorspace: Colorspace::Srgb };


        let inst = Instant::now();
        let encoded = encode_full(header, &file).unwrap();
        println!("kodim23 encoding done in {} ms", inst.elapsed().as_millis());
        let _write = std::fs::write("out/kodim23.qoi", &encoded);
    }
}