jam-std-common 0.1.29

use codec::{Error, Input};

/// [`Input`] wrapper which records the bytes read.
pub struct RecordedInput<'a, I, R = Vec<u8>> {
	input: &'a mut I,
	record: R,
	failed: bool,
}

impl<'a, I, R: Default> RecordedInput<'a, I, R> {
	pub fn new(input: &'a mut I) -> Self {
		Self::with_record(input, Default::default())
	}
}

impl<'a, I, R> RecordedInput<'a, I, R> {
	pub fn with_record(input: &'a mut I, record: R) -> Self {
		Self { input, record, failed: false }
	}

	pub fn record(&self) -> &R {
		&self.record
	}

	pub fn into_record(self) -> R {
		self.record
	}

	/// Returns `true` if the wrapped `Input` has reported an error.
	///
	/// In this case, the record may be incomplete.
	pub fn failed(&self) -> bool {
		self.failed
	}
}

impl<'a, I: Input, R: Extend<u8>> Input for RecordedInput<'a, I, R> {
	fn remaining_len(&mut self) -> Result<Option<usize>, Error> {
		self.input.remaining_len()
	}

	fn read(&mut self, into: &mut [u8]) -> Result<(), Error> {
		let res = self.input.read(into);
		match res {
			Ok(()) => self.record.extend(into.iter().copied()),
			Err(_) => self.failed = true,
		}
		res
	}

	fn read_byte(&mut self) -> Result<u8, Error> {
		let res = self.input.read_byte();
		match res {
			Ok(byte) => self.record.extend(std::iter::once(byte)),
			Err(_) => self.failed = true,
		}
		res
	}

	fn descend_ref(&mut self) -> Result<(), Error> {
		self.input.descend_ref()
	}

	fn ascend_ref(&mut self) {
		self.input.ascend_ref();
	}

	fn on_before_alloc_mem(&mut self, size: usize) -> Result<(), Error> {
		self.input.on_before_alloc_mem(size)
	}
}

#[cfg(test)]
mod tests {
	use super::*;

	#[test]
	fn basic_operation() {
		let mut input = [1, 2, 3, 4].as_slice();
		let mut recorded = RecordedInput::<_, Vec<_>>::new(&mut input);

		let mut buf = [0; 3];
		recorded.read(&mut buf).unwrap();
		assert_eq!(buf, [1, 2, 3]);
		assert_eq!(recorded.read_byte().unwrap(), 4);
		assert!(!recorded.failed());
		assert_eq!(recorded.record(), &[1, 2, 3, 4]);

		assert!(recorded.read_byte().is_err());
		assert!(recorded.failed());
	}
}