use std::io::{Result, Write};
pub struct Counter<W> {
inner: W,
count: usize,
}
impl<W> Counter<W> {
pub fn new(inner: W) -> Self {
Self { inner, count: 0 }
}
pub fn count(&self) -> usize {
self.count
}
pub fn into_inner(self) -> W {
self.inner
}
}
impl<W: Write> Write for Counter<W> {
fn write(&mut self, buf: &[u8]) -> Result<usize> {
let count = self.inner.write(buf)?;
self.count = self.count.saturating_add(count);
Ok(count)
}
fn flush(&mut self) -> Result<()> {
self.inner.flush()
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::io::Cursor;
#[test]
fn encode_header() -> Result<()> {
let data = "hello world".as_bytes();
let buf = Cursor::new(vec![]);
let mut counter = Counter::new(buf);
counter.write_all(data)?;
assert_eq!(counter.count(), data.len());
assert_eq!(counter.into_inner().into_inner(), data);
Ok(())
}
}