Skip to main content

irox_bits/
allocimpls.rs

1// SPDX-License-Identifier: MIT
2// Copyright 2025 IROX Contributors
3//
4
5extern crate alloc;
6use crate::bits::Bits;
7use crate::error::Error;
8use crate::mutbits::MutBits;
9use crate::BitsWrapper;
10use alloc::boxed::Box;
11use alloc::collections::VecDeque;
12use alloc::string::String;
13use alloc::sync::Arc;
14use alloc::vec::Vec;
15use core::sync::atomic::{AtomicU64, Ordering};
16
17macro_rules! impl_bits_pop {
18    ($($ty:tt)+) => {
19        impl Bits for $($ty)+ {
20            fn next_u8(&mut self) -> Result<Option<u8>, Error> {
21                if self.is_empty() {
22                    return Ok(None)
23                }
24                Ok(Some(self.remove(0) as u8))
25            }
26
27            fn read_some_into<T: MutBits>(&mut self, into: &mut T) -> Result<usize, Error> {
28                Ok(into.write_some_bytes(self.as_ref()))
29            }
30            fn remaining(&self) -> Option<usize> {
31                Some(self.len())
32            }
33        }
34    };
35}
36
37impl_bits_pop!(String);
38impl_bits_pop!(&mut String);
39impl_bits_pop!(Vec<u8>);
40impl_bits_pop!(&mut Vec<u8>);
41macro_rules! impl_bits_vecdeque {
42    ($($ty:tt)+) => {
43        impl Bits for $($ty)+ {
44            fn next_u8(&mut self) -> Result<Option<u8>, Error> {
45                Ok(self.pop_front())
46            }
47
48            fn read_some_into<T: MutBits>(&mut self, into: &mut T) -> Result<usize, Error> {
49                let mut wrote = 0;
50                while let Some(val) = self.pop_front() {
51                    let Ok(()) = into.write_u8(val) else {
52                        return Ok(wrote);
53                    };
54                    wrote += 1;
55                }
56                Ok(wrote)
57            }
58            fn remaining(&self) -> Option<usize> {
59                Some(self.len())
60            }
61        }
62    };
63}
64impl_bits_vecdeque!(VecDeque<u8>);
65impl_bits_vecdeque!(&mut VecDeque<u8>);
66
67macro_rules! impl_mutbits_vecdeque {
68    ($($ty:tt)+) => {
69        impl MutBits for $($ty)+ {
70            fn write_u8(&mut self, val: u8) -> Result<(), Error> {
71                self.push_back(val);
72                Ok(())
73            }
74        }
75    };
76}
77impl_mutbits_vecdeque!(&mut VecDeque<u8>);
78impl_mutbits_vecdeque!(VecDeque<u8>);
79
80macro_rules! impl_push {
81    ($cast:ty, $($ty:tt)+) => {
82        impl MutBits for $($ty)+ {
83            fn write_u8(&mut self, val: u8) -> Result<(), Error> {
84                self.push(val as $cast);
85                Ok(())
86            }
87        }
88    };
89}
90impl_push!(char, &mut String);
91impl_push!(char, String);
92impl_push!(u8, Vec<u8>);
93impl_push!(u8, &mut Vec<u8>);
94
95impl<T: Bits> Bits for Box<T> {
96    fn next_u8(&mut self) -> Result<Option<u8>, Error> {
97        T::next_u8(self)
98    }
99}
100impl<T: MutBits> MutBits for Box<T> {
101    fn write_u8(&mut self, val: u8) -> Result<(), Error> {
102        T::write_u8(self, val)
103    }
104}
105
106///
107/// A struct to count the number of bytes moving through it.
108pub struct SharedCountingBits<'a, B> {
109    inner: BitsWrapper<'a, B>,
110    count: Arc<AtomicU64>,
111}
112
113impl<'a, B> SharedCountingBits<'a, B> {
114    pub fn new(inner: BitsWrapper<'a, B>) -> Self {
115        Self {
116            inner,
117            count: Arc::new(AtomicU64::new(0)),
118        }
119    }
120    /// Get a reference to the counter
121    pub fn get_count(&self) -> SharedROCounter {
122        SharedROCounter::new(self.count.clone())
123    }
124}
125impl<B: Bits> Bits for SharedCountingBits<'_, B> {
126    fn next_u8(&mut self) -> Result<Option<u8>, Error> {
127        let res = self.inner.next_u8();
128        if let Ok(Some(_)) = &res {
129            self.count.fetch_add(1, Ordering::Relaxed);
130        }
131        res
132    }
133}
134impl<B: MutBits> MutBits for SharedCountingBits<'_, B> {
135    fn write_u8(&mut self, val: u8) -> Result<(), Error> {
136        self.inner.write_u8(val)?;
137        self.count.fetch_add(1, Ordering::Relaxed);
138        Ok(())
139    }
140}
141
142///
143/// A Read-Only counter that can be shared between threads.  The owner of the underlying counter
144/// is free to update the value, but users of this object alone may not.
145#[derive(Debug, Clone)]
146pub struct SharedROCounter {
147    counter: Arc<AtomicU64>,
148}
149
150impl SharedROCounter {
151    pub fn new(counter: Arc<AtomicU64>) -> Self {
152        SharedROCounter { counter }
153    }
154
155    ///
156    /// Returns the current value of the counter
157    pub fn get_count(&self) -> u64 {
158        self.counter.load(Ordering::Relaxed)
159    }
160}