initial commit

1 files changed, 328 insertions, 0 deletions

diff --git a/src/lib.rs b/src/lib.rs
new file mode 100644
index 0000000..a801f9a
--- /dev/null
+++ b/src/lib.rs
@@ -0,0 +1,328 @@
+//#![cfg_attr(not(feature = "std"), no_std)]
+
+use core::cmp::min;
+use core::fmt::{self, Debug, Display};
+
+pub trait IOError: Sized + Display + Debug {
+}
+impl<T: Sized + Display + Debug> IOError for T {
+}
+
+#[cfg(feature = "std")]
+impl IOError for std::io::Error {
+	fn unexpected_eof() -> Self {
+		std::io::ErrorKind::UnexpectedEof.into()
+	}
+}
+
+#[derive(Debug)]
+#[non_exhaustive]
+pub enum Error<I: IOError> {
+	IO(I),
+	NotPng,
+	BadIhdr,
+	UnrecognizedChunk([u8; 4]),
+	BadBlockType,
+	TooMuchData,
+	UnexpectedEob,
+	BadZlibHeader,
+}
+
+impl<I: IOError> From<I> for Error<I> {
+	fn from(value: I) -> Self {
+		Self::IO(value)
+	}
+}
+
+impl<I: IOError> Display for Error<I> {
+	fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+		match self {
+			Self::IO(e) => write!(f, "{e}"),
+			Self::NotPng => write!(f, "not a png file"),
+			Self::BadIhdr => write!(f, "bad IHDR chunk"),
+			Self::UnrecognizedChunk([a, b, c, d]) => write!(f, "unrecognized chunk type: {a} {b} {c} {d}"),
+			Self::BadBlockType => write!(f, "bad DEFLATE block type"),
+			Self::TooMuchData => write!(f, "decompressed data is larger than it should be"),
+			Self::UnexpectedEob => write!(f, "unexpected end of block"),
+			Self::BadZlibHeader => write!(f, "bad zlib header"),
+		}
+	}
+}
+
+#[cfg(feature = "std")]
+impl<I: IOError> std::error::Error for Error<I> {}
+
+pub trait Read {
+	type Error: IOError;
+	fn read(&mut self, buf: &mut [u8]) -> Result<(), Self::Error>;	
+	fn skip_bytes(&mut self, count: usize) -> Result<(), Self::Error> {
+		let mut count = count;
+		const BUF_LEN: usize = 128;
+		let mut buf = [0; BUF_LEN];
+		while count > 0 {
+			let c = min(BUF_LEN, count);
+			self.read(&mut buf[..c])?;
+			count -= c;
+		}
+		Ok(())
+	}
+}
+
+#[derive(Debug)]
+pub struct UnexpectedEofError;
+
+impl core::fmt::Display for UnexpectedEofError {
+	fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+		write!(f, "unexpected EOF")
+	}
+}
+
+impl<'a> Read for &'a [u8] {
+	type Error = UnexpectedEofError;
+	fn read(&mut self, buf: &mut [u8]) -> Result<(), Self::Error> {
+		if self.len() < buf.len() {
+			return Err(UnexpectedEofError);
+		}
+		buf.copy_from_slice(&self[..buf.len()]);
+		*self = &self[buf.len()..];
+		Ok(())
+	}
+	// TODO: skip_bytes implementation
+}
+
+struct BlockReader<'a, R: Read> {
+	inner: &'a mut R,
+	bytes_left: usize
+}
+
+impl<'a, R: Read> BlockReader<'_, R> {
+	fn read(&mut self, buf: &mut [u8]) -> Result<(), Error<R::Error>> {
+		if buf.len() > self.bytes_left {
+			return Err(Error::UnexpectedEob);
+		}
+		Ok(self.inner.read(buf)?)
+	}
+	
+	fn read_partial(&mut self, buf: &mut [u8]) -> Result<usize, Error<R::Error>> {
+		let count = min(self.bytes_left, buf.len());
+		self.read(&mut buf[..count])?;
+		Ok(count)
+	}
+}
+
+#[derive(Debug)]
+pub struct ImageHeader {
+	width: u32,
+	height: u32,
+	bit_depth: BitDepth,
+	color_type: ColorType,
+}
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
+pub enum BitDepth {
+	One,
+	Two,
+	Four,
+	Eight,
+	Sixteen,
+}
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
+pub enum ColorType {
+	Gray,
+	GrayAlpha,
+	Rgb,
+	Rgba,
+	Indexed,
+}
+
+impl BitDepth {
+	fn from_byte(x: u8) -> Option<Self> {
+		Some(match x {
+			1 => Self::One,
+			2 => Self::Two,
+			4 => Self::Four,
+			8 => Self::Eight,
+			16 => Self::Sixteen,
+			_ => return None
+		})
+	}
+}
+
+impl ColorType {
+	fn from_byte(x: u8) -> Option<Self> {
+		Some(match x {
+			0 => Self::Gray,
+			2 => Self::Rgb,
+			3 => Self::Indexed,
+			4 => Self::GrayAlpha,
+			6 => Self::Rgba,
+			_ => return None,
+		})
+	}
+}
+
+impl ImageHeader {
+	pub fn width(&self) -> u32 {
+		self.width
+	}
+	pub fn height(&self) -> u32 {
+		self.height
+	}
+	pub fn bit_depth(&self) -> BitDepth {
+		self.bit_depth
+	}
+	pub fn color_type(&self) -> ColorType {
+		self.color_type
+	}
+}
+
+/// number of bits to read in each [`Read::read`] call.
+///
+/// don't change this to something bigger than `u32`, since we don't want to overread past the zlib checksum.
+type ReadBits = u32;
+/// number of bits to store in the [`BitReader`] buffer.
+type Bits = u64;
+
+struct BitReader<'a, R: Read> {
+	inner: BlockReader<'a, R>,
+	bits: Bits,
+	bits_left: u8,
+}
+
+impl<'a, R: Read> From<BlockReader<'a, R>> for BitReader<'a, R> {
+	fn from(inner: BlockReader<'a, R>) -> Self {
+		Self {
+			inner,
+			bits: 0,
+			bits_left: 0,
+		}
+	}
+}
+
+impl<R: Read> BitReader<'_, R> {
+	fn peek_bits(&mut self, count: u8) -> Result<u32, Error<R::Error>> {
+		debug_assert!(count > 0 && u32::from(count) <= 31);
+		if self.bits_left < count {
+			// read more bits
+			let mut new_bits = [0; ReadBits::BITS as usize / 8];
+			self.inner.read_partial(&mut new_bits)?;
+			let new_bits = Bits::from(ReadBits::from_le_bytes(new_bits));
+			self.bits |= new_bits << self.bits_left;
+			self.bits_left += ReadBits::BITS as u8;
+		}
+		Ok((self.bits as u32) & ((1 << count) - 1))
+	}
+	
+	fn read_bits(&mut self, count: u8) -> Result<u32, Error<R::Error>> {
+		let bits = self.peek_bits(count)?;
+		self.bits_left -= count;
+		self.bits >>= count;
+		Ok(bits)
+	}
+}
+
+
+pub fn read_png_header<R: Read>(reader: &mut R) -> Result<ImageHeader, Error<R::Error>> {
+	let mut signature = [0; 8];
+	reader.read(&mut signature)?;
+	if signature != [137, 80, 78, 71, 13, 10, 26, 10] {
+		return Err(Error::NotPng);
+	}
+	let mut ihdr = [0; 25];
+	reader.read(&mut ihdr)?;
+	let ihdr_len = (u32::from_be_bytes([ihdr[0], ihdr[1], ihdr[2], ihdr[3]]) + 12) as usize;
+	if &ihdr[4..8] != b"IHDR" || ihdr_len < ihdr.len() {
+		return Err(Error::BadIhdr);
+	}
+	reader.skip_bytes(ihdr_len - ihdr.len())?;
+
+	let width = u32::from_be_bytes([ihdr[8], ihdr[9], ihdr[10], ihdr[11]]);
+	let height = u32::from_be_bytes([ihdr[12], ihdr[13], ihdr[14], ihdr[15]]);
+	let bit_depth = BitDepth::from_byte(ihdr[16]).ok_or(Error::BadIhdr)?;
+	let color_type = ColorType::from_byte(ihdr[17]).ok_or(Error::BadIhdr)?;
+	let compression = ihdr[18];
+	let filter = ihdr[19];
+	if compression != 0 || filter != 0 {
+		return Err(Error::BadIhdr);
+	}
+	
+	Ok(ImageHeader { width, height, bit_depth, color_type })
+}
+
+struct ByteWriter<'a> {
+	slice: &'a mut [u8],
+}
+
+fn read_idat<R: Read>(mut reader: BlockReader<'_, R>, writer: &mut ByteWriter) -> Result<(), Error<R::Error>> {
+	let mut zlib_header = [0; 2];
+	reader.read(&mut zlib_header)?;
+	
+	let mut reader = BitReader::from(reader);
+	loop {
+		let bfinal = reader.read_bits(1)?;
+		let btype = reader.read_bits(2)?;
+		if btype == 0 {
+			// uncompressed block
+			let len = reader.read_bits(16)? as usize;
+			reader.read_bits(16)?; // nlen
+			if writer.slice.len() < len {
+				return Err(Error::TooMuchData);
+			}
+			reader.inner.read(&mut writer.slice[..len])?;
+		} else if btype == 1 || btype == 2 {
+			// compressed block
+			todo!("{btype}")
+		} else {
+			// 0b11 is not a valid block type
+			return Err(Error::BadBlockType);
+		}
+		
+		if bfinal != 0 {
+			break;
+		}
+	}
+	
+	// skip checksum
+	let mut trash = [0; 4];
+	reader.inner.read(&mut trash[..4 - usize::from(reader.bits_left) / 8])?;
+	
+	Ok(())
+}
+
+pub fn read_png<R: Read>(
+	reader: &mut R,
+	header: Option<ImageHeader>,
+	buf: &mut [u8],
+) -> Result<(), Error<R::Error>> {
+	let _header = match header {
+		None => read_png_header(reader)?,
+		Some(h) => h,
+	};
+	let mut writer = ByteWriter { slice: buf };
+	loop {
+		let mut chunk_header = [0; 8];
+		reader.read(&mut chunk_header[..])?;
+		let chunk_len = u32::from_be_bytes([chunk_header[0], chunk_header[1], chunk_header[2], chunk_header[3]]) as usize;
+		let chunk_type = [chunk_header[4], chunk_header[5], chunk_header[6], chunk_header[7]];
+		if &chunk_type == b"IEND" {
+			break;
+		} else if &chunk_type == b"IDAT" {
+			read_idat(BlockReader {
+				inner: reader,
+				bytes_left: chunk_len
+			}, &mut writer)?;
+		} else if &chunk_type == b"PLTE" {
+			todo!();
+		} else if chunk_type[0].is_ascii_lowercase() {
+			// non-essential chunk
+		} else {
+			return Err(Error::UnrecognizedChunk(chunk_type));
+		}
+		
+		reader.skip_bytes(chunk_len + 4)?;
+	}
+	
+	Ok(())
+}
+