flatty_base/utils/
iter.rs

1use crate::{
2    error::{Error, ErrorKind},
3    traits::*,
4    utils::{ceil_mul, max},
5};
6use core::marker::PhantomData;
7
8pub trait TypeIter: Clone {
9    type Item: Flat + ?Sized;
10    fn align(&self) -> usize;
11    fn min_size(&self, pos: usize) -> usize;
12
13    fn check_align_and_min_size(&self, data: &[u8]) -> Result<(), Error> {
14        if data.as_ptr().align_offset(self.align()) != 0 {
15            Err(Error {
16                kind: ErrorKind::BadAlign,
17                pos: 0,
18            })
19        } else if data.len() < self.min_size(0) {
20            Err(Error {
21                kind: ErrorKind::InsufficientSize,
22                pos: 0,
23            })
24        } else {
25            Ok(())
26        }
27    }
28}
29
30#[derive(Debug)]
31pub struct SingleType<T: Flat + ?Sized> {
32    _phantom: PhantomData<T>,
33}
34impl<T: Flat + ?Sized> Clone for SingleType<T> {
35    fn clone(&self) -> Self {
36        Self { _phantom: PhantomData }
37    }
38}
39impl<T: Flat + ?Sized> SingleType<T> {
40    #[allow(clippy::new_without_default)]
41    pub fn new() -> Self {
42        Self { _phantom: PhantomData }
43    }
44}
45impl<T: Flat + ?Sized> TypeIter for SingleType<T> {
46    type Item = T;
47    fn align(&self) -> usize {
48        T::ALIGN
49    }
50    fn min_size(&self, pos: usize) -> usize {
51        ceil_mul(pos, T::ALIGN) + T::MIN_SIZE
52    }
53}
54
55#[derive(Debug)]
56pub struct TwoOrMoreTypes<T: Flat + Sized, I: TypeIter> {
57    _phantom: PhantomData<T>,
58    next: I,
59}
60impl<T: Flat + Sized, I: TypeIter> Clone for TwoOrMoreTypes<T, I> {
61    fn clone(&self) -> Self {
62        Self {
63            _phantom: PhantomData,
64            next: self.next.clone(),
65        }
66    }
67}
68impl<T: Flat + Sized, I: TypeIter> TwoOrMoreTypes<T, I> {
69    pub fn new(next: I) -> Self {
70        Self {
71            _phantom: PhantomData,
72            next,
73        }
74    }
75    pub fn next(self) -> I {
76        self.next
77    }
78}
79impl<T: Flat + Sized, I: TypeIter> TypeIter for TwoOrMoreTypes<T, I> {
80    type Item = T;
81    fn align(&self) -> usize {
82        max(T::ALIGN, self.next.align())
83    }
84    fn min_size(&self, pos: usize) -> usize {
85        self.next.min_size(ceil_mul(pos, T::ALIGN) + T::SIZE)
86    }
87}
88
89#[derive(Clone, Debug)]
90pub struct PosIter<I: TypeIter> {
91    pos: usize,
92    iter: I,
93}
94impl<I: TypeIter> PosIter<I> {
95    pub fn new(iter: I) -> Self {
96        Self { pos: 0, iter }
97    }
98    pub fn pos(&self) -> usize {
99        self.pos
100    }
101}
102impl<T: Flat + Sized, I: TypeIter> PosIter<TwoOrMoreTypes<T, I>> {
103    pub fn next(self) -> PosIter<I> {
104        PosIter {
105            pos: ceil_mul(self.pos + T::SIZE, I::Item::ALIGN),
106            iter: self.iter.next(),
107        }
108    }
109}
110impl<T: Flat + ?Sized> PosIter<SingleType<T>> {
111    pub fn assert_last(&self) {
112        // Nothing to do here, this is a static assert.
113    }
114}
115
116pub trait Data<'a>: Sized + 'a {
117    fn bytes(&self) -> &'_ [u8];
118    type Output<T: Flat + ?Sized + 'a>: Sized;
119    fn split(self, pos: usize) -> (Self, Self);
120    fn value<T: Flat + ?Sized + 'a>(self) -> Self::Output<T>;
121}
122
123impl<'a> Data<'a> for &'a [u8] {
124    fn bytes(&self) -> &'_ [u8] {
125        self
126    }
127    type Output<T: Flat + ?Sized + 'a> = &'a [u8];
128    fn split(self, pos: usize) -> (Self, Self) {
129        self.split_at(pos)
130    }
131    fn value<T: Flat + ?Sized + 'a>(self) -> Self::Output<T> {
132        self
133    }
134}
135
136impl<'a> Data<'a> for &'a mut [u8] {
137    fn bytes(&self) -> &'_ [u8] {
138        self
139    }
140    type Output<T: Flat + ?Sized + 'a> = &'a mut [u8];
141    fn split(self, pos: usize) -> (Self, Self) {
142        self.split_at_mut(pos)
143    }
144    fn value<T: Flat + ?Sized + 'a>(self) -> Self::Output<T> {
145        self
146    }
147}
148
149#[derive(Clone, Debug)]
150pub struct RefData<'a>(&'a [u8]);
151impl<'a> RefData<'a> {
152    pub unsafe fn new(data: &'a [u8]) -> Self {
153        Self(data)
154    }
155}
156impl<'a> Data<'a> for RefData<'a> {
157    fn bytes(&self) -> &'_ [u8] {
158        self.0.bytes()
159    }
160    type Output<T: Flat + ?Sized + 'a> = &'a T;
161    fn split(self, pos: usize) -> (Self, Self) {
162        let (a, b) = Data::split(self.0, pos);
163        (Self(a), Self(b))
164    }
165    fn value<T: Flat + ?Sized + 'a>(self) -> Self::Output<T> {
166        unsafe { T::from_bytes_unchecked(self.0.value::<T>()) }
167    }
168}
169
170#[derive(Debug)]
171pub struct MutData<'a>(&'a mut [u8]);
172impl<'a> MutData<'a> {
173    pub unsafe fn new(data: &'a mut [u8]) -> Self {
174        Self(data)
175    }
176}
177impl<'a> Data<'a> for MutData<'a> {
178    fn bytes(&self) -> &'_ [u8] {
179        self.0.bytes()
180    }
181    type Output<T: Flat + ?Sized + 'a> = &'a mut T;
182    fn split(self, pos: usize) -> (Self, Self) {
183        let (a, b) = Data::split(self.0, pos);
184        (Self(a), Self(b))
185    }
186    fn value<T: Flat + ?Sized + 'a>(self) -> Self::Output<T> {
187        unsafe { T::from_mut_bytes_unchecked(self.0.value::<T>()) }
188    }
189}
190
191#[derive(Debug)]
192pub struct DataIter<'a, D: Data<'a>, I: TypeIter> {
193    _ghost: PhantomData<&'a ()>,
194    data: D,
195    iter: PosIter<I>,
196}
197impl<'a, D: Data<'a> + Clone, I: TypeIter> Clone for DataIter<'a, D, I> {
198    fn clone(&self) -> Self {
199        Self {
200            _ghost: self._ghost,
201            data: self.data.clone(),
202            iter: self.iter.clone(),
203        }
204    }
205}
206impl<'a, D: Data<'a>, I: TypeIter> DataIter<'a, D, I> {
207    pub fn new(data: D, iter: I) -> Result<Self, Error> {
208        iter.check_align_and_min_size(data.bytes())?;
209        Ok(unsafe { Self::new_unchecked(data, iter) })
210    }
211    /// # Safety
212    ///
213    /// `data` must be aligned and have sufficient size.
214    pub unsafe fn new_unchecked(data: D, iter: I) -> Self {
215        Self {
216            _ghost: PhantomData,
217            data,
218            iter: PosIter::new(iter),
219        }
220    }
221    pub fn pos(&self) -> usize {
222        self.iter.pos()
223    }
224    pub fn value(self) -> D::Output<I::Item> {
225        self.data.value()
226    }
227}
228impl<'a, D: Data<'a>, T: Flat + Sized, I: TypeIter> DataIter<'a, D, TwoOrMoreTypes<T, I>> {
229    pub fn next(self) -> (DataIter<'a, D, I>, D::Output<T>) {
230        let prev_pos = self.iter.pos();
231        let iter = self.iter.next();
232        let next_pos = iter.pos();
233        let (prev_data, next_data) = self.data.split(next_pos - prev_pos);
234        (
235            DataIter {
236                _ghost: PhantomData,
237                data: next_data,
238                iter,
239            },
240            prev_data.value(),
241        )
242    }
243}
244impl<'a, D: Data<'a>, T: Flat + ?Sized> DataIter<'a, D, SingleType<T>> {
245    pub fn assert_last(&self) {
246        self.iter.assert_last()
247    }
248    pub fn finalize(self) -> D::Output<T> {
249        self.data.value()
250    }
251}
252
253pub type BytesIter<'a, I> = DataIter<'a, &'a [u8], I>;
254pub type BytesMutIter<'a, I> = DataIter<'a, &'a mut [u8], I>;
255pub type RefIter<'a, I> = DataIter<'a, RefData<'a>, I>;
256pub type MutIter<'a, I> = DataIter<'a, MutData<'a>, I>;
257
258pub trait ValidateIter {
259    fn validate_all(self) -> Result<(), Error>;
260}
261impl<'a, T: Flat + Sized + 'a, I: TypeIter> ValidateIter for BytesIter<'a, TwoOrMoreTypes<T, I>>
262where
263    BytesIter<'a, I>: ValidateIter,
264    I::Item: 'a,
265{
266    fn validate_all(self) -> Result<(), Error> {
267        unsafe { T::validate_unchecked(self.clone().value()) }.map_err(|e| e.offset(self.pos()))?;
268        self.next().0.validate_all()
269    }
270}
271
272impl<'a, T: Flat + ?Sized> ValidateIter for BytesIter<'a, SingleType<T>> {
273    fn validate_all(self) -> Result<(), Error> {
274        self.assert_last();
275        unsafe { T::validate_unchecked(self.clone().value()) }.map_err(|e| e.offset(self.pos()))?;
276        Ok(())
277    }
278}
279
280pub trait FoldSizeIter {
281    /// # Safety
282    ///
283    /// Internal data must be valid.
284    unsafe fn fold_size(self, size: usize) -> usize;
285}
286impl<'a, T: Flat + Sized + 'a, I: TypeIter> FoldSizeIter for BytesIter<'a, TwoOrMoreTypes<T, I>>
287where
288    BytesIter<'a, I>: FoldSizeIter,
289    I::Item: 'a,
290{
291    unsafe fn fold_size(self, size: usize) -> usize {
292        self.next().0.fold_size(ceil_mul(size, T::ALIGN) + T::SIZE)
293    }
294}
295impl<'a, T: Flat + ?Sized> FoldSizeIter for BytesIter<'a, SingleType<T>> {
296    unsafe fn fold_size(self, size: usize) -> usize {
297        ceil_mul(size, T::ALIGN) + (*T::ptr_from_bytes(self.finalize() as *const _ as *mut _)).size()
298    }
299}
300
301pub mod prelude {
302    pub use super::{FoldSizeIter, TypeIter, ValidateIter};
303}
304
305#[doc(hidden)]
306#[macro_export]
307macro_rules! type_list {
308    ($first_type:ty, $($types:ty),+ $(,)?) => {
309        $crate::utils::iter::TwoOrMoreTypes::<$first_type, _>::new($crate::utils::iter::type_list!($( $types ),*))
310    };
311    ($type:ty $(,)?) => {
312        $crate::utils::iter::SingleType::<$type>::new()
313    };
314}
315
316#[doc(hidden)]
317#[macro_export]
318macro_rules! fold_size {
319    ($accum:expr; $first_type:ty, $($types:ty),+ $(,)?) => {
320        $crate::utils::iter::fold_size!(
321            $crate::utils::ceil_mul($accum, <$first_type as $crate::traits::FlatBase>::ALIGN) + <$first_type as $crate::traits::FlatSized>::SIZE;
322            $( $types ),*
323        )
324    };
325    ($accum:expr; $type:ty $(,)?) => {
326        $crate::utils::ceil_mul($accum, <$type as $crate::traits::FlatBase>::ALIGN) + <$type as $crate::traits::FlatSized>::SIZE
327    };
328}
329
330#[doc(hidden)]
331#[macro_export]
332macro_rules! fold_min_size {
333    ($accum:expr; $first_type:ty, $($types:ty),+ $(,)?) => {
334        $crate::utils::iter::fold_min_size!(
335            $crate::utils::ceil_mul($accum, <$first_type as $crate::traits::FlatBase>::ALIGN) + <$first_type as $crate::traits::FlatSized>::SIZE;
336            $( $types ),*
337        )
338    };
339    ($accum:expr; $type:ty $(,)?) => {
340        $crate::utils::ceil_mul($accum, <$type as $crate::traits::FlatBase>::ALIGN) + <$type as $crate::traits::FlatBase>::MIN_SIZE
341    };
342}
343
344pub use {fold_min_size, fold_size, type_list};
345
346#[cfg(test)]
347mod tests {
348    use super::{type_list, PosIter};
349
350    #[test]
351    fn pos() {
352        let iter = PosIter::new(type_list!(u8, u16, u32));
353        assert_eq!(iter.pos(), 0);
354        let iter = iter.next();
355        assert_eq!(iter.pos(), 2);
356        let iter = iter.next();
357        assert_eq!(iter.pos(), 4);
358        iter.assert_last();
359    }
360}
flatty_base/utils/iter.rs

flatty_base/utils/
iter.rs
