diff options
| -rw-r--r-- | .gitignore | 1 | ||||
| -rw-r--r-- | Cargo.lock | 46 | ||||
| -rw-r--r-- | Cargo.toml | 5 | ||||
| -rw-r--r-- | byte_encode_derive/Cargo.lock | 46 | ||||
| -rw-r--r-- | byte_encode_derive/Cargo.toml | 12 | ||||
| -rw-r--r-- | byte_encode_derive/src/lib.rs | 128 | ||||
| -rw-r--r-- | src/bmp.rs | 83 | ||||
| -rw-r--r-- | src/byte_encode.rs | 27 | ||||
| -rw-r--r-- | src/crc.rs | 37 | ||||
| -rw-r--r-- | src/image.rs | 4 | ||||
| -rw-r--r-- | src/main.rs | 17 | ||||
| -rw-r--r-- | src/png.rs | 251 | ||||
| -rw-r--r-- | src/reader.rs | 69 | ||||
| -rw-r--r-- | src/zlib.rs | 258 |
14 files changed, 941 insertions, 43 deletions
@@ -1,2 +1,3 @@ /target +*/target .direnv @@ -3,5 +3,51 @@ version = 4 [[package]] +name = "byte_encode_derive" +version = "0.1.0" +dependencies = [ + "quote", + "syn", +] + +[[package]] +name = "proc-macro2" +version = "1.0.106" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8fd00f0bb2e90d81d1044c2b32617f68fcb9fa3bb7640c23e9c748e53fb30934" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "quote" +version = "1.0.46" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dfbc457d0c7a0759a614551b11a6409e5951f6c7537be1f1b7682b9ae9230368" +dependencies = [ + "proc-macro2", +] + +[[package]] +name = "syn" +version = "2.0.118" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1b9ae57f904213ebb649ce6895b8a66c66f0203b9319718f69a5612a065b1422" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] name = "testing" version = "0.1.0" +dependencies = [ + "byte_encode_derive", +] + +[[package]] +name = "unicode-ident" +version = "1.0.24" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75" @@ -1,4 +1,7 @@ [package] name = "testing" version = "0.1.0" -edition = "2024"
\ No newline at end of file +edition = "2024" + +[dependencies] +byte_encode_derive = { version = "0.1.0", path = "byte_encode_derive" } diff --git a/byte_encode_derive/Cargo.lock b/byte_encode_derive/Cargo.lock new file mode 100644 index 0000000..2616932 --- /dev/null +++ b/byte_encode_derive/Cargo.lock @@ -0,0 +1,46 @@ +# This file is automatically @generated by Cargo. +# It is not intended for manual editing. +version = 4 + +[[package]] +name = "byte_encode_derive" +version = "0.1.0" +dependencies = [ + "quote", + "syn", +] + +[[package]] +name = "proc-macro2" +version = "1.0.106" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8fd00f0bb2e90d81d1044c2b32617f68fcb9fa3bb7640c23e9c748e53fb30934" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "quote" +version = "1.0.46" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dfbc457d0c7a0759a614551b11a6409e5951f6c7537be1f1b7682b9ae9230368" +dependencies = [ + "proc-macro2", +] + +[[package]] +name = "syn" +version = "2.0.118" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1b9ae57f904213ebb649ce6895b8a66c66f0203b9319718f69a5612a065b1422" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "unicode-ident" +version = "1.0.24" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75" diff --git a/byte_encode_derive/Cargo.toml b/byte_encode_derive/Cargo.toml new file mode 100644 index 0000000..832dadd --- /dev/null +++ b/byte_encode_derive/Cargo.toml @@ -0,0 +1,12 @@ +[package] +name = "byte_encode_derive" +version = "0.1.0" +edition = "2024" + +[lib] +proc-macro = true + +[dependencies] + +syn = { version = "2", features = ["full", "parsing"] } +quote = "1"
\ No newline at end of file diff --git a/byte_encode_derive/src/lib.rs b/byte_encode_derive/src/lib.rs new file mode 100644 index 0000000..5970c66 --- /dev/null +++ b/byte_encode_derive/src/lib.rs @@ -0,0 +1,128 @@ +// Code mostly made by gemini somehow + +use proc_macro::TokenStream; +use quote::quote; +use syn::{parse_macro_input, Data, DeriveInput, Fields}; + +#[proc_macro_derive(ByteEncode)] +pub fn derive_byte_encode(input: TokenStream) -> TokenStream { + let input = parse_macro_input!(input as DeriveInput); + let name = input.ident; + + // We only support structs with named fields for this example + let fields = match input.data { + Data::Struct(data) => match data.fields { + Fields::Named(fields) => fields.named, + _ => panic!("Only named fields are supported"), + }, + _ => panic!("Only structs are supported"), + }; + + // --- 1. Generate from_le_bytes parsing logic --- + let mut current_offset = quote! { 0 }; + let from_le_bytes_initializers = fields.iter().map(|f| { + let field_name = &f.ident; + let field_type = &f.ty; + + // Calculate start and end bounds for the current field slice + let start = quote! { #current_offset }; + let end = quote! { #start + <#field_type>::SIZE }; + + // Update the offset tracking for the next iteration loop + current_offset = quote! { #end }; + + quote! { + #field_name: <#field_type as ByteEncode<{<#field_type>::SIZE}>>::from_le_bytes( + buffer[#start..#end].try_into().unwrap() + ) + } + }); + let mut current_offset = quote! { 0 }; + let from_be_bytes_initializers = fields.iter().map(|f| { + let field_name = &f.ident; + let field_type = &f.ty; + + // Calculate start and end bounds for the current field slice + let start = quote! { #current_offset }; + let end = quote! { #start + <#field_type>::SIZE }; + + // Update the offset tracking for the next iteration loop + current_offset = quote! { #end }; + + quote! { + #field_name: <#field_type as ByteEncode<{<#field_type>::SIZE}>>::from_be_bytes( + &mut buffer[#start..#end].try_into().unwrap() + ) + } + }); + + + // --- 2. Generate to_le_bytes serialization logic --- + let mut current_offset = quote! { 0 }; + let to_le_bytes_writers = fields.iter().map(|f| { + let field_name = &f.ident; + let field_type = &f.ty; + + let start = quote! { #current_offset }; + let end = quote! { #start + <#field_type>::SIZE }; + + current_offset = quote! { #end }; + + quote! { + res[#start..#end].copy_from_slice(&self.#field_name.to_le_bytes()); + } + }); + let mut current_offset = quote! { 0 }; + let to_be_bytes_writers = fields.iter().map(|f| { + let field_name = &f.ident; + let field_type = &f.ty; + + let start = quote! { #current_offset }; + let end = quote! { #start + <#field_type>::SIZE }; + + current_offset = quote! { #end }; + + quote! { + res[#start..#end].copy_from_slice(&self.#field_name.to_be_bytes()); + } + }); + + + // --- 3. Construct final token stream --- + // This dynamically sums up sizes like: 0 + Field1::SIZE + Field2::SIZE ... + let total_size_expr = fields.iter().fold(quote! { 0 }, |acc, f| { + let field_type = &f.ty; + quote! { #acc + <#field_type>::SIZE } + }); + + let expanded = quote! { + impl ByteEncode<{ #total_size_expr }> for #name { + fn from_le_bytes(buffer: &[u8; { #total_size_expr }]) -> Self { + Self { + #( #from_le_bytes_initializers, )* + } + } + + fn to_le_bytes(&self) -> [u8; { #total_size_expr }] { + let mut res = [0u8; { #total_size_expr }]; + #( #to_le_bytes_writers )* + res + } + + + fn from_be_bytes(buffer: &mut [u8; { #total_size_expr }]) -> Self { + Self { + #( #from_be_bytes_initializers, )* + } + } + + fn to_be_bytes(&self) -> [u8; { #total_size_expr }] { + let mut res = [0u8; { #total_size_expr }]; + #( #to_be_bytes_writers )* + res + } + } + }; + + TokenStream::from(expanded) +} @@ -1,13 +1,53 @@ -use crate::{image::{ColorType, Image, ImageBase, PixelArr, RGB, RGBA}, reader::FileReader, writer::FileWriter}; +use crate::{byte_encode::ByteEncode, image::{ColorType, Image, ImageBase, PixelArr, RGB, RGBA}, reader::FileReader, writer::FileWriter}; + +// extern crate byte_encode_derive; +use byte_encode_derive::ByteEncode; + + +trait BitmapHeader +{ + fn get_width(&self) -> u32; + fn get_height(&self) -> u32; + fn get_bpp(&self) -> u8; +} + +#[derive(Clone, ByteEncode)] +struct BITMAPINFOHEADER +{ + width: u32, + height: u32, + _color_planes: u16, + bpp: u16, + compression_method: u32, + _image_size: u32, + _horr_res: u32, + _vert_res: u32, + _palette_size: u32, + _important_colors: u32 +} + +impl BitmapHeader for BITMAPINFOHEADER +{ + fn get_width(&self) -> u32 { + self.width + } + fn get_height(&self) -> u32 { + self.height + } + fn get_bpp(&self) -> u8 { + self.bpp as u8 + } +} + impl<T: ColorType> Image<T> for BMPImage { - fn read_image(reader: &mut FileReader) -> Result<ImageBase<T>, &'static str> + fn read_image(reader: &mut FileReader) -> Result<ImageBase<T>, String> { let magic: [u8; 2] = reader.read_array(); if magic != [0x42, 0x4d] { - return Err("Not a bitmap!"); + return Err("Not a bitmap!".to_owned()); } let _file_size: u32 = reader.read(); @@ -22,39 +62,34 @@ impl<T: ColorType> Image<T> for BMPImage let height: u32; let bpp: u8; - if header_size == 40 + + if header_size == 12 // BITMAPCOREHEADER { - width = reader.read(); - height = reader.read(); + width = reader.read::<2, u16>() as u32; + height = reader.read::<2, u16>() as u32; let _color_planes: u16 = reader.read(); bpp = reader.read::<2, u16>() as u8; - let compression_method: u32 = reader.read(); - let _image_size: u32 = reader.read(); - let _horr_res: u32 = reader.read(); - let _vert_res: u32 = reader.read(); - let _color_count: u32 = reader.read(); - let _impotant_color_count: u32 = reader.read(); + } + else if header_size == 40 // BITMAPINFOHEADER + { + let header: BITMAPINFOHEADER = reader.read(); + width = header.get_width(); + height = header.get_height(); + bpp = header.get_bpp(); - if compression_method != 0 + if header.compression_method != 0 { - return Err("Compresssion method not supported"); + return Err(format!("Bitmap compresssion method {0} not supported", header.compression_method)); } } - else if header_size == 12 - { - width = reader.read::<2, u16>() as u32; - height = reader.read::<2, u16>() as u32; - let _color_planes: u16 = reader.read(); - bpp = reader.read::<2, u16>() as u8; - } else { - return Err("Header size not supported yet"); + return Err(format!("Bitmap header size {header_size} not supported yet")); } if bpp != 24 { - return Err("Only bpp of 24 is supported so far"); + return Err(format!("bpp {bpp} not supported yet, only 24 is supported")); } let row_size = (((bpp as u32) * width + 31)/32) * 4; @@ -76,7 +111,7 @@ impl<T: ColorType> Image<T> for BMPImage return Ok(ImageBase { - bit_depth: bpp, + bpp, pixels: pixel_arr, }) } diff --git a/src/byte_encode.rs b/src/byte_encode.rs index f5cad19..4419f7f 100644 --- a/src/byte_encode.rs +++ b/src/byte_encode.rs @@ -1,12 +1,20 @@ -pub trait ByteEncode<const N: usize>: Sized + Copy +pub trait ByteEncode<const N: usize>: Sized + Clone { fn from_le_bytes(buffer: &[u8; N]) -> Self; fn to_le_bytes(&self) -> [u8; N]; + const SIZE: usize = N; - // fn from_be_bytes(buffer: &[u8; N]) -> Self - // { - // Self::from_le_bytes(buffer.reversed()) - // } + fn from_be_bytes(buffer: &mut [u8; N]) -> Self + { + buffer.reverse(); + Self::from_le_bytes(buffer) + } + fn to_be_bytes(&self) -> [u8; N] + { + let mut buffer = Self::to_le_bytes(self); + buffer.reverse(); + buffer + } } impl ByteEncode<4> for u32 @@ -18,6 +26,15 @@ impl ByteEncode<4> for u32 u32::to_le_bytes(*self) } } +// impl ByteEncode<4> for usize +// { +// fn from_le_bytes(buffer: &[u8; 4]) -> Self { +// u32::from_le_bytes(*buffer) as usize +// } +// fn to_le_bytes(&self) -> [u8; 4] { +// u32::to_le_bytes(*self as u32) +// } +// } impl ByteEncode<2> for u16 { fn from_le_bytes(buffer: &[u8; 2]) -> Self { diff --git a/src/crc.rs b/src/crc.rs new file mode 100644 index 0000000..f24da02 --- /dev/null +++ b/src/crc.rs @@ -0,0 +1,37 @@ +const CRC_TABLE: [u32; 256] = make_crc_table(); + +const fn make_crc_table() -> [u32; 256] +{ + let mut table = [0u32; 256]; + let mut c: u32; + let mut k: i32; + let mut n: u32 = 0; + while n < 256 + { + c = n; + k = 0; + while k < 8 + { + if c & 1 == 1 + { + c = 0xedb88320 ^ (c >> 1); + } + else + { + c = c >> 1; + } + k += 1; + } + table[n as usize] = c; + n += 1; + } + table +} + +pub fn update_crc(curr_crc: &mut u32, buff: &[u8]) -> () +{ + for x in buff + { + *curr_crc = CRC_TABLE[((*curr_crc as u8 ^ x) & 0xff) as usize] ^ (*curr_crc >> 8); + } +} diff --git a/src/image.rs b/src/image.rs index 9b42b27..fa1a18d 100644 --- a/src/image.rs +++ b/src/image.rs @@ -131,12 +131,12 @@ impl<T: ColorType> ops::IndexMut<(usize, usize)> for PixelArr<T> #[derive(Debug)] pub struct ImageBase<T: ColorType> { - pub bit_depth: u8, + pub bpp: u8, pub pixels: PixelArr<T> } pub trait Image<T: ColorType> { - fn read_image(reader: &mut FileReader) -> Result<ImageBase<T>, &'static str>; + fn read_image(reader: &mut FileReader) -> Result<ImageBase<T>, String>; fn write_image(image: &ImageBase<T>, writer: &mut FileWriter)-> (); } diff --git a/src/main.rs b/src/main.rs index 0da4f13..df6b948 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,4 +1,4 @@ -use crate::{bmp::BMPImage, image::{Image, ImageBase, RGB}, reader::FileReader, writer::FileWriter}; +use crate::{bmp::BMPImage, png::PNGImage, image::{Image, ImageBase, RGB}, reader::FileReader, writer::FileWriter}; use std::{env, fs::File, io::{BufReader, BufWriter}}; mod reader; @@ -6,20 +6,23 @@ mod image; mod bmp; mod writer; mod byte_encode; +mod png; +mod crc; +mod zlib; fn main() { let args: Vec<String> = env::args().collect(); if args.len() != 3 { - panic!("USAGE: <program> infline.bmp outfile.bmp"); + panic!("USAGE: <program> infline.png outfile.bmp"); } - let file = File::open(args[1].clone()).expect("File does not exist"); - let mut reader = FileReader::new(BufReader::new(file)); - let _bmp: ImageBase<RGB<u8>> = BMPImage::read_image(&mut reader).expect("Bitmap had error"); + let infile = File::open(args[1].clone()).expect("infile does not exist"); + let mut reader = FileReader::new(BufReader::new(infile)); + let png: ImageBase<RGB<u8>> = PNGImage::read_image(&mut reader).expect("PNG had error"); - let out_file = File::create(args[2].clone()).expect("Cannot create file"); + let out_file = File::create(args[2].clone()).expect("Cannot create outfile"); let mut writer = FileWriter::new(BufWriter::new(out_file)); - BMPImage::write_image(&_bmp, &mut writer); + BMPImage::write_image(&png, &mut writer); // println!("{_bmp:?}"); } diff --git a/src/png.rs b/src/png.rs new file mode 100644 index 0000000..ef1aae0 --- /dev/null +++ b/src/png.rs @@ -0,0 +1,251 @@ +use std::collections::{self, HashMap}; + +use byte_encode_derive::ByteEncode; + +use crate::{byte_encode::ByteEncode, crc::update_crc, image::{ColorType, Image, ImageBase}, reader::FileReader, zlib::{zlib_decode, ByteBuffer}}; + +#[derive(Default)] +pub struct PNGImage +{ + IHDR: Option<IHDR>, + IDAT: Option<IDAT>, + IEND: Option<IEND>, + tEXt: Option<tEXt>, +} + +trait Mergable: Sized +{ + fn merge(a: Self, b: Self) -> Self; +} + +impl<T: Mergable + Sized + Clone> Mergable for Option<T> +{ + fn merge(a: Option<T>, b: Option<T>) -> Option<T> + { + if let Some(ref a_content) = a + { + if let Some(b_content) = b + { + Some(T::merge(a_content.clone(), b_content)) + } + else + { + a + } + } + else + { + b + } + } +} + +trait PNGChunk: Mergable +{ + fn read(reader: &mut FileReader, length: usize) -> Self; + const CHUNK_TYPE: [char; 4]; +} + +#[derive(ByteEncode, Clone)] +struct IHDR +{ + pub width: u32, + pub height: u32, + pub bit_depth: u8, + pub color_type: u8, + pub compression_method: u8, + pub filter_method: u8, + pub interlace_method: u8 +} +impl PNGChunk for IHDR +{ + const CHUNK_TYPE: [char; 4] = ['I', 'H', 'D', 'R']; + fn read(reader: &mut FileReader, length: usize) -> Self { + reader.read() + } +} +impl Mergable for IHDR { + fn merge(a: Self, b: Self) -> Self { + todo!() + } +} +#[derive(ByteEncode, Clone)] +struct IEND +{ +} +impl PNGChunk for IEND +{ + const CHUNK_TYPE: [char; 4] = ['I', 'E', 'N', 'D']; + fn read(reader: &mut FileReader, length: usize) -> Self { + Self {} + } +} +impl Mergable for IEND { + fn merge(a: Self, b: Self) -> Self { + todo!() + } +} + +#[derive(Clone)] +struct tEXt +{ + keywords: HashMap<String, String> +} +impl PNGChunk for tEXt +{ + const CHUNK_TYPE: [char; 4] = ['t', 'E', 'X', 't']; + fn read(reader: &mut FileReader, length: usize) -> Self + { + let mut i = 0; + let mut keyword = String::new(); + let mut c: u8 = reader.read(); + i += 1; + while c != 0 + { + keyword.push(c as char); + c = reader.read::<1, u8>(); + i += 1; + } + + let mut text_string = String::new(); + while i < length + { + text_string.push(reader.read::<1, u8>() as char); + i += 1; + } + + let mut keywords = HashMap::new(); + keywords.insert(keyword, text_string); + + Self + { + keywords + } + } +} +impl Mergable for tEXt { + fn merge(a: Self, b: Self) -> Self { + let mut keywords = HashMap::new(); + keywords.extend(a.keywords); + keywords.extend(b.keywords); + Self + { + keywords + } + } +} + +#[derive(Clone)] +struct IDAT +{ + data: Vec<u8> +} +impl PNGChunk for IDAT +{ + fn read(reader: &mut FileReader, length: usize) -> Self { + let data = reader.get_bytes(length); + Self + { + data + } + } + + const CHUNK_TYPE: [char; 4] = ['I', 'D', 'A', 'T']; +} +impl Mergable for IDAT +{ + fn merge(a: Self, b: Self) -> Self { + let mut data = Vec::new(); + data.extend(a.data); + data.extend(b.data); + Self + { + data + } + } +} + +impl<T: ColorType> Image<T> for PNGImage +{ + fn read_image(reader: &mut FileReader) -> Result<ImageBase<T>, String> { + let magic: [u8; 8] = reader.read_array(); + if magic != [0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A] + { + return Err("Not a PNG!".to_owned()); + } + + reader.set_endianness(false); + reader.set_checksum_handler(update_crc, 0xffffffffu32); + + let mut p = PNGImage::default(); + + while p.IEND.is_none() + { + let length: usize = reader.read::<4, u32>() as usize; + reader.reset_checksum(); + let chunk_type: [char; 4] = reader.read_array::<1, u8, 4>().map(|x| x as char); + let chunk_flags = chunk_type.map(|x| x.is_lowercase()); + println!("type: {chunk_type:?}"); + println!("length: {length}"); + match chunk_type + { + IHDR::CHUNK_TYPE => { p.IHDR = Option::<IHDR>::merge(p.IHDR, Some(IHDR::read(reader, length))) } + IDAT::CHUNK_TYPE => { p.IDAT = Option::<IDAT>::merge(p.IDAT, Some(IDAT::read(reader, length))) } + IEND::CHUNK_TYPE => { p.IEND = Option::<IEND>::merge(p.IEND, Some(IEND::read(reader, length))) } + tEXt::CHUNK_TYPE => { p.tEXt = Option::<tEXt>::merge(p.tEXt, Some(tEXt::read(reader, length))) } + _ => { + if chunk_flags[0] + { + println!("Skipping ancillary chunk"); + println!("{}", if chunk_flags[1] { "Private (can find definition)" } else { "Public" }) + } + else + { + println!("Skipping important chunk!!!"); + } + reader.skip(length) + } + } + let mut calculated_crc = reader.get_checksum(); + let crc: u32 = reader.read(); + if crc ^ calculated_crc != 0xffffffffu32 && !chunk_flags[0] + { + calculated_crc ^= 0xffffffffu32; + println!("Bad crc: {crc} {calculated_crc}"); + } + } + + if let Some(ref tEXt) = p.tEXt + { + for x in tEXt.keywords.iter() + { + println!("{0}: {1}", x.0, x.1); + } + } + + if p.IEND.is_none() + { + return Err("IEND not present".to_owned()); + } + + let Some(IHDR) = p.IHDR else { return Err("IHDR not present".to_owned()); }; + let Some(IDAT) = p.IDAT else { return Err("IDAT not present".to_owned()); }; + + let colors_per_pixel: u8 = if IHDR.color_type == 4 || IHDR.color_type == 6 { 4 } else { 3 }; + let bpp: u8 = IHDR.bit_depth * colors_per_pixel; + let width: u32 = IHDR.width; + let height: u32 = IHDR.height; + println!("{bpp}, {width}, {height}"); + + let compression_method: u8 = IDAT.data[0]; + let additional_flags: u8 = IDAT.data[1]; + + let decoded = zlib_decode(ByteBuffer::new(&IDAT.data[2..])); + + todo!("PNG not finished") + } + + fn write_image(image: &crate::image::ImageBase<T>, writer: &mut crate::writer::FileWriter)-> () { + todo!() + } +} diff --git a/src/reader.rs b/src/reader.rs index 6f18e0b..2871cae 100644 --- a/src/reader.rs +++ b/src/reader.rs @@ -4,7 +4,11 @@ use crate::byte_encode::ByteEncode; pub struct FileReader { - buf_reader: BufReader<File> + buf_reader: BufReader<File>, + le: bool, + checksum_func: Option<fn(&mut u32, &[u8]) -> ()>, + curr_checksum: u32, + checksum_start: u32, } impl FileReader { @@ -12,7 +16,11 @@ impl FileReader { { Self { - buf_reader + buf_reader, + le: true, + checksum_func: None, + curr_checksum: 0, + checksum_start: 0 } } @@ -25,10 +33,63 @@ impl FileReader { std::array::from_fn(|_| self.read()) } - pub fn read<const N: usize, T: ByteEncode<N>>(&mut self) -> T + pub fn get_bytes(&mut self, n: usize) -> Vec<u8> + { + let mut buff = vec![0u8; n]; + self.buf_reader.read_exact(&mut buff).expect("get_bytes failed to read"); + if let Some(ref f) = self.checksum_func + { + f(&mut self.curr_checksum, buff.as_slice()); + } + buff + } + + pub fn read_bytes<const N: usize>(&mut self) -> [u8; N] { let mut data = [0u8; N]; self.buf_reader.read_exact(&mut data).expect("Bad read"); - T::from_le_bytes(&data) + if let Some(ref f) = self.checksum_func + { + f(&mut self.curr_checksum, &data); + } + data + } + + pub fn read<const N: usize, T: ByteEncode<N>>(&mut self) -> T + { + let mut data = self.read_bytes::<N>(); + if self.le + { + T::from_le_bytes(&data) + } + else + { + T::from_be_bytes(&mut data) + } + } + + pub fn set_endianness(&mut self, le: bool) + { + self.le = le; + } + + pub fn set_checksum_handler(&mut self, checksum_func: fn(&mut u32, &[u8]) -> (), checksum_start: u32) + { + self.checksum_func = Some(checksum_func); + self.checksum_start = checksum_start; + self.curr_checksum = checksum_start; + + } + + pub fn reset_checksum(&mut self) -> u32 + { + let old_checksum = self.curr_checksum.clone(); + self.curr_checksum = self.checksum_start.clone(); + old_checksum + } + + pub fn get_checksum(&self) -> u32 + { + self.curr_checksum } } diff --git a/src/zlib.rs b/src/zlib.rs new file mode 100644 index 0000000..dcbdb55 --- /dev/null +++ b/src/zlib.rs @@ -0,0 +1,258 @@ +const MAX_BITS: usize = 15; +const CODE_TREE_LEN: usize = 288; +const DIST_TREE_LEN: usize = 32; + +const LEN_START: [usize; 29] = [3, 4, 5, 6, 7, 8, 9, 10, 11, 13, 15, 17, 19, 23, 27, 31, 35, 43, 51, 59, 67, 83, 99, 115, 131, 163, 195, 227, 258]; +const LEN_EXTRA: [usize; 29] = [0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3, 4, 4, 4, 4, 5, 5, 5, 5, 0]; + +const DIST_START: [usize; 30] = [1, 2, 3, 4, 5, 7, 9, 13, 17, 25, 33, 49, 65, 97, 129, 193, 257, 385, 513, 769, 1025, 1537, 2049, 3073, 4097, 6145, 8193, 12289, 16385, 24577]; +const DIST_EXTRA: [usize; 30] = [0, 0, 0, 0, 1, 1, 2, 2, 3, 3, 4, 4, 5, 5, 6, 6, 7, 7, 8, 8, 9, 9, 10, 10, 11, 11, 12, 12, 13, 13]; + +#[derive(Clone)] +struct HuffmanTree<const LENGTH: usize> +{ + pub len_counts: [u8; MAX_BITS + 1], + pub ordered_codes: [u16; LENGTH] +} + +impl<const LENGTH: usize> HuffmanTree<LENGTH> +{ + const fn new(code_lens: [u8; LENGTH]) -> Self + { + let mut len_counts = [0u8; MAX_BITS + 1]; + let mut ordered_codes = [0u16; LENGTH]; + + let mut code: usize = 0; + while code < LENGTH + { + len_counts[code_lens[code] as usize] += 1; + code += 1 + } + + let mut offsets = [0usize; MAX_BITS + 1]; + let mut len: usize = 1; + while len < MAX_BITS + { + offsets[len + 1] = offsets[len] + len_counts[len] as usize; + len += 1; + } + + let mut code = 0; + while code < LENGTH + { + if code_lens[code] != 0 + { + ordered_codes[offsets[code_lens[code] as usize]] = code as u16; + offsets[code_lens[code] as usize] += 1; + } + code += 1; + } + + Self + { + len_counts, + ordered_codes + } + } + + + + fn decode<T: Bitstream>(&self, bitstream: &mut T) -> u16 + { + let mut code: u16 = 0; + let mut first: u16 = 0; + let mut index: usize = 0; + + for len in 1..MAX_BITS + { + code |= bitstream.next_bit() as u16; + let count: u16 = self.len_counts[len] as u16; + if code - count < first + { + return self.ordered_codes[index + ((code - first) as usize)]; + } + index += count as usize; + first += count; + first <<= 1; + code <<= 1; + } + + panic!("No codes left") + } + + const fn get_static_code_lens() -> [u8; CODE_TREE_LEN] + { + let mut code_lens = [0u8; CODE_TREE_LEN]; + let mut i = 0; + while i < CODE_TREE_LEN + { + code_lens[i] = match i + { + 0..144 => 8, + 144..256 => 9, + 256..280 => 7, + 280..CODE_TREE_LEN => 8, + _ => panic!() + }; + i += 1 + } + code_lens + } + const fn get_static_dist_code_lens() -> [u8; DIST_TREE_LEN] + { + [5u8; DIST_TREE_LEN] + } +} + +static STATIC_HUFFMAN_TREE: HuffmanTree<CODE_TREE_LEN> = HuffmanTree::new(HuffmanTree::<CODE_TREE_LEN>::get_static_code_lens()); +static STATIC_HUFFMAN_DIST_TREE: HuffmanTree<DIST_TREE_LEN> = HuffmanTree::new(HuffmanTree::<DIST_TREE_LEN>::get_static_dist_code_lens()); + +#[derive(PartialEq)] +enum CompressionType +{ + None = 0, + StaticCodes = 1, + DynamicCodes = 2, + Reserved = 3, +} + +impl CompressionType +{ + fn new(x: u8) -> Self + { + match x + { + 0 => Self::None, + 1 => Self::StaticCodes, + 2 => Self::DynamicCodes, + 3 => Self::Reserved, + _ => panic!("Compression type should not be greater than 4") + } + } +} + +pub struct ByteBuffer<'a> +{ + buffer: &'a [u8], + pos: usize, +} + +impl<'a> ByteBuffer<'a> +{ + pub fn new(buffer: &'a [u8]) -> Self + { + Self + { + buffer, + pos: 0 + } + } +} + +pub trait Bitstream { + fn next_bit(&mut self) -> bool; + + fn next_bits(&mut self, c: usize) -> u16 + { + if c > 16 + { + panic!("Cannot fetch more than 16 bits at a time"); + } + let mut bits = 0u16; + + for i in 0..c + { + bits |= (self.next_bit() as u16) << i; + } + + bits + } + + fn skip_to_byte_boundary(&mut self); +} + +impl<'a> Bitstream for ByteBuffer<'a> +{ + fn next_bit(&mut self) -> bool + { + let bit = (self.buffer[self.pos/8 as usize] & 1 << self.pos%8) != 0; + self.pos += 1; + bit + } + + fn skip_to_byte_boundary(&mut self) + { + // println!("{}", (8 - self.pos%8)%8); + self.pos += (8 - self.pos%8)%8; + } +} + +pub fn zlib_decode<T: Bitstream>(mut instream: T) -> Vec<u8> +{ + let mut output_buffer: Vec<u8> = Vec::new(); + // let mut output_pos: usize = 0; + + let mut final_chunk = false; + while !final_chunk + { + final_chunk = instream.next_bit(); + let compression_type = CompressionType::new(instream.next_bits(2) as u8); + if compression_type == CompressionType::Reserved + { + panic!("Reserved compression type should not be used"); + } + else if compression_type == CompressionType::None + { + instream.skip_to_byte_boundary(); + let len = instream.next_bits(16); + let nlen = instream.next_bits(16); + if len != !nlen + { + panic!("len and nlen don't match"); + } + for _ in 0..len + { + output_buffer.push(instream.next_bits(8) as u8); + } + // todo!("No compression"); + } + else + { + let (code_tree, dist_tree) = + if compression_type == CompressionType::StaticCodes + { + (STATIC_HUFFMAN_TREE.clone(), STATIC_HUFFMAN_DIST_TREE.clone()) + } + else // Dynamic codes + { + todo!("Dynamic codes") + }; + + + let mut code = 0u16; + while code != 256 + { + code = code_tree.decode(&mut instream); + if code < 256 + { + output_buffer.push(code as u8); + } + else if code > 256 + { + let len = LEN_START[code as usize - 257] + instream.next_bits(LEN_EXTRA[code as usize - 257]) as usize; + + let dist_code = dist_tree.decode(&mut instream); + let dist = DIST_START[dist_code as usize] + instream.next_bits(DIST_EXTRA[dist_code as usize]) as usize; + + let out_pos = output_buffer.len() - 1; + for _ in 0..len + { + output_buffer.push(output_buffer[out_pos - dist]); + } + } + } + } + } + + output_buffer +} |