Skip to main content

jam_std_common/util/
counted_output.rs

1use codec::Output;
2
3/// [`Output`] wrapper which counts the number of bytes written.
4pub struct CountedOutput<'a, O: ?Sized> {
5	output: &'a mut O,
6	num_bytes_written: usize,
7}
8
9impl<'a, O: ?Sized> CountedOutput<'a, O> {
10	pub fn new(output: &'a mut O) -> Self {
11		Self { output, num_bytes_written: 0 }
12	}
13
14	pub fn num_bytes_written(&self) -> usize {
15		self.num_bytes_written
16	}
17}
18
19impl<O: Output + ?Sized> Output for CountedOutput<'_, O> {
20	fn write(&mut self, bytes: &[u8]) {
21		self.output.write(bytes);
22		self.num_bytes_written += bytes.len();
23	}
24
25	fn push_byte(&mut self, byte: u8) {
26		self.output.push_byte(byte);
27		self.num_bytes_written += 1;
28	}
29}
30
31#[cfg(test)]
32mod tests {
33	use super::*;
34
35	#[test]
36	fn basic_operation() {
37		let mut out = Vec::new();
38		let mut counted = CountedOutput::new(&mut out);
39		assert_eq!(counted.num_bytes_written(), 0);
40
41		counted.write(&[1, 2, 3]);
42		assert_eq!(counted.num_bytes_written(), 3);
43
44		counted.push_byte(4);
45		assert_eq!(counted.num_bytes_written(), 4);
46
47		counted.write(&[]);
48		assert_eq!(counted.num_bytes_written(), 4);
49
50		// The bytes are forwarded to the wrapped output unchanged
51		assert_eq!(out, vec![1, 2, 3, 4]);
52	}
53}