QOI - Quote OK Image Format (Rust Port)

qoi.rs

const INDEX: u8 = 0x0;
const RUN_8: u8 = 0x40;
const RUN_16: u8 = 0x60;
const DIFF_8: u8 = 0x80;
const DIFF_16: u8 = 0xc0;
const DIFF_24: u8 = 0xe0;
const COLOR: u8 = 0xf0;

const MASK_2: u8 = 0xc0;
const MASK_3: u8 = 0xe0;
const MASK_4: u8 = 0xf0;

const MAGIC: [u8; 4] = [b'q', b'o', b'i', b'f'];

#[repr(C)]
#[derive(Copy, Clone, PartialEq, Eq)]
struct Rgba {
    r: u8,
    g: u8,
    b: u8,
    a: u8,
}

impl Rgba {
    #[inline]
    const fn new(r: u8, g: u8, b: u8, a: u8) -> Self {
        Self { r, g, b, a }
    }

    fn hash(self) -> u8 {
        self.r ^ self.g ^ self.b ^ self.a
    }
}

#[derive(Copy, Clone, Eq, PartialEq, Debug)]
pub enum Channels {
    Rgb = 3,
    Rgba = 4,
}

pub fn encode(pixels: &[u8], w: usize, h: usize, channels: Channels) -> Option<Vec<u8>> {
    if w == 0 || h == 0 {
        return None;
    }

    let num_channels = channels as usize;
    let mut bytes = Vec::with_capacity(w * h * (num_channels + 1) + 16 + 4);

    bytes.extend_from_slice(&MAGIC);
    bytes.extend_from_slice(&(w as u16).to_le_bytes());
    bytes.extend_from_slice(&(h as u16).to_le_bytes());

    let size_ind = bytes.len();
    bytes.extend_from_slice(&(0 as i32).to_le_bytes()); // will be set at the end
    let size_end = bytes.len();

    let mut index: [Rgba; 64] = unsafe { std::mem::zeroed() };

    let mut run: u16 = 0;
    let mut px_prev = Rgba::new(0, 0, 0, 255);
    let mut px = px_prev;

    let px_len = w * h * num_channels;
    let px_end = px_len - num_channels;

    let mut px_pos = 0;
    while px_pos < px_len {
        unsafe {
            match channels {
                Channels::Rgb => {
                    px.r = *pixels.get_unchecked(px_pos);
                    px.g = *pixels.get_unchecked(px_pos + 1);
                    px.b = *pixels.get_unchecked(px_pos + 2);
                }
                Channels::Rgba => {
                    px.r = *pixels.get_unchecked(px_pos);
                    px.g = *pixels.get_unchecked(px_pos + 1);
                    px.b = *pixels.get_unchecked(px_pos + 2);
                    px.a = *pixels.get_unchecked(px_pos + 3);
                }
            }
        }

        if px == px_prev {
            run += 1;
        }

        if run > 0 && (run == 0x2020 || px != px_prev || px_pos == px_end) {
            if run < 33 {
                run -= 1;
                bytes.push(RUN_8 | (run as u8));
            } else {
                run -= 33;
                bytes.push(RUN_16 | ((run >> 8) as u8));
                bytes.push(run as u8);
            }
            run = 0;
        }

        if px != px_prev {
            let index_u8 = px.hash() % 64;
            let index_pos = index_u8 as usize;
            if index[index_pos] == px {
                bytes.push(INDEX | index_u8);
            } else {
                index[index_pos] = px;

                let vr = px.r as i32 - px_prev.r as i32;
                let vg = px.g as i32 - px_prev.g as i32;
                let vb = px.b as i32 - px_prev.b as i32;
                let va = px.a as i32 - px_prev.a as i32;

                if vr > -16
                    && vr < 17
                    && vg > -16
                    && vg < 17
                    && vb > -16
                    && vb < 17
                    && va > -16
                    && va < 17
                {
                    if va == 0 && vr > -2 && vr < 3 && vg > -2 && vg < 3 && vb > -2 && vb < 3 {
                        bytes.push(DIFF_8 | ((((vr + 1) << 4) | (vg + 1) << 2 | (vb + 1)) as u8));
                    } else if va == 0
                        && vr > -16
                        && vr < 17
                        && vg > -8
                        && vg < 9
                        && vb > -8
                        && vb < 9
                    {
                        bytes.push(DIFF_16 | ((vr + 15) as u8));
                        bytes.push((((vg + 7) << 4) | (vb + 7)) as u8);
                    } else {
                        bytes.push(DIFF_24 | (((vr + 15) >> 1) as u8));
                        bytes.push((((vr + 15) << 7) | ((vg + 15) << 2) | ((vb + 15) >> 3)) as u8);
                        bytes.push((((vb + 15) << 5) | (va + 15)) as u8);
                    }
                } else {
                    let tag = bytes.len();
                    bytes.push(COLOR);
                    if vr != 0 {
                        bytes.push(px.r);
                        bytes[tag] |= 8;
                    }
                    if vg != 0 {
                        bytes.push(px.g);
                        bytes[tag] |= 4;
                    }
                    if vb != 0 {
                        bytes.push(px.b);
                        bytes[tag] |= 2;
                    }
                    if va != 0 {
                        bytes.push(px.a);
                        bytes[tag] |= 1;
                    }
                }
            }
        }

        px_prev = px;
        px_pos += num_channels;
    }

    // Mark the end of the data block with 4 empty bytes
    bytes.extend_from_slice(&[0, 0, 0, 0]);

    // Go back and fill the size value with the size of the data block
    let size = (bytes.len() - size_end) as i32;
    bytes[size_ind..][..4].copy_from_slice(&size.to_le_bytes());

    Some(bytes)
}

pub fn decode(data: &[u8], channels: Channels) -> Option<(usize, usize, Vec<u8>)> {
    let full_len = data.len();
    if full_len < 12 {
        return None;
    }

    if &data[..4] != &MAGIC {
        return None;
    }

    let mut p = 4;

    let read_u8 = |p: &mut usize| {
        let n = data[*p];
        *p += 1;
        n
    };

    let read_u16 = |p: &mut usize| {
        let n = u16::from_le_bytes([data[*p], data[*p + 1]]);
        *p += 2;
        n
    };

    let read_i32 = |p: &mut usize| {
        let n = i32::from_le_bytes([data[*p], data[*p + 1], data[*p + 2], data[*p + 3]]);
        *p += 4;
        n
    };

    let w = read_u16(&mut p) as usize;
    let h = read_u16(&mut p) as usize;
    let data_len = read_i32(&mut p) as usize;

    if w == 0 || h == 0 || data_len + 12 != full_len {
        return None;
    }

    let num_channels = channels as usize;
    let px_len = w * h * num_channels;
    let mut pixels = Vec::with_capacity(px_len);

    let mut px = Rgba::new(0, 0, 0, 255);
    let mut index: [Rgba; 64] = unsafe { std::mem::zeroed() };

    let mut run: u16 = 0;
    let mut px_pos = 0;
    while px_pos < px_len && data.len() > 0 {
        if run > 0 {
            run -= 1;
        } else {
            let b1 = read_u8(&mut p);

            if (b1 & MASK_2) == INDEX {
                px = index[(b1 ^ INDEX) as usize];
            } else if (b1 & MASK_3) == RUN_8 {
                run = (b1 & 0x1f) as u16;
            } else if (b1 & MASK_3) == RUN_16 {
                let b2 = read_u8(&mut p);
                run = ((((b1 & 0x1f) as u16) << 8) | (b2 as u16)) + 32;
            } else if (b1 & MASK_2) == DIFF_8 {
                px.r = px.r.wrapping_add(((b1 >> 4) & 0x03).wrapping_sub(1));
                px.g = px.g.wrapping_add(((b1 >> 2) & 0x03).wrapping_sub(1));
                px.b = px.b.wrapping_add((b1 & 0x03).wrapping_sub(1));
            } else if (b1 & MASK_3) == DIFF_16 {
                let b2 = read_u8(&mut p);
                px.r = px.r.wrapping_add((b1 & 0x1f).wrapping_sub(15));
                px.g = px.g.wrapping_add((b2 >> 4).wrapping_sub(7));
                px.b = px.b.wrapping_add((b2 & 0x0f).wrapping_sub(7));
            } else if (b1 & MASK_4) == DIFF_24 {
                let b2 = read_u8(&mut p);
                let b3 = read_u8(&mut p);
                px.r =
                    px.r.wrapping_add((((b1 & 0x0f) << 1) | (b2 >> 7)).wrapping_sub(15));
                px.g = px.g.wrapping_add(((b2 & 0x7c) >> 2).wrapping_sub(15));
                px.b =
                    px.b.wrapping_add((((b2 & 0x03) << 3) | ((b3 & 0xe0) >> 5)).wrapping_sub(15));
                px.a = px.a.wrapping_add((b3 & 0x1f).wrapping_sub(15));
            } else if (b1 & MASK_4) == COLOR {
                if (b1 & 8) != 0 {
                    px.r = read_u8(&mut p);
                }
                if (b1 & 4) != 0 {
                    px.g = read_u8(&mut p);
                }
                if (b1 & 2) != 0 {
                    px.b = read_u8(&mut p);
                }
                if (b1 & 1) != 0 {
                    px.a = read_u8(&mut p);
                }
            }

            index[(px.hash() % 64) as usize] = px;
        }

        match channels {
            Channels::Rgb => pixels.extend_from_slice(&[px.r, px.g, px.b]),
            Channels::Rgba => pixels.extend_from_slice(&[px.r, px.g, px.b, px.a]),
        }

        px_pos += num_channels;
    }

    Some((w, h, pixels))
}