irox_tools/read/
counting.rs

1// SPDX-License-Identifier: MIT
2// Copyright 2025 IROX Contributors
3//
4
5extern crate alloc;
6
7use alloc::sync::Arc;
8use core::sync::atomic::{AtomicU64, Ordering};
9use irox_bits::SharedROCounter;
10use std::io::Read;
11
12///
13/// Wraps the reader, providing a convenient way to count the number of bytes read from the
14/// underlying impl.
15pub struct ReadCounting<T: Read> {
16    reader: T,
17    count: u64,
18}
19
20impl<T: Read> ReadCounting<T> {
21    pub fn new(reader: T) -> Self {
22        ReadCounting { reader, count: 0 }
23    }
24
25    /// Returns the number of bytes read to this point.
26    pub fn get_read_count(&self) -> u64 {
27        self.count
28    }
29}
30impl<T: Read> Read for ReadCounting<T> {
31    fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
32        let read = self.reader.read(buf)?;
33        self.count += read as u64;
34        Ok(read)
35    }
36}
37
38///
39/// A [`ReadCounting`] except can be shared between threads
40pub struct SharedReadCounting<T: Read> {
41    reader: T,
42    counter: Arc<AtomicU64>,
43}
44
45impl<T: Read> SharedReadCounting<T> {
46    pub fn new(reader: T) -> Self {
47        SharedReadCounting {
48            reader,
49            counter: Arc::new(AtomicU64::new(0)),
50        }
51    }
52
53    /// Returns a shared read-only copy of the counter
54    pub fn get_counter(&self) -> SharedROCounter {
55        SharedROCounter::new(self.counter.clone())
56    }
57}
58
59impl<T: Read> Read for SharedReadCounting<T> {
60    fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
61        let read = self.reader.read(buf)?;
62        self.counter.fetch_add(read as u64, Ordering::Relaxed);
63        Ok(read)
64    }
65}
irox_tools/read/counting.rs

irox_tools/read/
counting.rs
