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 + 63) / 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<'a> Iterator for U64GroupedBitmapIter<'a> {
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<'a, 'b> Iterator for U64GroupedBitmapDifference<'a, 'b> {
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[self.data_index];
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 + 63) / 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    #[allow(dead_code)]
350    pub fn iter(&self) -> U64GroupedBitmapIter {
351        U64GroupedBitmapIter::new(self.len, &self.data)
352    }
353
354    pub fn capacity(&self) -> u32 {
355        let len: u32 = self.data.len().try_into().unwrap();
356        len * u64::BITS
357    }
358
359    fn any_unset(&self) -> bool {
360        self.data.iter().any(|x| x.count_zeros() > 0)
361    }
362
363    fn first_unset(&self, start_bit: u32, end_bit: u32) -> Option<u32> {
364        assert_eq!(end_bit, (start_bit - start_bit % 64) + 64);
365
366        let (index, bit) = Self::data_index_of(start_bit);
367        let mask = (1 << bit) - 1;
368        let group = self.data[index];
369        let group = group | mask;
370        match group.trailing_ones() {
371            64 => None,
372            x => Some(start_bit + x - u32::try_from(bit).unwrap()),
373        }
374    }
375
376    pub fn len(&self) -> u32 {
377        self.len
378    }
379
380    // TODO: thread this through up to BuddyAllocator
381    #[allow(dead_code)]
382    pub fn resize(&mut self, new_len: u32) {
383        assert!(new_len < self.capacity());
384        self.len = new_len;
385    }
386
387    pub fn get(&self, bit: u32) -> bool {
388        assert!(bit < self.len);
389        let (index, bit_index) = Self::data_index_of(bit);
390        let group = self.data[index];
391        group & U64GroupedBitmap::select_mask(bit_index) != 0
392    }
393
394    // Returns true iff the bit's group is all set
395    pub fn set(&mut self, bit: u32) -> bool {
396        assert!(bit < self.len);
397        let (index, bit_index) = Self::data_index_of(bit);
398        let mut group = self.data[index];
399        group |= Self::select_mask(bit_index);
400        self.data[index] = group;
401
402        group == u64::MAX
403    }
404
405    pub fn clear(&mut self, bit: u32) {
406        assert!(bit < self.len, "{bit} must be less than {}", self.len);
407        let (index, bit_index) = Self::data_index_of(bit);
408        self.data[index] &= !Self::select_mask(bit_index);
409    }
410}
411
412#[cfg(test)]
413mod test {
414    use crate::tree_store::page_store::bitmap::{BtreeBitmap, U64GroupedBitmap};
415    use rand::prelude::IteratorRandom;
416    use rand::rngs::StdRng;
417    use rand::{Rng, SeedableRng};
418    use std::collections::HashSet;
419
420    #[test]
421    fn alloc() {
422        let num_pages = 2;
423        let mut allocator = BtreeBitmap::new(num_pages);
424        for i in 0..num_pages {
425            allocator.clear(i);
426        }
427        for i in 0..num_pages {
428            assert_eq!(i, allocator.alloc().unwrap());
429        }
430        assert!(allocator.alloc().is_none());
431    }
432
433    #[test]
434    fn record_alloc() {
435        let mut allocator = BtreeBitmap::new(2);
436        allocator.clear(0);
437        allocator.clear(1);
438        allocator.set(0);
439        assert_eq!(1, allocator.alloc().unwrap());
440        assert!(allocator.alloc().is_none());
441    }
442
443    #[test]
444    fn free() {
445        let mut allocator = BtreeBitmap::new(1);
446        allocator.clear(0);
447        assert_eq!(0, allocator.alloc().unwrap());
448        assert!(allocator.alloc().is_none());
449        allocator.clear(0);
450        assert_eq!(0, allocator.alloc().unwrap());
451    }
452
453    #[test]
454    fn reuse_lowest() {
455        let num_pages = 65;
456        let mut allocator = BtreeBitmap::new(num_pages);
457        for i in 0..num_pages {
458            allocator.clear(i);
459        }
460        for i in 0..num_pages {
461            assert_eq!(i, allocator.alloc().unwrap());
462        }
463        allocator.clear(5);
464        allocator.clear(15);
465        assert_eq!(5, allocator.alloc().unwrap());
466        assert_eq!(15, allocator.alloc().unwrap());
467        assert!(allocator.alloc().is_none());
468    }
469
470    #[test]
471    fn all_space_used() {
472        let num_pages = 65;
473        let mut allocator = BtreeBitmap::new(num_pages);
474        for i in 0..num_pages {
475            allocator.clear(i);
476        }
477        // Allocate everything
478        while allocator.alloc().is_some() {}
479        // The last u64 must be used, since the leaf layer is compact
480        assert_eq!(
481            u64::MAX,
482            *allocator.heights.last().unwrap().data.last().unwrap()
483        );
484    }
485
486    #[test]
487    fn find_free() {
488        let num_pages = 129;
489        let mut allocator = BtreeBitmap::new(num_pages);
490        assert!(allocator.find_first_unset().is_none());
491        allocator.clear(128);
492        assert_eq!(allocator.find_first_unset().unwrap(), 128);
493        allocator.clear(65);
494        assert_eq!(allocator.find_first_unset().unwrap(), 65);
495        allocator.clear(8);
496        assert_eq!(allocator.find_first_unset().unwrap(), 8);
497        allocator.clear(0);
498        assert_eq!(allocator.find_first_unset().unwrap(), 0);
499    }
500
501    #[test]
502    fn iter() {
503        let num_pages = 129;
504        let mut bitmap = U64GroupedBitmap::new_empty(num_pages, num_pages);
505        let values = [0, 1, 33, 63, 64, 65, 90, 126, 127, 128];
506        for x in values {
507            bitmap.set(x);
508        }
509        for (i, x) in bitmap.iter().enumerate() {
510            assert_eq!(values[i], x);
511        }
512        assert_eq!(bitmap.iter().count(), values.len());
513    }
514
515    #[test]
516    fn random_pattern() {
517        let seed = rand::thread_rng().gen();
518        // Print the seed to debug for reproducibility, in case this test fails
519        println!("seed={seed}");
520        let mut rng = StdRng::seed_from_u64(seed);
521
522        let num_pages = rng.gen_range(2..10000);
523        let mut allocator = BtreeBitmap::new(num_pages);
524        for i in 0..num_pages {
525            allocator.clear(i);
526        }
527        let mut allocated = HashSet::new();
528
529        for _ in 0..(num_pages * 2) {
530            if rng.gen_bool(0.75) {
531                if let Some(page) = allocator.alloc() {
532                    allocated.insert(page);
533                } else {
534                    assert_eq!(allocated.len(), num_pages as usize);
535                }
536            } else if let Some(to_free) = allocated.iter().choose(&mut rng).copied() {
537                allocator.clear(to_free);
538                allocated.remove(&to_free);
539            }
540        }
541
542        for _ in allocated.len()..(num_pages as usize) {
543            allocator.alloc().unwrap();
544        }
545        assert!(allocator.alloc().is_none());
546
547        for i in 0..num_pages {
548            allocator.clear(i);
549        }
550
551        for _ in 0..num_pages {
552            allocator.alloc().unwrap();
553        }
554        assert!(allocator.alloc().is_none());
555    }
556}
redb/tree_store/page_store/bitmap.rs

redb/tree_store/page_store/
bitmap.rs
