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        }
31    };
32}
33
34impl_bits_pop!(String);
35impl_bits_pop!(&mut String);
36impl_bits_pop!(Vec<u8>);
37impl_bits_pop!(&mut Vec<u8>);
38macro_rules! impl_bits_vecdeque {
39    ($($ty:tt)+) => {
40        impl Bits for $($ty)+ {
41            fn next_u8(&mut self) -> Result<Option<u8>, Error> {
42                Ok(self.pop_front())
43            }
44
45            fn read_some_into<T: MutBits>(&mut self, into: &mut T) -> Result<usize, Error> {
46                let mut wrote = 0;
47                while let Some(val) = self.pop_front() {
48                    let Ok(()) = into.write_u8(val) else {
49                        return Ok(wrote);
50                    };
51                    wrote += 1;
52                }
53                Ok(wrote)
54            }
55        }
56    };
57}
58impl_bits_vecdeque!(VecDeque<u8>);
59impl_bits_vecdeque!(&mut VecDeque<u8>);
60
61macro_rules! impl_mutbits_vecdeque {
62    ($($ty:tt)+) => {
63        impl MutBits for $($ty)+ {
64            fn write_u8(&mut self, val: u8) -> Result<(), Error> {
65                self.push_back(val);
66                Ok(())
67            }
68        }
69    };
70}
71impl_mutbits_vecdeque!(&mut VecDeque<u8>);
72impl_mutbits_vecdeque!(VecDeque<u8>);
73
74macro_rules! impl_push {
75    ($cast:ty, $($ty:tt)+) => {
76        impl MutBits for $($ty)+ {
77            fn write_u8(&mut self, val: u8) -> Result<(), Error> {
78                self.push(val as $cast);
79                Ok(())
80            }
81        }
82    };
83}
84impl_push!(char, &mut String);
85impl_push!(char, String);
86impl_push!(u8, Vec<u8>);
87impl_push!(u8, &mut Vec<u8>);
88
89impl<T: Bits> Bits for Box<T> {
90    fn next_u8(&mut self) -> Result<Option<u8>, Error> {
91        T::next_u8(self)
92    }
93}
94impl<T: MutBits> MutBits for Box<T> {
95    fn write_u8(&mut self, val: u8) -> Result<(), Error> {
96        T::write_u8(self, val)
97    }
98}
99
100///
101/// A struct to count the number of bytes moving through it.
102pub struct SharedCountingBits<'a, B> {
103    inner: BitsWrapper<'a, B>,
104    count: Arc<AtomicU64>,
105}
106
107impl<'a, B> SharedCountingBits<'a, B> {
108    pub fn new(inner: BitsWrapper<'a, B>) -> Self {
109        Self {
110            inner,
111            count: Arc::new(AtomicU64::new(0)),
112        }
113    }
114    /// Get a reference to the counter
115    pub fn get_count(&self) -> SharedROCounter {
116        SharedROCounter::new(self.count.clone())
117    }
118}
119impl<B: Bits> Bits for SharedCountingBits<'_, B> {
120    fn next_u8(&mut self) -> Result<Option<u8>, Error> {
121        let res = self.inner.next_u8();
122        if let Ok(Some(_)) = &res {
123            self.count.fetch_add(1, Ordering::Relaxed);
124        }
125        res
126    }
127}
128impl<B: MutBits> MutBits for SharedCountingBits<'_, B> {
129    fn write_u8(&mut self, val: u8) -> Result<(), Error> {
130        self.inner.write_u8(val)?;
131        self.count.fetch_add(1, Ordering::Relaxed);
132        Ok(())
133    }
134}
135
136///
137/// A Read-Only counter that can be shared between threads.  The owner of the underlying counter
138/// is free to update the value, but users of this object alone may not.
139#[derive(Debug, Clone)]
140pub struct SharedROCounter {
141    counter: Arc<AtomicU64>,
142}
143
144impl SharedROCounter {
145    pub fn new(counter: Arc<AtomicU64>) -> Self {
146        SharedROCounter { counter }
147    }
148
149    ///
150    /// Returns the current value of the counter
151    pub fn get_count(&self) -> u64 {
152        self.counter.load(Ordering::Relaxed)
153    }
154}