qoi_decode.rs

#![feature(let_else)]
#![no_std]

use core::{fmt, mem};

use bytemuck::{Pod, Zeroable};

#[derive(Debug, Copy, Clone, Pod, Zeroable)]
#[repr(C)]
pub struct Header {
    pub magic: [u8; 4],
    pub width: U32,
    pub height: U32,
    pub channels: u8,
    pub color_space: u8,
}

impl Header {
    pub fn read(data: &[u8]) -> Option<(&Self, &[u8])> {
        let (header, body) = try_split::<Header>(data)?;
        if header.magic != *b"qoif" {
            return None;
        }
        Some((header, body))
    }
}

pub trait Color: Clone {
    fn from_rgba(rgba: [u8; 4]) -> Self;
}

impl Color for [u8; 3] {
    fn from_rgba([r, g, b, _]: [u8; 4]) -> Self {
        [r, g, b]
    }
}

impl Color for [u8; 4] {
    fn from_rgba(rgba: [u8; 4]) -> Self {
        rgba
    }
}

pub fn decode_into(mut body: &[u8], mut dst_image: &mut [impl Color]) {
    // The color being a u32 gives a slight performance boost.
    let mut color = u32::from_ne_bytes([0, 0, 0, 255]);
    let mut colors = [0; 64];
    while let &[tag, ref rem @ ..] = body {
        body = rem;
        match tag {
            0b00000000..=0b00111111 => {
                // Index
                color = colors[tag as usize];

                // While we did decode a color, we don't need to hash and
                // re-encode it into the table again as that's where we got it
                // from. So we just write the pixel and continue onto the next
                // loop.
                let [pixel, rem @ ..] = dst_image else { return };
                dst_image = rem;
                *pixel = Color::from_rgba(color.to_ne_bytes());

                continue;
            }
            0b01000000..=0b01111111 => {
                // Diff
                let [diff_r, diff_g, diff_b] = [(tag >> 4) & 0b11, (tag >> 2) & 0b11, tag & 0b11];
                let [r, g, b, a] = color.to_ne_bytes();
                color = u32::from_ne_bytes([
                    r.wrapping_sub(2).wrapping_add(diff_r),
                    g.wrapping_sub(2).wrapping_add(diff_g),
                    b.wrapping_sub(2).wrapping_add(diff_b),
                    a,
                ]);
            }
            0b10000000..=0b10111111 => {
                // Luma
                let &[red_blue_diff, ref rem @ ..] = body else { return };
                body = rem;

                let [diff_g, diff_r, diff_b] =
                    [tag & 0b11_1111, red_blue_diff >> 4, red_blue_diff & 0xF];

                let red_blue_bias = diff_g.wrapping_sub(40);

                let [r, g, b, a] = color.to_ne_bytes();
                color = u32::from_ne_bytes([
                    r.wrapping_add(diff_r).wrapping_add(red_blue_bias),
                    g.wrapping_sub(32).wrapping_add(diff_g),
                    b.wrapping_add(diff_b).wrapping_add(red_blue_bias),
                    a,
                ]);
            }
            0b11000000..=0b11111101 => {
                // Run
                let count = ((tag & 63) + 1) as usize;
                if count > dst_image.len() {
                    return;
                }
                let (fill, rem) = dst_image.split_at_mut(count);
                dst_image = rem;
                fill.fill(Color::from_rgba(color.to_ne_bytes()));
                continue;
            }
            0b11111110 => {
                // RGB
                let &[new_r, new_g, new_b, ref rem @ ..] = body else { return };
                body = rem;
                let [_, _, _, a] = color.to_ne_bytes();
                color = u32::from_ne_bytes([new_r, new_g, new_b, a]);
            }
            0b11111111 => {
                // RGBA
                let &[new_r, new_g, new_b, new_a, ref rem @ ..] = body else { return };
                body = rem;
                color = u32::from_ne_bytes([new_r, new_g, new_b, new_a]);
            }
        }

        let [pixel, rem @ ..] = dst_image else { return };
        dst_image = rem;
        *pixel = Color::from_rgba(color.to_ne_bytes());

        let [r, g, b, a] = color.to_ne_bytes();
        let index = r
            .wrapping_mul(3)
            .wrapping_add(g.wrapping_mul(5))
            .wrapping_add(b.wrapping_mul(7))
            .wrapping_add(a.wrapping_mul(11))
            & 63;

        colors[index as usize] = color;
    }
}

fn try_split<T: Pod>(data: &[u8]) -> Option<(&T, &[u8])> {
    if data.len() >= mem::size_of::<T>() {
        let (data, rem) = data.split_at(mem::size_of::<T>());
        Some((bytemuck::from_bytes(data), rem))
    } else {
        None
    }
}

#[derive(Copy, Clone, Pod, Zeroable)]
#[repr(transparent)]
pub struct U32([u8; 4]);

impl U32 {
    pub fn get(&self) -> u32 {
        u32::from_be_bytes(self.0)
    }
}

impl fmt::Debug for U32 {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        fmt::Debug::fmt(&self.get(), f)
    }
}