1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107

use crate::error::Result;
use std::io::Read;

use super::util::read_u8;

const MAX_LITERAL_SIZE: usize = 128;
const MIN_REPEAT_SIZE: usize = 3;

pub struct ByteRleIter<R: Read> {
    reader: R,
    literals: [u8; MAX_LITERAL_SIZE],
    num_literals: usize,
    used: usize,
    repeat: bool,
}

impl<R: Read> ByteRleIter<R> {
    pub fn new(reader: R) -> Self {
        Self {
            reader,
            literals: [0u8; MAX_LITERAL_SIZE],
            num_literals: 0,
            used: 0,
            repeat: false,
        }
    }

    pub fn into_inner(self) -> R {
        self.reader
    }

    fn read_byte(&mut self) -> Result<u8> {
        read_u8(&mut self.reader)
    }

    fn read_values(&mut self) -> Result<()> {
        let control = self.read_byte()?;
        self.used = 0;
        if control < 0x80 {
            self.repeat = true;
            self.num_literals = control as usize + MIN_REPEAT_SIZE;
            let val = self.read_byte()?;
            self.literals[0] = val;
        } else {
            self.repeat = false;
            self.num_literals = 0x100 - control as usize;
            for i in 0..self.num_literals {
                let result = self.read_byte()?;
                self.literals[i] = result;
            }
        }
        Ok(())
    }
}

impl<R: Read> Iterator for ByteRleIter<R> {
    type Item = Result<u8>;

    fn next(&mut self) -> Option<Self::Item> {
        if self.used == self.num_literals {
            match self.read_values() {
                Ok(_) => {}
                Err(_err) => return None,
            }
        }

        let result = if self.repeat {
            self.literals[0]
        } else {
            self.literals[self.used]
        };
        self.used += 1;
        Some(Ok(result))
    }
}

#[cfg(test)]
mod test {
    use super::*;

    #[test]
    fn reader_test() {
        let data = [0x61u8, 0x00];

        let data = &mut data.as_ref();

        let iter = ByteRleIter::new(data).collect::<Result<Vec<_>>>().unwrap();

        assert_eq!(iter, vec![0; 100]);

        let data = [0x01, 0x01];

        let data = &mut data.as_ref();

        let iter = ByteRleIter::new(data).collect::<Result<Vec<_>>>().unwrap();

        assert_eq!(iter, vec![1; 4]);

        let data = [0xfe, 0x44, 0x45];

        let data = &mut data.as_ref();

        let iter = ByteRleIter::new(data).collect::<Result<Vec<_>>>().unwrap();

        assert_eq!(iter, vec![0x44, 0x45]);
    }
}