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: Box<[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(
105        bits: usize,
106        size: usize,
107        data: Option<Box<[u64]>>,
108    ) -> Result<Self, BitStorageError> {
109        if let Some(data) = &data {
110            // 0 bit storage
111            if data.is_empty() {
112                return Ok(BitStorage {
113                    data: Box::new([]),
114                    bits,
115                    size,
116                    ..Default::default()
117                });
118            }
119        }
120
121        debug_assert!((1..=32).contains(&bits));
122
123        let values_per_long = 64 / bits;
124        let magic_index = values_per_long - 1;
125        let (divide_mul, divide_add, divide_shift) = MAGIC[magic_index];
126        let calculated_length = size.div_ceil(values_per_long);
127
128        let mask = (1 << bits) - 1;
129
130        let using_data = if let Some(data) = data {
131            if data.len() != calculated_length {
132                return Err(BitStorageError::InvalidLength {
133                    got: data.len(),
134                    expected: calculated_length,
135                });
136            }
137            data
138        } else {
139            vec![0; calculated_length].into()
140        };
141
142        Ok(BitStorage {
143            data: using_data,
144            bits,
145            mask,
146            size,
147            values_per_long,
148            divide_mul: divide_mul as u32 as u64,
149            divide_add: divide_add as u32 as u64,
150            divide_shift,
151        })
152    }
153
154    pub fn cell_index(&self, index: u64) -> usize {
155        // as unsigned wrap
156        let first = self.divide_mul;
157        let second = self.divide_add;
158
159        (((index * first) + second) >> 32 >> self.divide_shift) as usize
160    }
161
162    /// Get the data at the given index.
163    ///
164    /// # Panics
165    ///
166    /// This function will panic if the given index is greater than or equal to
167    /// the size of this storage.
168    pub fn get(&self, index: usize) -> u64 {
169        assert!(
170            index < self.size,
171            "Index {index} out of bounds (must be less than {})",
172            self.size
173        );
174
175        // 0 bit storage
176        if self.data.is_empty() {
177            return 0;
178        }
179
180        let cell_index = self.cell_index(index as u64);
181        let cell = &self.data[cell_index];
182        let bit_index = (index - cell_index * self.values_per_long) * self.bits;
183        (cell >> bit_index) & self.mask
184    }
185
186    pub fn get_and_set(&mut self, index: usize, value: u64) -> u64 {
187        // 0 bit storage
188        if self.data.is_empty() {
189            return 0;
190        }
191
192        debug_assert!(index < self.size);
193        debug_assert!(value <= self.mask);
194        let cell_index = self.cell_index(index as u64);
195        let cell = &mut self.data[cell_index];
196        let bit_index = (index - cell_index * self.values_per_long) * self.bits;
197        let old_value = (*cell >> (bit_index as u64)) & self.mask;
198        *cell = (*cell & !(self.mask << bit_index)) | ((value & self.mask) << bit_index);
199        old_value
200    }
201
202    pub fn set(&mut self, index: usize, value: u64) {
203        // 0 bit storage
204        if self.data.is_empty() {
205            return;
206        }
207
208        debug_assert!(index < self.size);
209        debug_assert!(value <= self.mask);
210        let cell_index = self.cell_index(index as u64);
211        let cell = &mut self.data[cell_index];
212        let bit_index = (index - cell_index * self.values_per_long) * self.bits;
213        *cell = (*cell & !(self.mask << bit_index)) | ((value & self.mask) << bit_index);
214    }
215
216    /// The number of entries.
217    #[inline]
218    pub fn size(&self) -> usize {
219        self.size
220    }
221
222    pub fn iter(&self) -> BitStorageIter {
223        BitStorageIter {
224            storage: self,
225            index: 0,
226        }
227    }
228}
229
230pub struct BitStorageIter<'a> {
231    storage: &'a BitStorage,
232    index: usize,
233}
234
235impl Iterator for BitStorageIter<'_> {
236    type Item = u64;
237
238    fn next(&mut self) -> Option<Self::Item> {
239        if self.index >= self.storage.size {
240            return None;
241        }
242
243        let value = self.storage.get(self.index);
244        self.index += 1;
245        Some(value)
246    }
247}
248
249#[cfg(test)]
250mod tests {
251    use super::*;
252
253    #[test]
254    fn test_wikivg_example() {
255        let data = [
256            1, 2, 2, 3, 4, 4, 5, 6, 6, 4, 8, 0, 7, 4, 3, 13, 15, 16, 9, 14, 10, 12, 0, 2,
257        ];
258        let compact_data: [u64; 2] = [0x0020863148418841, 0x01018A7260F68C87];
259        let storage = BitStorage::new(5, data.len(), Some(Box::new(compact_data))).unwrap();
260
261        for (i, expected) in data.iter().enumerate() {
262            assert_eq!(storage.get(i), *expected);
263        }
264    }
265}
azalea_world/bit_storage.rs

azalea_world/
bit_storage.rs
