pile/crates/pile-flac/src/reader/reader.rs

use std::io::{Read, Seek, SeekFrom};

use crate::{
	FlacBlock, FlacDecodeError,
	blocks::{FlacAudioFrame, FlacMetablockHeader, FlacMetablockType},
};

// TODO: quickly skip blocks we do not need

/// The next block we expect to read
enum ReaderState {
	MagicBits,
	MetablockHeader { is_first: bool },
	MetaBlock { header: FlacMetablockHeader },
	AudioData,
	Done,
}

pub struct FlacReader<R: Read + Seek> {
	inner: R,
	state: ReaderState,
}

impl<R: Read + Seek> FlacReader<R> {
	const MIN_AUDIO_FRAME_LEN: usize = 5000;

	pub fn new(inner: R) -> Self {
		Self {
			inner,
			state: ReaderState::MagicBits,
		}
	}
}

impl<R: Read + Seek> Iterator for FlacReader<R> {
	type Item = Result<FlacBlock, FlacDecodeError>;

	fn next(&mut self) -> Option<Self::Item> {
		loop {
			match &mut self.state {
				ReaderState::Done => return None,

				ReaderState::MagicBits => {
					let mut data = [0u8; 4];
					if let Err(e) = self.inner.read_exact(&mut data[..4]) {
						self.state = ReaderState::Done;
						return Some(Err(e.into()));
					}

					if data != [0x66, 0x4C, 0x61, 0x43] {
						self.state = ReaderState::Done;
						return Some(Err(FlacDecodeError::BadMagicBytes));
					}

					self.state = ReaderState::MetablockHeader { is_first: true };
				}

				ReaderState::MetablockHeader { is_first } => {
					let mut data = [0u8; 4];
					if let Err(e) = self.inner.read_exact(&mut data[..]) {
						self.state = ReaderState::Done;
						return Some(Err(e.into()));
					}

					let header = match FlacMetablockHeader::decode(&data) {
						Ok(h) => h,
						Err(e) => {
							self.state = ReaderState::Done;
							return Some(Err(e));
						}
					};

					if *is_first && !matches!(header.block_type, FlacMetablockType::Streaminfo) {
						self.state = ReaderState::Done;
						return Some(Err(FlacDecodeError::BadFirstBlock));
					}

					self.state = ReaderState::MetaBlock { header };
				}

				ReaderState::MetaBlock { header } => {
					let mut data = vec![0u8; header.length as usize];
					if let Err(e) = self.inner.read_exact(&mut data) {
						self.state = ReaderState::Done;
						return Some(Err(e.into()));
					}

					let block = match FlacBlock::decode(header.block_type, &data) {
						Ok(b) => b,
						Err(e) => {
							self.state = ReaderState::Done;
							return Some(Err(e));
						}
					};

					if header.is_last {
						self.state = ReaderState::AudioData;
					} else {
						self.state = ReaderState::MetablockHeader { is_first: false };
					}

					return Some(Ok(block));
				}

				ReaderState::AudioData => {
					let mut data = Vec::new();
					loop {
						let mut byte = [0u8; 1];
						match self.inner.read_exact(&mut byte) {
							Ok(_) => {
								data.push(byte[0]);

								if data.len() >= Self::MIN_AUDIO_FRAME_LEN + 2 {
									let len = data.len();
									if data[len - 2] == 0b1111_1111
										&& data[len - 1] & 0b1111_1100 == 0b1111_1000
									{
										let frame_data = data[..len - 2].to_vec();

										if frame_data.len() < 2
											|| frame_data[0] != 0b1111_1111 || frame_data[1]
											& 0b1111_1100
											!= 0b1111_1000
										{
											self.state = ReaderState::Done;
											return Some(Err(FlacDecodeError::BadSyncBytes));
										}

										let audio_frame = match FlacAudioFrame::decode(&frame_data)
										{
											Ok(f) => f,
											Err(e) => {
												self.state = ReaderState::Done;
												return Some(Err(e));
											}
										};

										// Seek back 2 bytes so the next frame starts with the sync bytes
										if let Err(e) = self.inner.seek(SeekFrom::Current(-2)) {
											self.state = ReaderState::Done;
											return Some(Err(e.into()));
										}

										self.state = ReaderState::AudioData;

										return Some(Ok(FlacBlock::AudioFrame(audio_frame)));
									}
								}
							}
							Err(e) if e.kind() == std::io::ErrorKind::UnexpectedEof => {
								if data.len() > 2 {
									if data[0] != 0b1111_1111
										|| data[1] & 0b1111_1100 != 0b1111_1000
									{
										self.state = ReaderState::Done;
										return Some(Err(FlacDecodeError::BadSyncBytes));
									}

									let audio_frame = match FlacAudioFrame::decode(&data) {
										Ok(f) => f,
										Err(e) => {
											self.state = ReaderState::Done;
											return Some(Err(e));
										}
									};

									self.state = ReaderState::Done;
									return Some(Ok(FlacBlock::AudioFrame(audio_frame)));
								} else {
									self.state = ReaderState::Done;
									return None;
								}
							}
							Err(e) => {
								self.state = ReaderState::Done;
								return Some(Err(e.into()));
							}
						}
					}
				}
			}
		}
	}
}