gel_protogen/
arrays.rs

1#![allow(private_bounds)]
2
3use crate::prelude::*;
4
5pub use std::marker::PhantomData;
6
7/// Shared implementation for all array types.
8macro_rules! array_impl {
9    (#[$doc:meta] impl <$lt:lifetime, $generic:ident $(, $length_generic:ident)?> $ty:ident) => {
10        #[$doc]
11        #[derive(Copy, Clone, Default)]
12        pub struct $ty<$lt, $($length_generic,)? $generic>
13        where
14            $generic: DecoderFor<$lt, $generic>,
15        {
16            _phantom: PhantomData<( $generic , $( $length_generic)? )>,
17            buf: &'a [u8],
18            len: usize,
19        }
20
21        impl<$lt, $generic, $($length_generic)?> $ty<$lt, $($length_generic,)? $generic>
22        where
23            $generic: DecoderFor<$lt, $generic>,
24        {
25            #[inline(always)]
26            pub const fn new(buf: &$lt [u8], len: usize) -> Self {
27                Self {
28                    buf,
29                    len,
30                    _phantom: PhantomData,
31                }
32            }
33
34            #[inline(always)]
35            pub const fn empty() -> Self {
36                Self {
37                    buf: &[],
38                    len: 0,
39                    _phantom: PhantomData,
40                }
41            }
42
43            #[inline(always)]
44            pub const fn len(&self) -> usize {
45                self.len
46            }
47
48            #[inline(always)]
49            pub const fn is_empty(&self) -> bool {
50                self.len == 0
51            }
52
53            #[inline(always)]
54            pub const fn into_slice(self) -> &'a [u8] {
55                self.buf
56            }
57        }
58
59        impl<$lt, $generic, $($length_generic)?> std::fmt::Debug for $ty<$lt, $($length_generic,)? $generic>
60        where
61            $generic: DecoderFor<$lt, $generic>,
62            $generic: std::fmt::Debug,
63        {
64            fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
65                f.debug_list().entries(self).finish()
66            }
67        }
68
69        // ZTArrays of type [`u8`] are special-cased to return a slice of bytes.
70        impl<$lt, $generic, $($length_generic)?> ArrayExt<$lt> for $ty<$lt, $($length_generic,)? $generic>
71        where
72            $generic: DecoderFor<$lt, $generic>,
73            $( $length_generic: $lt )?
74        {
75            #[inline(always)]
76            fn into_slice(self) -> &'a [u8] {
77                self.buf
78            }
79        }
80
81        impl<$lt, $generic, $($length_generic)?> AsRef<[u8]> for $ty<$lt, $($length_generic,)? $generic>
82        where
83            $generic: DecoderFor<$lt, $generic>,
84        {
85            #[inline(always)]
86            fn as_ref(&self) -> &[u8] {
87                self.buf
88            }
89        }
90
91        impl<$lt, $generic, $($length_generic)?> IntoIterator for $ty<$lt, $($length_generic,)? $generic>
92        where
93            $generic: DecoderFor<$lt, $generic>,
94        {
95            type Item = $generic;
96            type IntoIter = ArrayIter<'a, $generic>;
97            fn into_iter(self) -> Self::IntoIter {
98                Self::IntoIter {
99                    _phantom: PhantomData,
100                    buf: self.buf,
101                    len: self.len,
102                }
103            }
104        }
105
106        impl<$lt, $generic, $($length_generic)?> IntoIterator for &$ty<$lt, $($length_generic,)? $generic>
107        where
108            $generic: DecoderFor<$lt, $generic>,
109        {
110            type Item = $generic;
111            type IntoIter = ArrayIter<'a, $generic>;
112            fn into_iter(self) -> Self::IntoIter {
113                Self::IntoIter {
114                    _phantom: PhantomData,
115                    buf: self.buf,
116                    len: self.len,
117                }
118            }
119        }
120
121        // Arrays of fixed-size elements can extract elements in O(1).
122        impl<$lt, $generic, $($length_generic)?> $ty<$lt, $($length_generic,)? $generic>
123        where
124            $generic: DataTypeFixedSize + DecoderFor<$lt, $generic>,
125        {
126            #[inline]
127            pub fn get(&self, index: impl TryInto<usize>) -> Option<$generic> {
128                let Ok(index) = index.try_into() else {
129                    return None;
130                };
131                let index: usize = index;
132                if index >= self.len as _ {
133                    None
134                } else {
135                    let mut segment = &self.buf[T::SIZE * index..T::SIZE * (index + 1)];
136                    // As we've normally pre-scanned all items, this will not panic
137                    Some(T::decode_for(&mut segment).unwrap())
138                }
139            }
140        }
141
142        /// Arrays of `u8` can be indexed.
143        impl<$lt, $($length_generic)?> std::ops::Index<usize> for $ty<$lt, $($length_generic,)? u8> {
144            type Output = u8;
145            #[inline(always)]
146            fn index(&self, index: usize) -> &Self::Output {
147                &self.as_ref()[index]
148            }
149        }
150
151        /// Arrays of `u8` can be compared to slices.
152        impl<$lt, $($length_generic)?> PartialEq<&[u8]> for $ty<$lt, $($length_generic,)? u8>
153        {
154            fn eq(&self, other: &&[u8]) -> bool {
155                self.as_ref() == *other
156            }
157        }
158
159        /// Arrays of `u8` can be compared to fixed-size slices.
160        impl<$lt, $($length_generic, )? const N: usize> PartialEq<&[u8; N]> for $ty<$lt, $($length_generic,)? u8>
161        {
162            fn eq(&self, other: &&[u8; N]) -> bool {
163                self.as_ref() == *other
164            }
165        }
166    };
167}
168
169/// Shared trait for all array types.
170pub trait ArrayExt<'a>: 'a {
171    /// Convert the array into a slice of bytes.
172    fn into_slice(self) -> &'a [u8];
173}
174
175array_impl!(
176    /// A zero-terminated array.
177    impl <'a, T> ZTArray
178);
179array_impl!(
180    /// A count-prefixed array.
181    impl <'a, T, L> Array
182);
183array_impl!(
184    /// A rest array: consumes the remainder of the buffer.
185    impl <'a, T> RestArray
186);
187
188/// [`ZTArray`], [`Array`], and [`RestArray`] [`Iterator`] for values of type `T`.
189#[derive(Copy, Clone, Default)]
190pub struct ArrayIter<'a, T> {
191    _phantom: PhantomData<T>,
192    buf: &'a [u8],
193    len: usize,
194}
195
196impl<'a, T> Iterator for ArrayIter<'a, T>
197where
198    T: DecoderFor<'a, T>,
199{
200    type Item = T;
201    fn next(&mut self) -> Option<Self::Item> {
202        if self.len == 0 {
203            return None;
204        }
205        self.len -= 1;
206        let value = T::decode_for(&mut self.buf).ok()?;
207        Some(value)
208    }
209
210    #[inline(always)]
211    fn size_hint(&self) -> (usize, Option<usize>) {
212        (self.len, Some(self.len))
213    }
214}
215
216impl<'a, T> ExactSizeIterator for ArrayIter<'a, T>
217where
218    T: DecoderFor<'a, T>,
219{
220    #[inline(always)]
221    fn len(&self) -> usize {
222        self.len as usize
223    }
224}
225
226#[cfg(test)]
227mod tests {
228    use super::*;
229
230    #[test]
231    fn test_rest_array_u8() {
232        // Test with u8 values
233        let data = vec![1, 2, 3, 4, 5];
234        let mut buf = &data[..];
235        let rest_array = RestArray::<u8>::decode_for(&mut buf).unwrap();
236
237        assert_eq!(rest_array.len(), 5);
238        assert!(!rest_array.is_empty());
239        assert_eq!(buf.len(), 0); // Buffer should be consumed entirely
240
241        let collected: Vec<u8> = rest_array.into_iter().collect();
242        assert_eq!(collected, vec![1, 2, 3, 4, 5]);
243    }
244
245    #[test]
246    fn test_rest_array_u32() {
247        let data = vec![
248            0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x02, 0x00, 0x00, 0x00, 0x03,
249        ];
250
251        let mut buf = &data[..];
252        let rest_array = RestArray::<u32>::decode_for(&mut buf).unwrap();
253
254        assert_eq!(rest_array.len(), 3);
255        assert!(!rest_array.is_empty());
256        assert_eq!(buf.len(), 0); // Buffer should be consumed entirely
257
258        let collected: Vec<u32> = rest_array.into_iter().collect();
259        assert_eq!(collected, vec![1, 2, 3]);
260    }
261
262    #[test]
263    fn test_rest_array_empty() {
264        // Test with empty buffer
265        let data: Vec<u8> = vec![];
266        let mut buf = &data[..];
267        let rest_array = RestArray::<u8>::decode_for(&mut buf).unwrap();
268
269        assert_eq!(rest_array.len(), 0);
270        assert!(rest_array.is_empty());
271        assert_eq!(buf.len(), 0);
272
273        let collected: Vec<u8> = rest_array.into_iter().collect();
274        assert_eq!(collected, vec![]);
275    }
276
277    #[test]
278    fn test_rest_array_get() {
279        // Test get method for fixed-size elements
280        let data = vec![1u8, 2, 3, 4, 5];
281        let mut buf = &data[..];
282        let rest_array = RestArray::<u8>::decode_for(&mut buf).unwrap();
283
284        assert_eq!(rest_array.get(0), Some(1));
285        assert_eq!(rest_array.get(2), Some(3));
286        assert_eq!(rest_array.get(4), Some(5));
287        assert_eq!(rest_array.get(5), None); // Out of bounds
288    }
289
290    #[test]
291    fn test_array_u32() {
292        let data = vec![
293            0x00, 0x00, 0x00, 0x03, // Length prefix
294            0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x02, 0x00, 0x00, 0x00, 0x03,
295        ];
296
297        let mut buf = &data[..];
298        let array = Array::<u32, u32>::decode_for(&mut buf).unwrap();
299
300        assert_eq!(array.len(), 3);
301        assert!(!array.is_empty());
302        assert_eq!(buf.len(), 0);
303
304        let collected: Vec<u32> = array.into_iter().collect();
305        assert_eq!(collected, vec![1, 2, 3]);
306    }
307
308    #[test]
309    fn test_array_invalid_length() {
310        let data = vec![
311            0xFF, 0xFF, 0xFF, 0xFF, // Invalid length
312            0x00, 0x00, 0x00, 0x01,
313        ];
314
315        let mut buf = &data[..];
316        let result = Array::<u32, u32>::decode_for(&mut buf);
317        assert!(result.is_err());
318
319        let mut buf = [].as_slice();
320        let result = Array::<u32, u32>::decode_for(&mut buf);
321        assert!(result.is_err());
322    }
323
324    #[test]
325    fn test_zt_array() {
326        let data = vec![
327            0x01, 0x02, 0x03, 0x00, // Zero-terminated array
328        ];
329
330        let mut buf = &data[..];
331        let array = ZTArray::<u8>::decode_for(&mut buf).unwrap();
332
333        assert_eq!(array.len(), 3);
334        assert!(!array.is_empty());
335        assert_eq!(buf.len(), 0);
336
337        let collected: Vec<u8> = array.into_iter().collect();
338        assert_eq!(collected, vec![1, 2, 3]);
339    }
340
341    #[test]
342    fn test_zt_array_u32() {
343        // Unlikely, but helps test our primitive fast path
344        let data = vec![0xff, 0xff, 0xff, 0xff, 0xfe, 0xfe, 0xfe, 0xfe, 0];
345
346        let mut buf = &data[..];
347        let array = ZTArray::<u32>::decode_for(&mut buf).unwrap();
348
349        assert_eq!(array.len(), 2);
350        assert!(!array.is_empty());
351        assert_eq!(buf.len(), 0);
352    }
353
354    #[test]
355    fn test_zt_array_string() {
356        let data = vec![
357            b'h', b'e', b'l', b'l', b'o', b'\0', b'w', b'o', b'r', b'l', b'd', b'\0',
358            b'\0', // Zero-terminated array
359        ];
360
361        let mut buf = &data[..];
362        let array = ZTArray::<ZTString>::decode_for(&mut buf).unwrap();
363
364        assert_eq!(array.len(), 2);
365        assert!(!array.is_empty());
366        assert_eq!(buf.len(), 0);
367
368        let collected: Vec<_> = array.into_iter().collect();
369        assert_eq!(collected, vec!["hello", "world"]);
370    }
371
372    #[test]
373    fn test_zt_array_missing_terminator() {
374        let data = vec![0x01, 0x02, 0x03]; // No zero terminator
375
376        let mut buf = &data[..];
377        let result = ZTArray::<u8>::decode_for(&mut buf);
378        assert!(result.is_err());
379
380        // Test with empty arrays
381        let mut buf = [].as_slice();
382        assert!(ZTArray::<u8>::decode_for(&mut buf).is_err());
383        assert!(ZTArray::<u32>::decode_for(&mut buf).is_err());
384        assert!(ZTArray::<ZTString>::decode_for(&mut buf).is_err());
385    }
386
387    #[test]
388    fn test_zt_array_empty() {
389        let data = vec![0x00]; // Just terminator
390
391        let mut buf = &data[..];
392        let array = ZTArray::<u8>::decode_for(&mut buf).unwrap();
393
394        assert_eq!(array.len(), 0);
395        assert!(array.is_empty());
396        assert_eq!(buf.len(), 0);
397
398        let collected: Vec<u8> = array.into_iter().collect();
399        assert_eq!(collected, vec![]);
400    }
401}