gistfile1.txt

use std::{
    mem::transmute,
    ptr::{copy_nonoverlapping, write_bytes},
};

/*

QOI - The “Quite OK Image” format for fast, lossless image compression

Dominic Szablewski - https://phoboslab.org


-- LICENSE: The MIT License(MIT)

Copyright(c) 2021 Dominic Szablewski

Permission is hereby granted, free of charge, to any person obtaining a copy of
this software and associated documentation files(the "Software"), to deal in
the Software without restriction, including without limitation the rights to
use, copy, modify, merge, publish, distribute, sublicense, and / or sell copies
of the Software, and to permit persons to whom the Software is furnished to do
so, subject to the following conditions :
The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.


-- About

QOI encodes and decodes images in a lossless format. An encoded QOI image is
usually around 10--30% larger than a decently optimized PNG image.

QOI outperforms simpler PNG encoders in compression ratio and performance. QOI
images are typically 20% smaller than PNGs written with stbi_image but 10%
larger than with libpng. Encoding is 25-50x faster and decoding is 3-4x faster
than stbi_image or libpng.

-- Data Format

A QOI file has a 12 byte header, followed by any number of data "chunks".

struct qoi_header_t {
    char [4];              // magic bytes "qoif"
    unsigned short width;  // image width in pixels (BE)
    unsigned short height; // image height in pixels (BE)
    unsigned int size;     // number of data bytes following this header (BE)
};

The decoder and encoder start with {r: 0, g: 0, b: 0, a: 255} as the previous
pixel value. Pixels are either encoded as
 - a run of the previous pixel
 - an index into a previously seen pixel
 - a difference to the previous pixel value in r,g,b,a
 - full r,g,b,a values

A running array[64] of previously seen pixel values is maintained by the encoder
and decoder. Each pixel that is seen by the encoder and decoder is put into this
array at the position (r^g^b^a) % 64. In the encoder, if the pixel value at this
index matches the current pixel, this index position is written to the stream.

Each chunk starts with a 2, 3 or 4 bit tag, followed by a number of data bits.
The bit length of chunks is divisible by 8 - i.e. all chunks are byte aligned.

QOI_INDEX {
    u8 tag  :  2;   // b00
    u8 idx  :  6;   // 6-bit index into the color index array: 0..63
}

QOI_RUN_8 {
    u8 tag  :  3;   // b010
    u8 run  :  5;   // 5-bit run-length repeating the previous pixel: 1..32
}

QOI_RUN_16 {
    u8 tag  :  3;   // b011
    u16 run : 13;   // 13-bit run-length repeating the previous pixel: 33..8224
}

QOI_DIFF_8 {
    u8 tag  :  2;   // b10
    u8 dr   :  2;   // 2-bit   red channel difference: -1..2
    u8 dg   :  2;   // 2-bit green channel difference: -1..2
    u8 db   :  2;   // 2-bit  blue channel difference: -1..2
}

QOI_DIFF_16 {
    u8 tag  :  3;   // b110
    u8 dr   :  5;   // 5-bit   red channel difference: -15..16
    u8 dg   :  4;   // 4-bit green channel difference:  -7.. 8
    u8 db   :  4;   // 4-bit  blue channel difference:  -7.. 8
}

QOI_DIFF_24 {
    u8 tag  :  4;   // b1110
    u8 dr   :  5;   // 5-bit   red channel difference: -15..16
    u8 dg   :  5;   // 5-bit green channel difference: -15..16
    u8 db   :  5;   // 5-bit  blue channel difference: -15..16
    u8 da   :  5;   // 5-bit alpha channel difference: -15..16
}

QOI_COLOR {
    u8 tag  :  4;   // b1111
    u8 has_r:  1;   //   red byte follows
    u8 has_g:  1;   // green byte follows
    u8 has_b:  1;   //  blue byte follows
    u8 has_a:  1;   // alpha byte follows
    u8 r;           //   red value if has_r == 1: 0..255
    u8 g;           // green value if has_g == 1: 0..255
    u8 b;           //  blue value if has_b == 1: 0..255
    u8 a;           // alpha value if has_a == 1: 0..255
}

The byte stream is padded with 4 zero bytes. Size the longest chunk we can
encounter is 5 bytes (QOI_COLOR with RGBA set), with this padding we just have
to check for an overrun once per decode loop iteration.

*/

const QOI_INDEX: u8 = 0x00; // 00xxxxxx
const QOI_RUN_8: u8 = 0x40; // 010xxxxx
const QOI_RUN_16: u8 = 0x60; // 011xxxxx
const QOI_DIFF_8: u8 = 0x80; // 10xxxxxx
const QOI_DIFF_16: u8 = 0xc0; // 110xxxxx
const QOI_DIFF_24: u8 = 0xe0; // 1110xxxx
const QOI_COLOR: u8 = 0xf0; // 1111xxxx

const QOI_MASK_2: u8 = 0xc0; // 11000000
const QOI_MASK_3: u8 = 0xe0; // 11100000
const QOI_MASK_4: u8 = 0xf0; // 11110000

#[inline(always)]
const fn qui_color_hash(c: Rgba) -> usize {
    ((c.r ^ c.g ^ c.b ^ c.a) & 63) as usize
}

const QOI_MAGIC: u32 = u32::from_be_bytes(*b"qoif");
const QOI_HEADER_SIZE: usize = 12;
const QOI_PADDING: usize = 4;

pub struct UnsafeWriter {
    ptr: *mut u8,
}

impl UnsafeWriter {
    pub unsafe fn fill(&mut self, v: u8, count: usize) {
        write_bytes(self.ptr, v, count);
        self.ptr = self.ptr.add(count);
    }

    pub unsafe fn write_8(&mut self, v: u8) {
        *self.ptr = v;
        self.ptr = self.ptr.add(1);
    }

    pub unsafe fn write_16(&mut self, v: u16) {
        copy_nonoverlapping(v.to_be_bytes().as_ptr(), self.ptr, 2);
        self.ptr = self.ptr.add(2);
    }

    pub unsafe fn write_32(&mut self, v: u32) {
        copy_nonoverlapping(v.to_be_bytes().as_ptr(), self.ptr, 4);
        self.ptr = self.ptr.add(4);
    }
}

pub struct UnsafeReader {
    ptr: *const u8,
    len: usize,
}

impl UnsafeReader {
    pub unsafe fn read_8(&mut self) -> u8 {
        let v = *self.ptr;
        self.ptr = self.ptr.add(1);
        self.len -= 1;
        v
    }

    pub unsafe fn read_16(&mut self) -> u16 {
        let v = u16::from_be_bytes(*(self.ptr as *const [u8; 2]));
        self.ptr = self.ptr.add(2);
        self.len -= 2;
        v
    }

    pub unsafe fn read_32(&mut self) -> u32 {
        let v = u32::from_be_bytes(*(self.ptr as *const [u8; 4]));
        self.ptr = self.ptr.add(4);
        self.len -= 4;
        v
    }
}

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

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

    #[inline(always)]
    const fn new_opaque() -> Self {
        Rgba {
            r: 0,
            g: 0,
            b: 0,
            a: 255,
        }
    }

    #[inline(always)]
    fn read_rgba(bytes: &[u8]) -> Self {
        let mut v = [0; 4];
        v.copy_from_slice(&bytes[..4]);
        unsafe { transmute(v) }
    }

    #[inline(always)]
    fn read_rgb(bytes: &[u8]) -> Self {
        let mut v = [255; 4];
        v.copy_from_slice(&bytes[..3]);
        unsafe { transmute(v) }
    }

    #[inline(always)]
    fn write_rgba(&self, bytes: &mut [u8]) {
        bytes.copy_from_slice(&u32::to_le_bytes(unsafe { transmute(*self) }))
    }

    #[inline(always)]
    fn write_rgb(&self, bytes: &mut [u8]) {
        bytes.copy_from_slice(&u32::to_le_bytes(unsafe { transmute(*self) })[..3])
    }

    fn var(&self, prev: &Self) -> Var {
        Var {
            r: self.r as i16 - prev.r as i16,
            g: self.g as i16 - prev.g as i16,
            b: self.b as i16 - prev.b as i16,
            a: self.a as i16 - prev.a as i16,
        }
    }
}

#[derive(Clone, Copy)]
struct Var {
    r: i16,
    g: i16,
    b: i16,
    a: i16,
}

#[derive(Debug)]
pub enum EncodeError {
    NotEnoughPixels,
    ImageTooLarge { width: u16, height: u16 },
}

/// Encode raw RGB or RGBA pixels into a QOI image in memory. w and h denote the
/// width and height of the pixel data. channels must be either 3 for RGB data
/// or 4 for RGBA.
/// The function either returns NULL on failure (invalid parameters or malloc
/// failed) or a pointer to the encoded data on success. On success the out_len
/// is set to the size in bytes of the encoded data.
/// The returned qoi data should be free()d after user.
pub fn qoi_encode(
    pixels: &[u8],
    w: u16,
    h: u16,
    has_alpha: bool,
) -> Result<Box<[u8]>, EncodeError> {
    let px_len = w as usize * h as usize * (has_alpha as usize + 3);
    if pixels.len() < px_len {
        return Err(EncodeError::NotEnoughPixels);
    }

    let max_size =
        w as usize * h as usize * (has_alpha as usize + 4) + QOI_HEADER_SIZE + QOI_PADDING;

    if u32::try_from(max_size).is_err() {
        return Err(EncodeError::ImageTooLarge {
            width: w,
            height: h,
        });
    }

    let mut bytes = vec![0u8; max_size];

    let mut writer = UnsafeWriter {
        ptr: bytes.as_mut_ptr(),
    };

    unsafe {
        writer.write_32(QOI_MAGIC);
        writer.write_16(w);
        writer.write_16(h);
        writer.write_32(0); // size, will be set later
    }

    let mut index = [Rgba::new(); 64];

    let mut run = 0u16;
    let mut px_prev = Rgba::new_opaque();
    let mut px;

    let px_len = w as usize * h as usize * (has_alpha as usize + 3);
    let px_end = px_len - (has_alpha as usize + 3);
    for px_pos in (0..px_len).step_by(has_alpha as usize + 3) {
        match has_alpha {
            true => px = Rgba::read_rgba(&pixels[px_pos..][..4]),
            false => px = Rgba::read_rgb(&pixels[px_pos..][..3]),
        }

        let v = px.var(&px_prev);

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

        if run > 0 && (run == 0x2020 || px != px_prev || px_pos == px_end) {
            if run < 33 {
                run -= 1;
                unsafe { writer.write_8(QOI_RUN_8 | run as u8) }
            } else {
                run -= 33;
                unsafe {
                    writer.write_8(QOI_RUN_16 | (run >> 8) as u8);
                    writer.write_8(run as u8);
                }
            }
            run = 0;
        }

        if px != px_prev {
            let index_pos = qui_color_hash(px);

            if index[index_pos] == px {
                unsafe { writer.write_8(QOI_INDEX | index_pos as u8) }
            } else {
                index[index_pos] = px;

                if v.a == 0 && v.r > -2 && v.r < 3 && v.g > -2 && v.g < 3 && v.b > -2 && v.b < 3 {
                    unsafe {
                        writer.write_8(
                            QOI_DIFF_8
                                | ((v.r + 1) << 4) as u8
                                | ((v.g + 1) << 2) as u8
                                | (v.b + 1) as u8,
                        );
                    }
                } else if v.a == 0
                    && v.r > -16
                    && v.r < 17
                    && v.g > -8
                    && v.g < 9
                    && v.b > -8
                    && v.b < 9
                {
                    unsafe {
                        writer.write_8(QOI_DIFF_16 | (v.r + 15) as u8);
                        writer.write_8(((v.g + 7) << 4) as u8 | (v.b + 7) as u8);
                    }
                } else if v.r > -16
                    && v.r < 17
                    && v.g > -16
                    && v.g < 17
                    && v.b > -16
                    && v.b < 17
                    && v.a > -16
                    && v.a < 17
                {
                    {
                        unsafe {
                            writer.write_8(QOI_DIFF_24 | ((v.r + 15) >> 1) as u8);
                            writer.write_8(
                                ((v.r + 15) << 7) as u8
                                    | ((v.g + 15) << 2) as u8
                                    | ((v.b + 15) >> 3) as u8,
                            );
                            writer.write_8(((v.b + 15) << 5) as u8 | (v.a + 15) as u8);
                        }
                    }
                } else {
                    unsafe {
                        writer.write_8(
                            QOI_COLOR
                                | (if v.r != 0 { 8 } else { 0 })
                                | (if v.g != 0 { 4 } else { 0 })
                                | (if v.b != 0 { 2 } else { 0 })
                                | (if v.a != 0 { 1 } else { 0 }),
                        );
                        if v.r != 0 {
                            writer.write_8(px.r);
                        }
                        if v.g != 0 {
                            writer.write_8(px.g);
                        }
                        if v.b != 0 {
                            writer.write_8(px.b);
                        }
                        if v.a != 0 {
                            writer.write_8(px.a);
                        }
                    }
                }
            }
        }
        px_prev = px;
    }

    unsafe { writer.fill(0, QOI_PADDING) };

    let size = unsafe { writer.ptr.offset_from(bytes.as_ptr()) } as usize;

    bytes.truncate(size);

    let data_len = size - QOI_HEADER_SIZE;

    bytes[8..12].copy_from_slice(&(data_len as u32).to_be_bytes());
    Ok(bytes.into_boxed_slice())
}

#[derive(Debug)]
pub enum DecodeError {
    MissingHeader,
    DataIsTooSmall,
    InvalidMagic,
}

// Decode a QOI image from memory into either raw RGB (channels=3) or RGBA
// (channels=4) pixel data.
// The function either returns NULL on failure (invalid parameters or malloc
// failed) or a pointer to the decoded pixels. On success out_w and out_h will
// be set to the width and height of the decoded image.
// The returned pixel data should be free()d after use.
pub fn qoi_decode(
    bytes: &[u8],
    has_alpha: bool,
    out_w: &mut u16,
    out_h: &mut u16,
) -> Result<Box<[u8]>, DecodeError> {
    if bytes.len() < QOI_HEADER_SIZE {
        return Err(DecodeError::MissingHeader);
    }

    let mut reader = UnsafeReader {
        ptr: bytes.as_ptr(),
        len: bytes.len(),
    };

    let magic = unsafe { reader.read_32() };
    let w = unsafe { reader.read_16() };
    let h = unsafe { reader.read_16() };
    let data_len = unsafe { reader.read_32() };

    *out_w = w;
    *out_h = h;

    if magic != QOI_MAGIC {
        return Err(DecodeError::InvalidMagic);
    }

    if bytes.len() != data_len as usize + QOI_HEADER_SIZE {
        return Err(DecodeError::DataIsTooSmall);
    }

    if w == 0 || h == 0 {
        return Ok(Box::new([]));
    }

    let px_len = w as usize * h as usize * (has_alpha as usize + 3);
    let mut pixels = vec![0; px_len];

    let mut px = Rgba::new_opaque();
    let mut index = [Rgba::new(); 64];

    let mut run = 0u16;
    for px_pos in (0..px_len).step_by(has_alpha as usize + 3) {
        if run > 0 {
            run -= 1;
        } else if reader.len >= QOI_PADDING {
            let b1 = unsafe { reader.read_8() };

            if (b1 & QOI_MASK_2) == QOI_INDEX {
                px = index[(b1 ^ QOI_INDEX) as usize];
            } else if (b1 & QOI_MASK_3) == QOI_RUN_8 {
                run = (b1 & 0x1f) as u16;
            } else if (b1 & QOI_MASK_3) == QOI_RUN_16 {
                let b2 = unsafe { reader.read_8() };
                run = ((((b1 & 0x1f) as u16) << 8) | (b2 as u16)) + 32;
            } else if (b1 & QOI_MASK_2) == QOI_DIFF_8 {
                px.r = (px.r as i16 + ((b1 >> 4) & 0x03) as i16 - 1) as u8;
                px.g = (px.g as i16 + ((b1 >> 2) & 0x03) as i16 - 1) as u8;
                px.b = (px.b as i16 + (b1 & 0x03) as i16 - 1) as u8;
            } else if (b1 & QOI_MASK_3) == QOI_DIFF_16 {
                let b2 = unsafe { reader.read_8() };
                px.r = (px.r as i16 + (b1 & 0x1f) as i16 - 15) as u8;
                px.g = (px.g as i16 + (b2 >> 4) as i16 - 7) as u8;
                px.b = (px.b as i16 + (b2 & 0x0f) as i16 - 7) as u8;
            } else if (b1 & QOI_MASK_4) == QOI_DIFF_24 {
                let b2 = unsafe { reader.read_8() };
                let b3 = unsafe { reader.read_8() };
                px.r = (px.r as i16 + (((b1 & 0x0f) << 1) | (b2 >> 7)) as i16 - 15) as u8;
                px.g = (px.g as i16 + ((b2 & 0x7c) >> 2) as i16 - 15) as u8;
                px.b = (px.b as i16 + (((b2 & 0x03) << 3) | ((b3 & 0xe0) >> 5)) as i16 - 15) as u8;
                px.a = (px.a as i16 + (b3 & 0x1f) as i16 - 15) as u8;
            } else if (b1 & QOI_MASK_4) == QOI_COLOR {
                if (b1 & 8) != 0 {
                    px.r = unsafe { reader.read_8() };
                }
                if (b1 & 4) != 0 {
                    px.g = unsafe { reader.read_8() };
                }
                if (b1 & 2) != 0 {
                    px.b = unsafe { reader.read_8() };
                }
                if (b1 & 1) != 0 {
                    px.a = unsafe { reader.read_8() };
                }
            }

            index[qui_color_hash(px)] = px;
        }

        match has_alpha {
            true => px.write_rgba(&mut pixels[px_pos..][..4]),
            false => px.write_rgb(&mut pixels[px_pos..][..3]),
        }
    }

    return Ok(pixels.into_boxed_slice());
}