qoi-lib/src/encoder.rs

use std::mem::discriminant;

use itertools::Itertools;

use crate::common::*;


impl Ops {
    // @TODO don't alloc here, get a mut reference and extend here
    fn format (self) -> 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 => vec![],
            Ops::Rgb([r, g, b]) => vec![254, r, g, b],
            Ops::Rgba([r, g, b, a]) => vec![255, r, g, b, a],
            Ops::Index(idx) => vec![0b1111 & idx as u8],
            Ops::Run(len) => vec![(3 << 6) | (len - 1)],
            Ops::Diff([dr, dg, db]) => vec![(1 << 6) | ((cd(dr) << 4) | (cd(dg) << 2) | cd(db)) ],
            Ops::Luma([dr, dg, db]) => vec![(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![];
    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(1);
                    out.extend_from_slice(&Ops::Run(62).format());
                },
                // 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)) {
                out.extend_from_slice(&cur.format());
            }

            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], [last[0], last[1], last[2]]) {
                cur = ops;
            } else {
                cur = header.channels.ops_color(&rgba[..header.channels.num()])?;
            }
            out.extend_from_slice(&cur.format());
            last = rgba;
            known[hash] = rgba;
        }
    }

    Some(out)
}


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);
    }
}