Skip to main content

jam_std_common/util/
recorded_input.rs

1use codec::{Error, Input};
2
3/// [`Input`] wrapper which records the bytes read.
4pub struct RecordedInput<'a, I, R = Vec<u8>> {
5	input: &'a mut I,
6	record: R,
7	failed: bool,
8}
9
10impl<'a, I, R: Default> RecordedInput<'a, I, R> {
11	pub fn new(input: &'a mut I) -> Self {
12		Self::with_record(input, Default::default())
13	}
14}
15
16impl<'a, I, R> RecordedInput<'a, I, R> {
17	pub fn with_record(input: &'a mut I, record: R) -> Self {
18		Self { input, record, failed: false }
19	}
20
21	pub fn record(&self) -> &R {
22		&self.record
23	}
24
25	pub fn into_record(self) -> R {
26		self.record
27	}
28
29	/// Returns `true` if the wrapped `Input` has reported an error.
30	///
31	/// In this case, the record may be incomplete.
32	pub fn failed(&self) -> bool {
33		self.failed
34	}
35}
36
37impl<'a, I: Input, R: Extend<u8>> Input for RecordedInput<'a, I, R> {
38	fn remaining_len(&mut self) -> Result<Option<usize>, Error> {
39		self.input.remaining_len()
40	}
41
42	fn read(&mut self, into: &mut [u8]) -> Result<(), Error> {
43		let res = self.input.read(into);
44		match res {
45			Ok(()) => self.record.extend(into.iter().copied()),
46			Err(_) => self.failed = true,
47		}
48		res
49	}
50
51	fn read_byte(&mut self) -> Result<u8, Error> {
52		let res = self.input.read_byte();
53		match res {
54			Ok(byte) => self.record.extend(std::iter::once(byte)),
55			Err(_) => self.failed = true,
56		}
57		res
58	}
59
60	fn descend_ref(&mut self) -> Result<(), Error> {
61		self.input.descend_ref()
62	}
63
64	fn ascend_ref(&mut self) {
65		self.input.ascend_ref();
66	}
67
68	fn on_before_alloc_mem(&mut self, size: usize) -> Result<(), Error> {
69		self.input.on_before_alloc_mem(size)
70	}
71}
72
73#[cfg(test)]
74mod tests {
75	use super::*;
76
77	#[test]
78	fn basic_operation() {
79		let mut input = [1, 2, 3, 4].as_slice();
80		let mut recorded = RecordedInput::<_, Vec<_>>::new(&mut input);
81
82		let mut buf = [0; 3];
83		recorded.read(&mut buf).unwrap();
84		assert_eq!(buf, [1, 2, 3]);
85		assert_eq!(recorded.read_byte().unwrap(), 4);
86		assert!(!recorded.failed());
87		assert_eq!(recorded.record(), &[1, 2, 3, 4]);
88
89		assert!(recorded.read_byte().is_err());
90		assert!(recorded.failed());
91	}
92}