orc_format/read/decode/
boolean_rle.rs

1use std::io::Read;
2
3use crate::error::Error;
4
5use super::read_u8;
6
7#[derive(Debug, Copy, Clone, PartialEq)]
8#[allow(clippy::large_enum_variant)]
9pub enum BooleanRun {
10    Run(u8, u16),
11    Literals([u8; 255]),
12}
13
14pub struct BooleanRleRunIter<R: Read> {
15    reader: R,
16}
17
18impl<R: Read> BooleanRleRunIter<R> {
19    pub fn new(reader: R) -> Self {
20        Self { reader }
21    }
22
23    pub fn into_inner(self) -> R {
24        self.reader
25    }
26}
27
28fn read_literals<R: Read>(reader: &mut R, header: i8) -> Result<[u8; 255], Error> {
29    let length = (-header) as usize;
30
31    let mut literals = [0u8; 255];
32
33    reader
34        .take(length as u64)
35        .read_exact(&mut literals[..length])?;
36
37    Ok(literals)
38}
39
40impl<R: Read> Iterator for BooleanRleRunIter<R> {
41    type Item = Result<BooleanRun, Error>;
42
43    #[inline]
44    fn next(&mut self) -> Option<Self::Item> {
45        let header = read_u8(&mut self.reader);
46        let header = match header {
47            Ok(header) => header as i8,
48            Err(e) => return Some(Err(e.into())),
49        };
50        if header < 0 {
51            Some(read_literals(&mut self.reader, header).map(BooleanRun::Literals))
52        } else {
53            let length = header as u16 + 3;
54            // this is not ok - it may require more than one byte
55            let value = read_u8(&mut self.reader);
56            let value = match value {
57                Ok(value) => value,
58                Err(e) => return Some(Err(e.into())),
59            };
60            Some(Ok(BooleanRun::Run(value, length)))
61        }
62    }
63}
64
65pub struct BooleanIter<R: Read> {
66    iter: BooleanRleRunIter<R>,
67    current: Option<BooleanRun>,
68    position: u8,
69    byte_position: usize,
70    remaining: usize,
71}
72
73impl<'a, R: Read> BooleanIter<R> {
74    pub fn new(reader: R, length: usize) -> Self {
75        Self {
76            iter: BooleanRleRunIter::new(reader),
77            current: None,
78            position: 0,
79            byte_position: 0,
80            remaining: length,
81        }
82    }
83
84    pub fn into_inner(self) -> R {
85        self.iter.into_inner()
86    }
87}
88
89impl<R: Read> Iterator for BooleanIter<R> {
90    type Item = Result<bool, Error>;
91
92    #[inline]
93    fn next(&mut self) -> Option<Self::Item> {
94        if let Some(run) = &self.current {
95            match run {
96                BooleanRun::Run(value, repetitions) => {
97                    let repetitions = *repetitions;
98                    let mask = 128u8 >> self.position;
99                    let result = value & mask == mask;
100                    self.position += 1;
101                    if self.remaining == 0 {
102                        self.current = None;
103                        return None;
104                    } else {
105                        self.remaining -= 1;
106                    }
107                    if self.position == 8 {
108                        if repetitions == 0 {
109                            self.current = None;
110                        } else {
111                            self.current = Some(BooleanRun::Run(*value, repetitions - 1));
112                        }
113                        self.position = 0;
114                    }
115                    Some(Ok(result))
116                }
117                BooleanRun::Literals(bytes) => {
118                    let mask = 128u8 >> self.position;
119                    let result = bytes[self.byte_position] & mask == mask;
120                    self.position += 1;
121                    if self.remaining == 0 {
122                        self.current = None;
123                        return None;
124                    } else {
125                        self.remaining -= 1;
126                    }
127                    if self.position == 8 {
128                        if bytes.len() == 1 {
129                            self.current = None;
130                            self.byte_position = 0;
131                        } else {
132                            self.byte_position += 1;
133                        }
134                        self.position = 0;
135                    }
136                    Some(Ok(result))
137                }
138            }
139        } else if self.remaining > 0 {
140            match self.iter.next()? {
141                Ok(run) => {
142                    self.current = Some(run);
143                    self.next()
144                }
145                Err(e) => {
146                    self.remaining = 0;
147                    Some(Err(e))
148                }
149            }
150        } else {
151            None
152        }
153    }
154
155    fn size_hint(&self) -> (usize, Option<usize>) {
156        (self.remaining, Some(self.remaining))
157    }
158}
159
160#[cfg(test)]
161mod test {
162    use super::*;
163
164    #[test]
165    fn basic() {
166        let data = [0x61u8, 0x00];
167
168        let data = &mut data.as_ref();
169
170        let iter = BooleanIter::new(data, 100)
171            .collect::<Result<Vec<_>, Error>>()
172            .unwrap();
173        assert_eq!(iter, vec![false; 100])
174    }
175
176    #[test]
177    fn literals() {
178        let data = [0xfeu8, 0b01000100, 0b01000101];
179
180        let data = &mut data.as_ref();
181
182        let iter = BooleanIter::new(data, 16)
183            .collect::<Result<Vec<_>, Error>>()
184            .unwrap();
185        assert_eq!(
186            iter,
187            vec![
188                false, true, false, false, false, true, false, false, // 0b01000100
189                false, true, false, false, false, true, false, true, // 0b01000101
190            ]
191        )
192    }
193
194    #[test]
195    fn another() {
196        // "For example, the byte sequence [0xff, 0x80] would be one true followed by seven false values."
197        let data = [0xff, 0x80];
198
199        let data = &mut data.as_ref();
200
201        let iter = BooleanIter::new(data, 8)
202            .collect::<Result<Vec<_>, Error>>()
203            .unwrap();
204        assert_eq!(
205            iter,
206            vec![true, false, false, false, false, false, false, false,]
207        )
208    }
209}