Skip to main content

sharky_arrayvec/
deku.rs

1use crate::ArrayVec;
2use deku::ctx::{Limit, Order, ReadExact};
3use deku::prelude::*;
4use std::io::prelude::*;
5use std::mem;
6
7impl<const N: usize> DekuReader<'_, ReadExact> for ArrayVec<N, u8> {
8    fn from_reader_with_ctx<R: Read + Seek>(
9        reader: &mut Reader<R>,
10        exact: ReadExact,
11    ) -> Result<Self, DekuError>
12    where
13        Self: Sized, {
14        assert!(exact.0 <= N);
15        let mut bytes = [0x00; N];
16        let _ = reader.read_bytes(exact.0, &mut bytes[..exact.0], Order::Lsb0)?;
17
18        Ok(Self::from_slice(&bytes[..exact.0]).unwrap())
19    }
20}
21
22/// Read `T`s into a vec until a given predicate returns true
23/// * `capacity` - an optional capacity to pre-allocate the vector with
24/// * `ctx` - The context required by `T`. It will be passed to every `T` when
25///   constructing.
26/// * `predicate` - the predicate that decides when to stop reading `T`s The
27///   predicate takes two parameters: the number of bits that have been read so
28///   far, and a borrow of the latest value to have been read. It should return
29///   `true` if reading should now stop, and `false` otherwise
30fn reader_vec_with_predicate<'a, const N: usize, T, Ctx, Predicate, R: Read + Seek>(
31    reader: &mut Reader<R>,
32    ctx: Ctx,
33    mut predicate: Predicate,
34) -> Result<ArrayVec<N, T>, DekuError>
35where
36    T: DekuReader<'a, Ctx>,
37    Ctx: Copy,
38    Predicate: FnMut(usize, &T) -> bool, {
39    // ZST detected, return empty vec
40    if mem::size_of::<T>() == 0 {
41        return Ok(ArrayVec::new());
42    }
43
44    let mut res = ArrayVec::new();
45
46    let start_read = reader.bits_read;
47
48    loop {
49        let val = <T>::from_reader_with_ctx(reader, ctx)?;
50        res.push(val);
51
52        // This unwrap is safe as we are pushing to the vec immediately before it,
53        // so there will always be a last element
54        if predicate(reader.bits_read - start_read, res.last().unwrap()) {
55            break;
56        }
57    }
58
59    Ok(res)
60}
61
62fn reader_vec_to_end<'a, const N: usize, T, Ctx, R: Read + Seek>(
63    reader: &mut Reader<R>,
64    ctx: Ctx,
65) -> Result<ArrayVec<N, T>, DekuError>
66where
67    T: DekuReader<'a, Ctx>,
68    Ctx: Copy, {
69    // ZST detected, return empty vec
70    if mem::size_of::<T>() == 0 {
71        return Ok(ArrayVec::new());
72    }
73
74    let mut res = ArrayVec::new();
75    loop {
76        if reader.end() {
77            break;
78        }
79        let val = <T>::from_reader_with_ctx(reader, ctx)?;
80        res.push(val);
81    }
82
83    Ok(res)
84}
85
86impl<'a, const N: usize, T, Ctx, Predicate> DekuReader<'a, (Limit<T, Predicate>, Ctx)>
87    for ArrayVec<N, T>
88where
89    T: DekuReader<'a, Ctx>,
90    Ctx: Copy,
91    Predicate: FnMut(&T) -> bool,
92{
93    fn from_reader_with_ctx<R: Read + Seek>(
94        reader: &mut Reader<R>,
95        (limit, inner_ctx): (Limit<T, Predicate>, Ctx),
96    ) -> Result<Self, DekuError>
97    where
98        Self: Sized, {
99        match limit {
100            // Read a given count of elements
101            Limit::Count(mut count) => {
102                assert!(count <= N);
103
104                // Handle the trivial case of reading an empty vector
105                if count == 0 {
106                    return Ok(ArrayVec::new());
107                }
108
109                // Otherwise, read until we have read `count` elements
110                reader_vec_with_predicate(reader, inner_ctx, move |_, _| {
111                    count -= 1;
112                    count == 0
113                })
114            }
115
116            // Read until a given predicate returns true
117            Limit::Until(mut predicate, _) => {
118                reader_vec_with_predicate(reader, inner_ctx, move |_, value| predicate(value))
119            }
120
121            // Read until a given quantity of bits have been read
122            Limit::BitSize(size) => {
123                let bit_size = size.0;
124
125                // Handle the trivial case of reading an empty vector
126                if bit_size == 0 {
127                    return Ok(ArrayVec::new());
128                }
129
130                reader_vec_with_predicate(reader, inner_ctx, move |read_bits, _| {
131                    read_bits == bit_size
132                })
133            }
134
135            // Read until a given quantity of bytes have been read
136            Limit::ByteSize(size) => {
137                let bit_size = size.0 * 8;
138
139                // Handle the trivial case of reading an empty vector
140                if bit_size == 0 {
141                    return Ok(ArrayVec::new());
142                }
143
144                reader_vec_with_predicate(reader, inner_ctx, move |read_bits, _| {
145                    read_bits == bit_size
146                })
147            }
148
149            Limit::End => reader_vec_to_end(reader, inner_ctx),
150        }
151    }
152}
153
154impl<'a, const N: usize, T: DekuReader<'a>, Predicate: FnMut(&T) -> bool>
155    DekuReader<'a, Limit<T, Predicate>> for ArrayVec<N, T>
156{
157    /// Read `T`s until the given limit from input for types which don't require
158    /// context.
159    fn from_reader_with_ctx<R: Read + Seek>(
160        reader: &mut Reader<R>,
161        limit: Limit<T, Predicate>,
162    ) -> Result<Self, DekuError>
163    where
164        Self: Sized, {
165        ArrayVec::from_reader_with_ctx(reader, (limit, ()))
166    }
167}
168
169impl<const N: usize, T: DekuWriter<Ctx>, Ctx: Copy> DekuWriter<Ctx> for ArrayVec<N, T> {
170    fn to_writer<W: Write + Seek>(
171        &self,
172        writer: &mut Writer<W>,
173        inner_ctx: Ctx,
174    ) -> Result<(), DekuError> {
175        for v in self.iter() {
176            v.to_writer(writer, inner_ctx)?;
177        }
178        Ok(())
179    }
180}
181
182#[cfg(test)]
183mod tests {
184
185    use super::*;
186
187    #[deku_derive(DekuRead, DekuWrite)]
188    #[derive(Debug, Clone, PartialEq, Eq)]
189    pub struct Cuesheet {
190        #[deku(endian = "big", count = "128")]
191        pub media_catalog_number: ArrayVec<128, u8>,
192    }
193}