azalea_world/
bit_storage.rs

1use std::{error::Error, fmt};
2
3#[rustfmt::skip]
4const MAGIC: [(i32, i32, i32); 64] = [
5    // divide_mul, divide_add, divide_shift
6    (-1, -1, 0),
7    (-0b10000000000000000000000000000000, 0, 0),
8    (0b1010101010101010101010101010101, 0b1010101010101010101010101010101, 0),
9    (-0b10000000000000000000000000000000, 0, 1),
10    (0b110011001100110011001100110011, 0b110011001100110011001100110011, 0),
11    (0b101010101010101010101010101010, 0b101010101010101010101010101010, 0),
12    (0b100100100100100100100100100100, 0b100100100100100100100100100100, 0),
13    (-0b10000000000000000000000000000000, 0, 2),
14    (0b11100011100011100011100011100, 0b11100011100011100011100011100, 0),
15    (0b11001100110011001100110011001, 0b11001100110011001100110011001, 0),
16    (0b10111010001011101000101110100, 0b10111010001011101000101110100, 0),
17    (0b10101010101010101010101010101, 0b10101010101010101010101010101, 0),
18    (0b10011101100010011101100010011, 0b10011101100010011101100010011, 0),
19    (0b10010010010010010010010010010, 0b10010010010010010010010010010, 0),
20    (0b10001000100010001000100010001, 0b10001000100010001000100010001, 0),
21    (-0b10000000000000000000000000000000, 0, 3),
22    (0b1111000011110000111100001111, 0b1111000011110000111100001111, 0),
23    (0b1110001110001110001110001110, 0b1110001110001110001110001110, 0),
24    (0b1101011110010100001101011110, 0b1101011110010100001101011110, 0),
25    (0b1111111111111111111111111111000, 0b1111111111111111111111111111000, 0),
26    (0b1100001100001100001100001100, 0b1100001100001100001100001100, 0),
27    (0b1011101000101110100010111010, 0b1011101000101110100010111010, 0),
28    (0b1011001000010110010000101100, 0b1011001000010110010000101100, 0),
29    (0b1010101010101010101010101010, 0b1010101010101010101010101010, 0),
30    (0b1010001111010111000010100011, 0b1010001111010111000010100011, 0),
31    (0b1001110110001001110110001001, 0b1001110110001001110110001001, 0),
32    (0b1001011110110100001001011110, 0b1001011110110100001001011110, 0),
33    (0b1001001001001001001001001001, 0b1001001001001001001001001001, 0),
34    (0b1000110100111101110010110000, 0b1000110100111101110010110000, 0),
35    (0b1000100010001000100010001000, 0b1000100010001000100010001000, 0),
36    (0b1000010000100001000010000100, 0b1000010000100001000010000100, 0),
37    (-0b10000000000000000000000000000000, 0, 4),
38    (0b111110000011111000001111100, 0b111110000011111000001111100, 0),
39    (0b111100001111000011110000111, 0b111100001111000011110000111, 0),
40    (0b111010100000111010100000111, 0b111010100000111010100000111, 0),
41    (0b111000111000111000111000111, 0b111000111000111000111000111, 0),
42    (0b110111010110011111001000101, 0b110111010110011111001000101, 0),
43    (0b110101111001010000110101111, 0b110101111001010000110101111, 0),
44    (0b110100100000110100100000110, 0b110100100000110100100000110, 0),
45    (0b110011001100110011001100110, 0b110011001100110011001100110, 0),
46    (0b110001111100111000001100011, 0b110001111100111000001100011, 0),
47    (0b110000110000110000110000110, 0b110000110000110000110000110, 0),
48    (0b101111101000001011111010000, 0b101111101000001011111010000, 0),
49    (0b101110100010111010001011101, 0b101110100010111010001011101, 0),
50    (0b101101100000101101100000101, 0b101101100000101101100000101, 0),
51    (0b101100100001011001000010110, 0b101100100001011001000010110, 0),
52    (0b101011100100110001000001010, 0b101011100100110001000001010, 0),
53    (0b101010101010101010101010101, 0b101010101010101010101010101, 0),
54    (0b101001110010111100000101001, 0b101001110010111100000101001, 0),
55    (0b101000111101011100001010001, 0b101000111101011100001010001, 0),
56    (0b101000001010000010100000101, 0b101000001010000010100000101, 0),
57    (0b100111011000100111011000100, 0b100111011000100111011000100, 0),
58    (0b100110101001000011100111110, 0b100110101001000011100111110, 0),
59    (0b100101111011010000100101111, 0b100101111011010000100101111, 0),
60    (0b100101001111001000001001010, 0b100101001111001000001001010, 0),
61    (0b100100100100100100100100100, 0b100100100100100100100100100, 0),
62    (0b100011111011100000100011111, 0b100011111011100000100011111, 0),
63    (0b100011010011110111001011000, 0b100011010011110111001011000, 0),
64    (0b100010101101100011110010111, 0b100010101101100011110010111, 0),
65    (0b100010001000100010001000100, 0b100010001000100010001000100, 0),
66    (0b100001100100101110001010011, 0b100001100100101110001010011, 0),
67    (0b100001000010000100001000010, 0b100001000010000100001000010, 0),
68    (0b100000100000100000100000100, 0b100000100000100000100000100, 0),
69    (-0b10000000000000000000000000000000, 0, 5),
70];
71
72/// A compact list of integers with the given number of bits per entry.
73#[derive(Clone, Debug, Default)]
74pub struct BitStorage {
75    pub data: Vec<u64>,
76    bits: usize,
77    mask: u64,
78    size: usize,
79    values_per_long: usize,
80    divide_mul: u64,
81    divide_add: u64,
82    divide_shift: i32,
83}
84
85#[derive(Debug)]
86pub enum BitStorageError {
87    InvalidLength { got: usize, expected: usize },
88}
89impl fmt::Display for BitStorageError {
90    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
91        match self {
92            BitStorageError::InvalidLength { got, expected } => write!(
93                f,
94                "Invalid length given for storage, got: {got}, but expected: {expected}",
95            ),
96        }
97    }
98}
99impl Error for BitStorageError {}
100
101impl BitStorage {
102    /// Create a new BitStorage with the given number of bits per entry.
103    /// `size` is the number of entries in the BitStorage.
104    pub fn new(bits: usize, size: usize, data: Option<Vec<u64>>) -> Result<Self, BitStorageError> {
105        if let Some(data) = &data {
106            // 0 bit storage
107            if data.is_empty() {
108                return Ok(BitStorage {
109                    data: Vec::new(),
110                    bits,
111                    size,
112                    ..Default::default()
113                });
114            }
115        }
116
117        debug_assert!((1..=32).contains(&bits));
118
119        let values_per_long = 64 / bits;
120        let magic_index = values_per_long - 1;
121        let (divide_mul, divide_add, divide_shift) = MAGIC[magic_index];
122        let calculated_length = size.div_ceil(values_per_long);
123
124        let mask = (1 << bits) - 1;
125
126        let using_data = if let Some(data) = data {
127            if data.len() != calculated_length {
128                return Err(BitStorageError::InvalidLength {
129                    got: data.len(),
130                    expected: calculated_length,
131                });
132            }
133            data
134        } else {
135            vec![0; calculated_length]
136        };
137
138        Ok(BitStorage {
139            data: using_data,
140            bits,
141            mask,
142            size,
143            values_per_long,
144            divide_mul: divide_mul as u32 as u64,
145            divide_add: divide_add as u32 as u64,
146            divide_shift,
147        })
148    }
149
150    pub fn cell_index(&self, index: u64) -> usize {
151        // as unsigned wrap
152        let first = self.divide_mul;
153        let second = self.divide_add;
154
155        (((index * first) + second) >> 32 >> self.divide_shift) as usize
156    }
157
158    /// Get the data at the given index.
159    ///
160    /// # Panics
161    ///
162    /// This function will panic if the given index is greater than or equal to
163    /// the size of this storage.
164    pub fn get(&self, index: usize) -> u64 {
165        assert!(
166            index < self.size,
167            "Index {index} out of bounds (must be less than {})",
168            self.size
169        );
170
171        // 0 bit storage
172        if self.data.is_empty() {
173            return 0;
174        }
175
176        let cell_index = self.cell_index(index as u64);
177        let cell = &self.data[cell_index];
178        let bit_index = (index - cell_index * self.values_per_long) * self.bits;
179        (cell >> bit_index) & self.mask
180    }
181
182    pub fn get_and_set(&mut self, index: usize, value: u64) -> u64 {
183        // 0 bit storage
184        if self.data.is_empty() {
185            return 0;
186        }
187
188        debug_assert!(index < self.size);
189        debug_assert!(value <= self.mask);
190        let cell_index = self.cell_index(index as u64);
191        let cell = &mut self.data[cell_index];
192        let bit_index = (index - cell_index * self.values_per_long) * self.bits;
193        let old_value = (*cell >> (bit_index as u64)) & self.mask;
194        *cell = (*cell & !(self.mask << bit_index)) | ((value & self.mask) << bit_index);
195        old_value
196    }
197
198    pub fn set(&mut self, index: usize, value: u64) {
199        // 0 bit storage
200        if self.data.is_empty() {
201            return;
202        }
203
204        debug_assert!(index < self.size);
205        debug_assert!(value <= self.mask);
206        let cell_index = self.cell_index(index as u64);
207        let cell = &mut self.data[cell_index];
208        let bit_index = (index - cell_index * self.values_per_long) * self.bits;
209        *cell = (*cell & !(self.mask << bit_index)) | ((value & self.mask) << bit_index);
210    }
211
212    /// The number of entries.
213    #[inline]
214    pub fn size(&self) -> usize {
215        self.size
216    }
217
218    pub fn iter(&self) -> BitStorageIter {
219        BitStorageIter {
220            storage: self,
221            index: 0,
222        }
223    }
224}
225
226pub struct BitStorageIter<'a> {
227    storage: &'a BitStorage,
228    index: usize,
229}
230
231impl Iterator for BitStorageIter<'_> {
232    type Item = u64;
233
234    fn next(&mut self) -> Option<Self::Item> {
235        if self.index >= self.storage.size {
236            return None;
237        }
238
239        let value = self.storage.get(self.index);
240        self.index += 1;
241        Some(value)
242    }
243}
244
245#[cfg(test)]
246mod tests {
247    use super::*;
248
249    #[test]
250    fn test_wikivg_example() {
251        let data = [
252            1, 2, 2, 3, 4, 4, 5, 6, 6, 4, 8, 0, 7, 4, 3, 13, 15, 16, 9, 14, 10, 12, 0, 2,
253        ];
254        let compact_data: [u64; 2] = [0x0020863148418841, 0x01018A7260F68C87];
255        let storage = BitStorage::new(5, data.len(), Some(compact_data.to_vec())).unwrap();
256
257        for (i, expected) in data.iter().enumerate() {
258            assert_eq!(storage.get(i), *expected);
259        }
260    }
261}
azalea_world/bit_storage.rs

azalea_world/
bit_storage.rs
