use codec::Output;
pub struct CountedOutput<'a, O: ?Sized> {
output: &'a mut O,
num_bytes_written: usize,
}
impl<'a, O: ?Sized> CountedOutput<'a, O> {
pub fn new(output: &'a mut O) -> Self {
Self { output, num_bytes_written: 0 }
}
pub fn num_bytes_written(&self) -> usize {
self.num_bytes_written
}
}
impl<O: Output + ?Sized> Output for CountedOutput<'_, O> {
fn write(&mut self, bytes: &[u8]) {
self.output.write(bytes);
self.num_bytes_written += bytes.len();
}
fn push_byte(&mut self, byte: u8) {
self.output.push_byte(byte);
self.num_bytes_written += 1;
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn basic_operation() {
let mut out = Vec::new();
let mut counted = CountedOutput::new(&mut out);
assert_eq!(counted.num_bytes_written(), 0);
counted.write(&[1, 2, 3]);
assert_eq!(counted.num_bytes_written(), 3);
counted.push_byte(4);
assert_eq!(counted.num_bytes_written(), 4);
counted.write(&[]);
assert_eq!(counted.num_bytes_written(), 4);
assert_eq!(out, vec![1, 2, 3, 4]);
}
}