1
use std::fmt;
2
use std::marker::PhantomData;
3

            
4
/// A vector data structure that stores objects in a byte compressed format
5
#[derive(Debug, Default)]
6
pub struct ByteCompressedVec<T> {
7
    data: Vec<u8>,
8
    bytes_per_entry: usize,
9
    _marker: PhantomData<T>,
10
}
11

            
12
impl<T: Entry + fmt::Debug> ByteCompressedVec<T> {
13
3
    pub fn new() -> ByteCompressedVec<T> {
14
3
        ByteCompressedVec {
15
3
            data: Vec::new(),
16
3
            bytes_per_entry: 0,
17
3
            _marker: PhantomData,
18
3
        }
19
3
    }
20

            
21
    /// Adds a new entry to the vector.
22
202
    pub fn push(&mut self, entry: T) {
23
202
        self.resize_entries(entry.bytes_required());
24
202

            
25
202
        // Add the new entry to the end of the vector.
26
202
        let old_len = self.data.len();
27
202
        self.data.resize(old_len + self.bytes_per_entry, 0);
28
202
        entry.to_bytes(&mut self.data[old_len..]);
29
202
    }
30

            
31
    /// Returns the entry at the given index.
32
5253
    pub fn get(&self, index: usize) -> T {
33
5253
        let start = index * self.bytes_per_entry;
34
5253
        let end = start + self.bytes_per_entry;
35
5253
        T::from_bytes(&self.data[start..end])
36
5253
    }
37

            
38
    /// Sets the entry at the given index.
39
100
    pub fn set(&mut self, index: usize, entry: T) {
40
100
        self.resize_entries(entry.bytes_required());
41
100

            
42
100
        let start = index * self.bytes_per_entry;
43
100
        let end = start + self.bytes_per_entry;
44
100
        entry.to_bytes(&mut self.data[start..end]);
45
100
    }
46

            
47
    /// Returns the number of elements in the vector.
48
5359
    pub fn len(&self) -> usize {
49
5359
        if self.bytes_per_entry == 0 {
50
3
            0
51
        } else {
52
5356
            debug_assert!(self.data.len() % self.bytes_per_entry == 0);
53
5356
            self.data.len() / self.bytes_per_entry
54
        }
55
5359
    }
56

            
57
    /// Returns true if the vector is empty.
58
    pub fn is_empty(&self) -> bool {
59
        self.len() == 0
60
    }
61

            
62
    /// Returns an iterator over the elements in the vector.
63
103
    pub fn iter(&self) -> ByteCompressedVecIterator<T> {
64
103
        ByteCompressedVecIterator {
65
103
            vector: self,
66
103
            current: 0,
67
103
        }
68
103
    }
69

            
70
    /// Resizes all entries in the vector to the given length.
71
302
    fn resize_entries(&mut self, new_bytes_required: usize) {
72
302
        if new_bytes_required > self.bytes_per_entry {
73
5
            let mut new_data: Vec<u8> = vec![0; self.len() * new_bytes_required];
74
5

            
75
5
            if self.bytes_per_entry > 0 {
76
                // Resize all the existing elements because the new entry requires more bytes.
77
101
                for (index, entry) in self.iter().enumerate() {
78
101
                    let start = index * new_bytes_required;
79
101
                    let end = start + new_bytes_required;
80
101
                    entry.to_bytes(&mut new_data[start..end]);
81
101
                }
82
3
            }
83

            
84
5
            self.bytes_per_entry = new_bytes_required;
85
5
            self.data = new_data;
86
297
        }
87
302
    }
88
}
89

            
90
impl<T: Entry + Clone + fmt::Debug> ByteCompressedVec<T> {
91
1
    pub fn from_elem(entry: T, n: usize) -> ByteCompressedVec<T> {
92
1
        let mut vec = ByteCompressedVec::new();
93
100
        for _ in 0..n {
94
100
            vec.push(entry.clone());
95
100
        }
96
1
        vec
97
1
    }
98
}
99

            
100
pub struct ByteCompressedVecIterator<'a, T> {
101
    vector: &'a ByteCompressedVec<T>,
102
    current: usize,
103
}
104

            
105
impl<T: Entry + fmt::Debug> Iterator for ByteCompressedVecIterator<'_, T> {
106
    type Item = T;
107

            
108
5352
    fn next(&mut self) -> Option<Self::Item> {
109
5352
        if self.current < self.vector.len() {
110
5251
            let result = self.vector.get(self.current);
111
5251
            self.current += 1;
112
5251
            Some(result)
113
        } else {
114
101
            None
115
        }
116
5352
    }
117
}
118

            
119
pub trait Entry {
120
    // Returns the entry as a byte vector
121
    fn to_bytes(&self, bytes: &mut [u8]);
122

            
123
    // Creates an entry from a byte vector
124
    fn from_bytes(bytes: &[u8]) -> Self;
125

            
126
    // Returns the number of bytes required to store the current entry
127
    fn bytes_required(&self) -> usize;
128
}
129

            
130
impl Entry for usize {
131
503
    fn to_bytes(&self, bytes: &mut [u8]) {
132
503
        let array = &self.to_le_bytes();
133
2705
        for (i, byte) in bytes.iter_mut().enumerate() {
134
2705
            *byte = array[i];
135
2705
        }
136
503
    }
137

            
138
5353
    fn from_bytes(bytes: &[u8]) -> Self {
139
5353
        let mut array = [0; 8];
140
41505
        for (i, byte) in bytes.iter().enumerate() {
141
41505
            array[i] = *byte;
142
41505
        }
143
5353
        usize::from_le_bytes(array)
144
5353
    }
145

            
146
402
    fn bytes_required(&self) -> usize {
147
402
        ((self + 1).ilog2() / u8::BITS) as usize + 1
148
402
    }
149
}
150

            
151
#[cfg(test)]
152
mod tests {
153
    use crate::bytevec;
154
    use rand::distr::Uniform;
155
    use rand::Rng;
156
    use test_log::test;
157

            
158
    use super::*;
159

            
160
1
    #[test]
161
    fn test_index_bytevector() {
162
        let mut vec = ByteCompressedVec::new();
163
        vec.push(1);
164
        assert_eq!(vec.len(), 1);
165

            
166
        vec.push(1024);
167
        assert_eq!(vec.len(), 2);
168

            
169
        assert_eq!(vec.get(0), 1);
170
        assert_eq!(vec.get(1), 1024);
171
    }
172

            
173
1
    #[test]
174
    fn test_random_bytevector() {
175
        let rng = rand::rng();
176

            
177
        let range = Uniform::new(0, usize::MAX).unwrap();
178
        let expected_vector: Vec<usize> = rng.sample_iter(range).take(100).collect();
179
        let mut vector = ByteCompressedVec::new();
180

            
181
        for element in &expected_vector {
182
            vector.push(*element);
183

            
184
            for (expected, element) in expected_vector.iter().zip(vector.iter()) {
185
                assert_eq!(*expected, element);
186
            }
187
        }
188
    }
189

            
190
1
    #[test]
191
    fn test_random_setting_bytevector() {
192
        let rng = rand::rng();
193

            
194
        let range = Uniform::new(0, usize::MAX).unwrap();
195
        let expected_vector: Vec<usize> = rng.sample_iter(range).take(100).collect();
196
        let mut vector = bytevec![0; 100];
197

            
198
        for (index, element) in expected_vector.iter().enumerate() {
199
            vector.set(index, *element);
200
        }
201

            
202
        for (expected, element) in expected_vector.iter().zip(vector.iter()) {
203
            assert_eq!(*expected, element);
204
        }
205
    }
206

            
207
1
    #[test]
208
    fn test_usize_entry() {
209
        let mut rng = rand::rng();
210

            
211
        for _ in 0..100 {
212
            let value = rng.random_range(0..1024);
213
            assert!(value.bytes_required() <= 2);
214

            
215
            let mut bytes = [0; 2];
216
            value.to_bytes(&mut bytes);
217
            assert_eq!(usize::from_bytes(&bytes), value);
218
        }
219
    }
220
}