aurora_db/storage/
index_storage.rs

1use roaring::RoaringBitmap;
2use serde::{Deserialize, Serialize};
3
4#[derive(Debug, Clone, Serialize, Deserialize)]
5pub enum IndexStorage {
6    Single(u32),
7    List(Vec<u32>),
8    Bitmap(RoaringBitmap),
9}
10
11pub enum IndexStorageIter<'a> {
12    Single(std::iter::Once<u32>),
13    List(std::slice::Iter<'a, u32>),
14    Bitmap(roaring::bitmap::Iter<'a>),
15}
16
17impl<'a> Iterator for IndexStorageIter<'a> {
18    type Item = u32;
19    fn next(&mut self) -> Option<Self::Item> {
20        match self {
21            IndexStorageIter::Single(i) => i.next(),
22            IndexStorageIter::List(i) => i.next().copied(),
23            IndexStorageIter::Bitmap(i) => i.next(),
24        }
25    }
26}
27
28impl IndexStorage {
29    pub fn new() -> Self {
30        IndexStorage::Bitmap(RoaringBitmap::new())
31    }
32
33    pub fn add(&mut self, id: u32) {
34        match self {
35            IndexStorage::Single(old_id) => {
36                if *old_id != id {
37                    *self = IndexStorage::List(vec![*old_id, id]);
38                }
39            }
40            IndexStorage::List(list) => {
41                if !list.contains(&id) {
42                    list.push(id);
43                    if list.len() > 16 {
44                        let mut bitmap = RoaringBitmap::new();
45                        for &i in list.iter() {
46                            bitmap.insert(i);
47                        }
48                        bitmap.insert(id);
49                        *self = IndexStorage::Bitmap(bitmap);
50                    }
51                }
52            }
53            IndexStorage::Bitmap(bitmap) => {
54                bitmap.insert(id);
55            }
56        }
57    }
58
59    pub fn remove(&mut self, id: u32) {
60        match self {
61            IndexStorage::Single(old_id) => {
62                if *old_id == id {
63                    *self = IndexStorage::Bitmap(RoaringBitmap::new());
64                }
65            }
66            IndexStorage::List(list) => {
67                list.retain(|&i| i != id);
68            }
69            IndexStorage::Bitmap(bitmap) => {
70                bitmap.remove(id);
71            }
72        }
73    }
74
75    pub fn len(&self) -> usize {
76        match self {
77            IndexStorage::Single(_) => 1,
78            IndexStorage::List(list) => list.len(),
79            IndexStorage::Bitmap(bitmap) => bitmap.len() as usize,
80        }
81    }
82
83    pub fn is_empty(&self) -> bool {
84        self.len() == 0
85    }
86
87    pub fn to_bitmap(&self) -> RoaringBitmap {
88        match self {
89            IndexStorage::Single(id) => {
90                let mut b = RoaringBitmap::new();
91                b.insert(*id);
92                b
93            }
94            IndexStorage::List(list) => {
95                let mut b = RoaringBitmap::new();
96                for &i in list {
97                    b.insert(i);
98                }
99                b
100            }
101            IndexStorage::Bitmap(bitmap) => bitmap.clone(),
102        }
103    }
104
105    pub fn iter(&self) -> IndexStorageIter<'_> {
106        match self {
107            IndexStorage::Single(id) => IndexStorageIter::Single(std::iter::once(*id)),
108            IndexStorage::List(list) => IndexStorageIter::List(list.iter()),
109            IndexStorage::Bitmap(bitmap) => IndexStorageIter::Bitmap(bitmap.iter()),
110        }
111    }
112}
113
114impl Default for IndexStorage {
115    fn default() -> Self {
116        Self::new()
117    }
118}
aurora_db/storage/index_storage.rs

aurora_db/storage/
index_storage.rs
