redb/tree_store/page_store/
bitmap.rs

Help
1use std::mem::size_of;
2
3const HEIGHT_OFFSET: usize = 0;
4const END_OFFSETS: usize = HEIGHT_OFFSET + size_of::<u32>();
5
6pub(crate) struct BtreeBitmap {
7    heights: Vec<U64GroupedBitmap>,
8}
9
10// Stores a 64-way bit-tree of allocated ids.
11//
12// Data structure format:
13// height: u32
14// layer_ends: array of u32, ending offset in bytes of layers.
15// layer data: u64s
16// ...consecutive layers. Except for the last level, all sub-trees of the root must be complete
17impl BtreeBitmap {
18    pub(crate) fn count_unset(&self) -> u32 {
19        self.get_level(self.get_height() - 1).count_unset()
20    }
21
22    pub(crate) fn has_unset(&self) -> bool {
23        self.get_level(self.get_height() - 1).any_unset()
24    }
25
26    pub(crate) fn get(&self, i: u32) -> bool {
27        self.get_level(self.get_height() - 1).get(i)
28    }
29
30    pub(crate) fn capacity(&self) -> u32 {
31        self.get_level(self.get_height() - 1).capacity()
32    }
33
34    pub(crate) fn len(&self) -> u32 {
35        self.get_level(self.get_height() - 1).len()
36    }
37
38    pub(crate) fn find_first_unset(&self) -> Option<u32> {
39        if let Some(mut entry) = self.get_level(0).first_unset(0, 64) {
40            let mut height = 0;
41
42            while height < self.get_height() - 1 {
43                height += 1;
44                entry *= 64;
45                entry = self
46                    .get_level(height)
47                    .first_unset(entry, entry + 64)
48                    .unwrap();
49            }
50
51            Some(entry)
52        } else {
53            None
54        }
55    }
56
57    fn get_level(&self, i: u32) -> &U64GroupedBitmap {
58        assert!(i < self.get_height());
59        &self.heights[i as usize]
60    }
61
62    fn get_height(&self) -> u32 {
63        self.heights.len().try_into().unwrap()
64    }
65
66    pub(crate) fn to_vec(&self) -> Vec<u8> {
67        let mut result = vec![];
68        let height: u32 = self.heights.len().try_into().unwrap();
69        result.extend(height.to_le_bytes());
70
71        let vecs: Vec<Vec<u8>> = self.heights.iter().map(|x| x.to_vec()).collect();
72        let mut data_offset = END_OFFSETS + self.heights.len() * size_of::<u32>();
73        let end_metadata = data_offset;
74        for serialized in &vecs {
75            data_offset += serialized.len();
76            let offset_u32: u32 = data_offset.try_into().unwrap();
77            result.extend(offset_u32.to_le_bytes());
78        }
79
80        assert_eq!(end_metadata, result.len());
81        for serialized in &vecs {
82            result.extend(serialized);
83        }
84
85        result
86    }
87
88    pub(crate) fn from_bytes(data: &[u8]) -> Self {
89        let height = u32::from_le_bytes(
90            data[HEIGHT_OFFSET..(HEIGHT_OFFSET + size_of::<u32>())]
91                .try_into()
92                .unwrap(),
93        );
94
95        let mut metadata = END_OFFSETS;
96        let mut data_start = END_OFFSETS + (height as usize) * size_of::<u32>();
97
98        let mut heights = vec![];
99        for _ in 0..height {
100            let data_end = u32::from_le_bytes(
101                data[metadata..(metadata + size_of::<u32>())]
102                    .try_into()
103                    .unwrap(),
104            ) as usize;
105            heights.push(U64GroupedBitmap::from_bytes(&data[data_start..data_end]));
106            data_start = data_end;
107            metadata += size_of::<u32>();
108        }
109
110        Self { heights }
111    }
112
113    // Initializes a new allocator, with no ids free
114    pub(crate) fn new(mut capacity: u32) -> Self {
115        let mut heights = vec![];
116
117        // Build from the leaf to root
118        loop {
119            heights.push(U64GroupedBitmap::new_full(capacity, capacity));
120            if capacity <= 64 {
121                break;
122            }
123            capacity = capacity.div_ceil(64);
124        }
125
126        // Reverse so that the root as index 0
127        heights.reverse();
128
129        Self { heights }
130    }
131
132    // Returns the first unset id, and sets it
133    pub(crate) fn alloc(&mut self) -> Option<u32> {
134        let entry = self.find_first_unset()?;
135        self.set(entry);
136        Some(entry)
137    }
138
139    pub(crate) fn set(&mut self, i: u32) {
140        let full = self.get_level_mut(self.get_height() - 1).set(i);
141        self.update_to_root(i, full);
142    }
143
144    pub(crate) fn clear(&mut self, i: u32) {
145        self.get_level_mut(self.get_height() - 1).clear(i);
146        self.update_to_root(i, false);
147    }
148
149    fn get_level_mut(&mut self, i: u32) -> &mut U64GroupedBitmap {
150        assert!(i < self.get_height());
151        &mut self.heights[i as usize]
152    }
153
154    // Recursively update to the root, starting at the given entry in the given height
155    // full parameter must be set if all bits in the entry's group of u64 are full
156    fn update_to_root(&mut self, i: u32, mut full: bool) {
157        if self.get_height() == 1 {
158            return;
159        }
160
161        let mut parent_height = self.get_height() - 2;
162        let mut parent_entry = i / 64;
163        loop {
164            full = if full {
165                self.get_level_mut(parent_height).set(parent_entry)
166            } else {
167                self.get_level_mut(parent_height).clear(parent_entry);
168                false
169            };
170
171            if parent_height == 0 {
172                break;
173            }
174            parent_height -= 1;
175            parent_entry /= 64;
176        }
177    }
178}
179
180pub(crate) struct U64GroupedBitmapIter<'a> {
181    len: u32,
182    data: &'a [u64],
183    data_index: usize,
184    current: u64,
185}
186
187impl<'a> U64GroupedBitmapIter<'a> {
188    fn new(len: u32, data: &'a [u64]) -> Self {
189        Self {
190            len,
191            data,
192            data_index: 0,
193            current: data[0],
194        }
195    }
196}
197
198impl Iterator for U64GroupedBitmapIter<'_> {
199    type Item = u32;
200
201    fn next(&mut self) -> Option<Self::Item> {
202        let data_index_u32: u32 = self.data_index.try_into().unwrap();
203        if data_index_u32 * u64::BITS >= self.len {
204            return None;
205        }
206        if self.current != 0 {
207            let mut result: u32 = self.data_index.try_into().unwrap();
208            result *= u64::BITS;
209            let bit = self.current.trailing_zeros();
210            result += bit;
211            self.current &= !U64GroupedBitmap::select_mask(bit as usize);
212            if result >= self.len {
213                return None;
214            }
215            return Some(result);
216        }
217        self.data_index += 1;
218        while self.data_index < self.data.len() {
219            let next = self.data[self.data_index];
220            if next != 0 {
221                self.current = next;
222                return self.next();
223            }
224            self.data_index += 1;
225        }
226        None
227    }
228}
229
230pub(crate) struct U64GroupedBitmapDifference<'a, 'b> {
231    data: &'a [u64],
232    exclusion_data: &'b [u64],
233    data_index: usize,
234    current: u64,
235}
236
237impl<'a, 'b> U64GroupedBitmapDifference<'a, 'b> {
238    fn new(data: &'a [u64], exclusion_data: &'b [u64]) -> Self {
239        assert_eq!(data.len(), exclusion_data.len());
240        let base = data[0];
241        let exclusion = exclusion_data[0];
242        Self {
243            data,
244            exclusion_data,
245            data_index: 0,
246            current: base & (!exclusion),
247        }
248    }
249}
250
251impl Iterator for U64GroupedBitmapDifference<'_, '_> {
252    type Item = u32;
253
254    fn next(&mut self) -> Option<Self::Item> {
255        if self.current != 0 {
256            let mut result: u32 = self.data_index.try_into().unwrap();
257            result *= u64::BITS;
258            let bit = self.current.trailing_zeros();
259            result += bit;
260            self.current &= !U64GroupedBitmap::select_mask(bit as usize);
261            return Some(result);
262        }
263        self.data_index += 1;
264        while self.data_index < self.data.len() {
265            let next = self.data[self.data_index];
266            let exclusion = *self.exclusion_data.get(self.data_index).unwrap_or(&0);
267            let next = next & (!exclusion);
268            if next != 0 {
269                self.current = next;
270                return self.next();
271            }
272            self.data_index += 1;
273        }
274        None
275    }
276}
277
278// A bitmap which groups consecutive groups of 64bits together
279pub(crate) struct U64GroupedBitmap {
280    len: u32,
281    data: Vec<u64>,
282}
283
284impl U64GroupedBitmap {
285    fn required_words(elements: u32) -> usize {
286        let words = elements.div_ceil(64);
287        words as usize
288    }
289
290    pub fn new_full(len: u32, capacity: u32) -> Self {
291        let data = vec![u64::MAX; Self::required_words(capacity)];
292        Self { len, data }
293    }
294
295    pub fn new_empty(len: u32, capacity: u32) -> Self {
296        let data = vec![0; Self::required_words(capacity)];
297        Self { len, data }
298    }
299
300    // Format:
301    // 4 bytes: number of elements
302    // n bytes: serialized groups
303    pub fn to_vec(&self) -> Vec<u8> {
304        let mut result = vec![];
305        result.extend(self.len.to_le_bytes());
306        for x in &self.data {
307            result.extend(x.to_le_bytes());
308        }
309        result
310    }
311
312    pub fn from_bytes(serialized: &[u8]) -> Self {
313        assert_eq!(0, (serialized.len() - size_of::<u32>()) % size_of::<u64>());
314        let mut data = vec![];
315        let len = u32::from_le_bytes(serialized[..size_of::<u32>()].try_into().unwrap());
316        let words = (serialized.len() - size_of::<u32>()) / size_of::<u64>();
317        for i in 0..words {
318            let start = size_of::<u32>() + i * size_of::<u64>();
319            let value = u64::from_le_bytes(
320                serialized[start..(start + size_of::<u64>())]
321                    .try_into()
322                    .unwrap(),
323            );
324            data.push(value);
325        }
326
327        Self { len, data }
328    }
329
330    fn data_index_of(bit: u32) -> (usize, usize) {
331        ((bit as usize) / 64, (bit as usize) % 64)
332    }
333
334    fn select_mask(bit: usize) -> u64 {
335        1u64 << (bit as u64)
336    }
337
338    fn count_unset(&self) -> u32 {
339        self.data.iter().map(|x| x.count_zeros()).sum()
340    }
341
342    pub fn difference<'a0, 'b0>(
343        &'a0 self,
344        exclusion: &'b0 U64GroupedBitmap,
345    ) -> U64GroupedBitmapDifference<'a0, 'b0> {
346        U64GroupedBitmapDifference::new(&self.data, &exclusion.data)
347    }
348
349    pub fn iter(&self) -> U64GroupedBitmapIter {
350        U64GroupedBitmapIter::new(self.len, &self.data)
351    }
352
353    pub fn capacity(&self) -> u32 {
354        let len: u32 = self.data.len().try_into().unwrap();
355        len * u64::BITS
356    }
357
358    fn any_unset(&self) -> bool {
359        self.data.iter().any(|x| x.count_zeros() > 0)
360    }
361
362    fn first_unset(&self, start_bit: u32, end_bit: u32) -> Option<u32> {
363        assert_eq!(end_bit, (start_bit - start_bit % 64) + 64);
364
365        let (index, bit) = Self::data_index_of(start_bit);
366        let mask = (1 << bit) - 1;
367        let group = self.data[index];
368        let group = group | mask;
369        match group.trailing_ones() {
370            64 => None,
371            x => Some(start_bit + x - u32::try_from(bit).unwrap()),
372        }
373    }
374
375    pub fn len(&self) -> u32 {
376        self.len
377    }
378
379    // TODO: thread this through up to BuddyAllocator
380    #[allow(dead_code)]
381    pub fn resize(&mut self, new_len: u32) {
382        assert!(new_len < self.capacity());
383        self.len = new_len;
384    }
385
386    pub fn get(&self, bit: u32) -> bool {
387        assert!(bit < self.len);
388        let (index, bit_index) = Self::data_index_of(bit);
389        let group = self.data[index];
390        group & U64GroupedBitmap::select_mask(bit_index) != 0
391    }
392
393    // Returns true iff the bit's group is all set
394    pub fn set(&mut self, bit: u32) -> bool {
395        assert!(bit < self.len);
396        let (index, bit_index) = Self::data_index_of(bit);
397        let mut group = self.data[index];
398        group |= Self::select_mask(bit_index);
399        self.data[index] = group;
400
401        group == u64::MAX
402    }
403
404    pub fn clear(&mut self, bit: u32) {
405        assert!(bit < self.len, "{bit} must be less than {}", self.len);
406        let (index, bit_index) = Self::data_index_of(bit);
407        self.data[index] &= !Self::select_mask(bit_index);
408    }
409}
410
411#[cfg(test)]
412mod test {
413    use crate::tree_store::page_store::bitmap::{BtreeBitmap, U64GroupedBitmap};
414    use rand::prelude::IteratorRandom;
415    use rand::rngs::StdRng;
416    use rand::{Rng, SeedableRng};
417    use std::collections::HashSet;
418
419    #[test]
420    fn alloc() {
421        let num_pages = 2;
422        let mut allocator = BtreeBitmap::new(num_pages);
423        for i in 0..num_pages {
424            allocator.clear(i);
425        }
426        for i in 0..num_pages {
427            assert_eq!(i, allocator.alloc().unwrap());
428        }
429        assert!(allocator.alloc().is_none());
430    }
431
432    #[test]
433    fn record_alloc() {
434        let mut allocator = BtreeBitmap::new(2);
435        allocator.clear(0);
436        allocator.clear(1);
437        allocator.set(0);
438        assert_eq!(1, allocator.alloc().unwrap());
439        assert!(allocator.alloc().is_none());
440    }
441
442    #[test]
443    fn free() {
444        let mut allocator = BtreeBitmap::new(1);
445        allocator.clear(0);
446        assert_eq!(0, allocator.alloc().unwrap());
447        assert!(allocator.alloc().is_none());
448        allocator.clear(0);
449        assert_eq!(0, allocator.alloc().unwrap());
450    }
451
452    #[test]
453    fn reuse_lowest() {
454        let num_pages = 65;
455        let mut allocator = BtreeBitmap::new(num_pages);
456        for i in 0..num_pages {
457            allocator.clear(i);
458        }
459        for i in 0..num_pages {
460            assert_eq!(i, allocator.alloc().unwrap());
461        }
462        allocator.clear(5);
463        allocator.clear(15);
464        assert_eq!(5, allocator.alloc().unwrap());
465        assert_eq!(15, allocator.alloc().unwrap());
466        assert!(allocator.alloc().is_none());
467    }
468
469    #[test]
470    fn all_space_used() {
471        let num_pages = 65;
472        let mut allocator = BtreeBitmap::new(num_pages);
473        for i in 0..num_pages {
474            allocator.clear(i);
475        }
476        // Allocate everything
477        while allocator.alloc().is_some() {}
478        // The last u64 must be used, since the leaf layer is compact
479        assert_eq!(
480            u64::MAX,
481            *allocator.heights.last().unwrap().data.last().unwrap()
482        );
483    }
484
485    #[test]
486    fn find_free() {
487        let num_pages = 129;
488        let mut allocator = BtreeBitmap::new(num_pages);
489        assert!(allocator.find_first_unset().is_none());
490        allocator.clear(128);
491        assert_eq!(allocator.find_first_unset().unwrap(), 128);
492        allocator.clear(65);
493        assert_eq!(allocator.find_first_unset().unwrap(), 65);
494        allocator.clear(8);
495        assert_eq!(allocator.find_first_unset().unwrap(), 8);
496        allocator.clear(0);
497        assert_eq!(allocator.find_first_unset().unwrap(), 0);
498    }
499
500    #[test]
501    fn iter() {
502        let num_pages = 129;
503        let mut bitmap = U64GroupedBitmap::new_empty(num_pages, num_pages);
504        let values = [0, 1, 33, 63, 64, 65, 90, 126, 127, 128];
505        for x in values {
506            bitmap.set(x);
507        }
508        for (i, x) in bitmap.iter().enumerate() {
509            assert_eq!(values[i], x);
510        }
511        assert_eq!(bitmap.iter().count(), values.len());
512    }
513
514    #[test]
515    fn random_pattern() {
516        let seed = rand::rng().random();
517        // Print the seed to debug for reproducibility, in case this test fails
518        println!("seed={seed}");
519        let mut rng = StdRng::seed_from_u64(seed);
520
521        let num_pages = rng.random_range(2..10000);
522        let mut allocator = BtreeBitmap::new(num_pages);
523        for i in 0..num_pages {
524            allocator.clear(i);
525        }
526        let mut allocated = HashSet::new();
527
528        for _ in 0..(num_pages * 2) {
529            if rng.random_bool(0.75) {
530                if let Some(page) = allocator.alloc() {
531                    allocated.insert(page);
532                } else {
533                    assert_eq!(allocated.len(), num_pages as usize);
534                }
535            } else if let Some(to_free) = allocated.iter().choose(&mut rng).copied() {
536                allocator.clear(to_free);
537                allocated.remove(&to_free);
538            }
539        }
540
541        for _ in allocated.len()..(num_pages as usize) {
542            allocator.alloc().unwrap();
543        }
544        assert!(allocator.alloc().is_none());
545
546        for i in 0..num_pages {
547            allocator.clear(i);
548        }
549
550        for _ in 0..num_pages {
551            allocator.alloc().unwrap();
552        }
553        assert!(allocator.alloc().is_none());
554    }
555}
redb/tree_store/page_store/bitmap.rs

redb/tree_store/page_store/
bitmap.rs
