zardzewialy-dekoder-cif/src/main.rs

#![feature(slice_internals)]
#![feature(read_buf)]
#![feature(new_uninit)]
#![feature(maybe_uninit_slice)]

use unsafe_std::bufread::{UnsafeBufRead, UnsafeLines};
use unsafe_std::bufreader::UnsafeBufReader;

use arrayvec::{ArrayString, ArrayVec};
use png::{BitDepth, ColorType, EncodingError};
use std::fs::File;
use std::io::BufWriter;
use std::{env, io, process};
use thiserror::Error;

mod nummap;
mod nummatch;
mod unsafe_std;

#[cfg(test)]
mod test;

#[derive(Error, Debug)]
#[allow(clippy::enum_variant_names)]
pub enum CifError {
    #[error("The CIF number string was too long")]
    CapacityError,
    #[error("CIF Number is invalid: {0:?}")]
    GrammarError(String),
    #[error("An IO error has occurred: {0}")]
    IoError(io::Error),
    #[error("CIF header is invalid")]
    InvalidHeaderError,
    #[error("CIF size is invalid: {0:?}")]
    InvalidSizeError(String),
    #[error("CIF metadata corrupted: {0:?}")]
    InvalidMetadataError(String),
    #[error("Invalid CIF pixel: {0}")]
    InvalidPixelError(String),
    #[error("CIF image doesn't contain any pixels")]
    NoDataError,
    #[error("PNG error: {0}")]
    PngError(EncodingError),
    #[error("Illegal whitespace: {0:?}")]
    IllegalWhitespaceError(String),
}

#[allow(clippy::enum_glob_use)]
use crate::CifError::*;

impl<T> From<arrayvec::CapacityError<T>> for CifError {
    fn from(_: arrayvec::CapacityError<T>) -> Self {
        CapacityError
    }
}

impl From<io::Error> for CifError {
    fn from(e: io::Error) -> Self {
        IoError(e)
    }
}

impl From<EncodingError> for CifError {
    fn from(e: EncodingError) -> Self {
        PngError(e)
    }
}

type Result<T> = std::result::Result<T, CifError>;

#[inline]
fn join_to_arraystring<const C: usize>(elems: &[&str], str: &mut ArrayString<C>) -> Result<()> {
    for (i, item) in elems.iter().enumerate() {
        if i > 0 {
            str.try_push(' ')?;
        }
        str.try_push_str(item)?;
        //str.push_str(item);
    }
    Ok(())
}

fn lines_read_noempty<T: UnsafeBufRead>(lines: &mut UnsafeLines<T>) -> Result<Option<String>> {
    while let Some(s) = lines.next().transpose()? {
        if s.is_empty() {
            continue
        }
        assert_trailing_leading(&s)?;
        return Ok(Some(s))
    }
    Ok(None)
}

#[inline]
fn assert_trailing_leading(s: &str) -> Result<()> {
    if s.is_empty() || s.as_bytes()[0] == 32 || s.as_bytes()[s.len() - 1] == 32 {
        return Err(IllegalWhitespaceError(s.to_string()))
    }
    Ok(())
}

fn unfuck_cif_number(number: &str) -> Result<u32> {
    assert_trailing_leading(number)?;
    let mut buf = ArrayString::<1024>::new();
    let split = number.split_ascii_whitespace().collect::<ArrayVec<_, 32>>();
    let splitlen = split.len();
    for (i, word) in split.iter().enumerate() {
        if !word.starts_with("ty") {
            continue
        }
        let thousands = match *word {
            "tysiąc" => {
                if i == 0 {
                    &1_u16
                } else {
                    return Err(GrammarError(number.to_string()))
                }
            }
            "tysięcy" | "tysiące" => {
                if i > 0 {
                    join_to_arraystring(&split[..i], &mut buf)?;
                    nummap::U16_NUMMAP.get(&buf).ok_or_else(|| GrammarError(number.to_string()))?
                    //nummatch::u_16(&buf)?
                } else {
                    return Err(GrammarError(number.to_string()))
                }
            }
            _ => return Err(GrammarError(number.to_string())),
        };
        let remainder = if i + 1 < splitlen {
            buf.clear();
            join_to_arraystring(&split[i + 1..], &mut buf)?;
            nummap::U16_NUMMAP.get(&buf).ok_or_else(|| GrammarError(number.to_string()))?
            //nummatch::u_16(&buf)?
        } else {
            &0_u16
        };
        return Ok(u32::from(*thousands) * 1000 + u32::from(*remainder))
    }
    Ok((*nummap::U16_NUMMAP.get(number).ok_or_else(|| GrammarError(number.to_string()))?).into())
}

fn unfuck_cif_number_fast(number: &str) -> Result<u8> {
    assert_trailing_leading(number)?;
    let mut buf = ArrayString::<1024>::new();
    let split = number.split_ascii_whitespace().take(8).collect::<ArrayVec<_, 3>>();

    join_to_arraystring(&split, &mut buf)?;
    let number = *(nummap::U8_NUMMAP.get(&buf).ok_or_else(|| InvalidPixelError(number.to_string()))?);
    //let number = nummatch::u_8(&buf)?;

    Ok(number)
}

fn unfuck_color_type(color_type: &str) -> Result<ColorType> {
    assert_trailing_leading(color_type)?;
    let mut buf = ArrayString::<1024>::new();
    let split = color_type.split_ascii_whitespace().take(8).collect::<ArrayVec<_, 2>>();

    join_to_arraystring(&split, &mut buf)?;
    //let color_type = *(nummap::COLOR_TYPE_NUMMAP.get(&buf).ok_or_else(|| InvalidSizeError(color_type.to_string()))?);
    let color_type = nummatch::color_type(&buf)?;

    Ok(color_type)
}

fn split_version(version_str: String) -> Result<()> {
    let version_split = version_str.split_ascii_whitespace().take(2).collect::<ArrayVec<_, 2>>();
    if version_split.len() != 2 || version_split[0] != "WERSJA" || version_split[1] != "jeden" {
        return Err(InvalidHeaderError)
    }
    Ok(())
}

fn spliz_size(size_string: &str) -> Result<ArrayVec<&str, 3>> {
    let size_split = size_string
        .strip_prefix("ROZMIAR")
        .ok_or_else(|| InvalidSizeError(size_string.to_string()))?
        .split(',')
        .take(3)
        .collect::<ArrayVec<_, 3>>();
    if size_split.len() != 3 {
        return Err(InvalidSizeError(size_string.to_string()))
    }
    Ok(size_split)
}

fn parse_metadata<B: std::io::Read>(lines: &mut UnsafeLines<UnsafeBufReader<B>>) -> Result<(String, ArrayVec<String, 124>, ArrayVec<String, 124>)> {
    let mut md_keylock = false;
    let mut md_key = ArrayString::<1024>::new();
    let mut md_key_vec = ArrayVec::<String, 124>::new();
    let mut md_val = ArrayVec::<String, 124>::new();

    loop {
        if let Some(line) = lines_read_noempty(lines)? {
            if md_keylock {
                let split = line.splitn(2, ' ').collect::<ArrayVec<_, 2>>();
                md_key.try_push_str(split[0])?;
                if split.len() != 2 {
                    continue
                }
                md_key_vec.push(md_key.to_string());
                md_val.push(split[1].trim_start().to_string());
                md_keylock = false;
                md_key.clear();
                continue
            }
            if let Some(sline) = line.strip_prefix("METADANE ") {
                let split = sline.trim_start().splitn(2, ' ').collect::<ArrayVec<_, 2>>();
                match split.len() {
                    0 => return Err(InvalidMetadataError(line.to_string())),
                    1 => {
                        md_key.try_push_str(&sline)?;
                        md_keylock = true;
                    }
                    _ => {
                        md_key.try_push_str(&sline)?;
                        md_key_vec.push(split[0].trim_start().to_string());
                        md_val.push(split[1].trim_start().to_string());
                        md_key.clear();
                    }
                }
                continue
            }
            return Ok((line, md_key_vec, md_val))
        }
        return Err(NoDataError)
    }
}

fn push_to_pixmap(s: String, bytes_per_pixel: u8, pixmap: &mut Vec<u8>) -> Result<()> {
    let px_split = s.split(';').collect::<ArrayVec<_, 4>>();
    if px_split.len() != bytes_per_pixel as usize {
        return Err(InvalidPixelError(s.to_string()))
    }
    for px in px_split {
        pixmap.push(unfuck_cif_number_fast(px.trim_start())?);
    }
    Ok(())
}

#[allow(clippy::too_many_lines)] // bruh
fn decoder(inputfile: &str) -> Result<(u32, u32, ColorType, ArrayVec<String, 124_usize>, ArrayVec<String, 124_usize>, Vec<u8>)> {
    let infile = File::open(inputfile)?;
    let mut lines = UnsafeBufReader::new(infile).lines();

    // "CIF: <flags>"
    if !lines.next().transpose()?.ok_or(InvalidHeaderError)?.starts_with("CIF:") {
        return Err(InvalidHeaderError)
    }

    // "WERSJA jeden"
    let version_str = lines_read_noempty(&mut lines)?.ok_or(InvalidHeaderError)?;
    split_version(version_str)?;

    // "ROZMIAR szerokość: osiem, wysokość: osiem, bitów_na_piksel: trzydzieści dwa"
    let size_string = lines_read_noempty(&mut lines)?.ok_or(InvalidHeaderError)?;
    let size_split = spliz_size(&size_string)?;

    let width = unfuck_cif_number(
        size_split[0]
            .trim_start()
            .strip_prefix("szerokość:")
            .ok_or_else(|| InvalidSizeError(size_string.to_string()))?
            .trim_start(),
    )?;
    let height = unfuck_cif_number(
        size_split[1]
            .trim_start()
            .strip_prefix("wysokość:")
            .ok_or_else(|| InvalidSizeError(size_string.to_string()))?
            .trim_start(),
    )?;
    let format = unfuck_color_type(
        size_split[2]
            .trim_start()
            .strip_prefix("bitów_na_piksel:")
            .ok_or_else(|| InvalidSizeError(size_string.to_string()))?
            .trim_start(),
    )?;

    if width == 0 || height == 0 {
        return Err(NoDataError)
    }

    let bytes_per_pixel = if format == ColorType::Rgb { 3_u8 } else { 4_u8 };
    let mut pixmap = Vec::with_capacity((bytes_per_pixel as u32 * width * height) as usize);

    let (first, md_key_vec, md_val) = parse_metadata(&mut lines)?;

    if !first.is_empty() {
        push_to_pixmap(first, bytes_per_pixel, &mut pixmap)?;
    }

    while let Some(line) = lines_read_noempty(&mut lines)? {
        push_to_pixmap(line, bytes_per_pixel, &mut pixmap)?;
    }

    Ok((width, height, format, md_key_vec, md_val, pixmap))
}

fn write_png(outfile: &str, width: u32, height: u32, format: ColorType, md_key: ArrayVec<String, 124>, md_val: ArrayVec<String, 124>, pixmap: Vec<u8>) -> Result<()> {
    let outfile = File::create(outfile)?;
    let rawriter = BufWriter::new(outfile);
    let mut outpng = png::Encoder::new(rawriter, width, height);

    for i in 0..md_key.len() {
        if md_key[i].is_empty() {
            break
        }
        outpng.add_itxt_chunk(md_key[i].clone(), md_val[i].clone())?;
    }

    outpng.set_color(format);
    outpng.set_depth(BitDepth::Eight);

    let mut outpng = outpng.write_header()?;
    outpng.write_image_data(&pixmap)?;
    Ok(())
}

fn main() {
    let argv = env::args().skip(1).take(2).collect::<ArrayVec<_, 2>>();
    if argv.len() != 2 {
        eprintln!("Invalid usage:\nUsage: zardzewialy-dekoder-cif <input cif file> <output bmp file>");
        process::exit(1);
    }

    let data = decoder(&argv[0]);
    if let Err(e) = data {
        eprintln!("Error while parsing cif: {}", e);
        process::exit(1);
    }
    let data = data.unwrap();
    if let Err(e) = write_png(&argv[1], data.0, data.1, data.2, data.3, data.4, data.5) {
        eprintln!("Error while writing image: {}", e);
    }
}