1use fxhash::FxHashMap;
11use uni_common::core::id::Vid;
12
13#[derive(Debug, Clone)]
26pub struct IdMap {
27 slot_to_vid: Vec<Vid>,
29 vid_to_slot: Option<FxHashMap<Vid, u32>>,
31}
32
33impl IdMap {
34 pub fn new() -> Self {
36 Self {
37 slot_to_vid: Vec::new(),
38 vid_to_slot: Some(FxHashMap::default()),
39 }
40 }
41
42 pub fn with_capacity(capacity: usize) -> Self {
44 Self {
45 slot_to_vid: Vec::with_capacity(capacity),
46 vid_to_slot: Some(FxHashMap::with_capacity_and_hasher(
47 capacity,
48 Default::default(),
49 )),
50 }
51 }
52
53 pub fn insert(&mut self, vid: Vid) -> u32 {
58 let map = self
59 .vid_to_slot
60 .as_mut()
61 .expect("Cannot insert into compacted IdMap");
62
63 if let Some(&slot) = map.get(&vid) {
64 return slot;
65 }
66
67 let slot = self.slot_to_vid.len() as u32;
68 self.slot_to_vid.push(vid);
69 map.insert(vid, slot);
70 slot
71 }
72
73 #[inline]
77 pub fn to_slot(&self, vid: Vid) -> Option<u32> {
78 if let Some(map) = &self.vid_to_slot {
79 map.get(&vid).copied()
80 } else {
81 self.slot_to_vid.binary_search(&vid).ok().map(|i| i as u32)
82 }
83 }
84
85 #[inline]
87 pub fn to_vid(&self, slot: u32) -> Option<Vid> {
88 self.slot_to_vid.get(slot as usize).copied()
89 }
90
91 #[inline]
93 pub fn to_vid_unchecked(&self, slot: u32) -> Vid {
94 self.slot_to_vid[slot as usize]
95 }
96
97 #[inline]
99 pub fn len(&self) -> usize {
100 self.slot_to_vid.len()
101 }
102
103 #[inline]
105 pub fn is_empty(&self) -> bool {
106 self.slot_to_vid.is_empty()
107 }
108
109 #[inline]
111 pub fn contains(&self, vid: Vid) -> bool {
112 self.to_slot(vid).is_some()
113 }
114
115 pub fn compact(&mut self) {
120 self.vid_to_slot = None;
125 }
126
127 pub fn iter(&self) -> impl Iterator<Item = (u32, Vid)> + '_ {
129 self.slot_to_vid
130 .iter()
131 .enumerate()
132 .map(|(slot, &vid)| (slot as u32, vid))
133 }
134
135 pub fn memory_size(&self) -> usize {
137 let map_size = self.vid_to_slot.as_ref().map_or(0, |m| {
138 m.len() * (std::mem::size_of::<Vid>() + std::mem::size_of::<u32>() + 8)
139 });
140 self.slot_to_vid.len() * std::mem::size_of::<Vid>() + map_size
141 }
142}
143
144impl Default for IdMap {
145 fn default() -> Self {
146 Self::new()
147 }
148}
149
150impl FromIterator<Vid> for IdMap {
151 fn from_iter<I: IntoIterator<Item = Vid>>(iter: I) -> Self {
152 let iter = iter.into_iter();
153 let (lower, upper) = iter.size_hint();
154 let mut map = Self::with_capacity(upper.unwrap_or(lower));
155
156 for vid in iter {
157 map.insert(vid);
158 }
159
160 map
161 }
162}
163
164#[cfg(test)]
165mod tests {
166 use super::*;
167
168 #[test]
169 fn test_insert_and_lookup() {
170 let mut map = IdMap::new();
171
172 let vid1 = Vid::new(100);
173 let vid2 = Vid::new(200);
174 let vid3 = Vid::new(50);
175
176 assert_eq!(map.insert(vid1), 0);
177 assert_eq!(map.insert(vid2), 1);
178 assert_eq!(map.insert(vid3), 2);
179
180 assert_eq!(map.insert(vid1), 0);
182
183 assert_eq!(map.to_slot(vid1), Some(0));
184 assert_eq!(map.to_slot(vid2), Some(1));
185 assert_eq!(map.to_slot(vid3), Some(2));
186
187 assert_eq!(map.to_vid(0), Some(vid1));
188 assert_eq!(map.to_vid(1), Some(vid2));
189 assert_eq!(map.to_vid(2), Some(vid3));
190 assert_eq!(map.to_vid(3), None);
191 }
192
193 #[test]
194 fn test_compact() {
195 let mut map = IdMap::new();
196 let vid1 = Vid::new(100);
198 let vid2 = Vid::new(200);
199
200 map.insert(vid1);
201 map.insert(vid2);
202
203 assert!(map.vid_to_slot.is_some());
204
205 map.compact();
206 assert!(map.vid_to_slot.is_none());
207
208 assert_eq!(map.to_slot(vid1), Some(0));
209 assert_eq!(map.to_slot(vid2), Some(1));
210 assert_eq!(map.to_slot(Vid::new(150)), None);
211 }
212}