1use crate::packed_map::PackedMap;
6use crate::{PackedItem, PackedVec};
7use std::borrow::Borrow;
8use std::collections::BTreeMap;
9
10pub struct PackedMapBuilder<K: ?Sized + Ord + PackedItem, V>
39where
40 K: ToOwned,
41 K::Owned: Ord,
42{
43 packed: PackedMap<K, V>,
44 unpacked: BTreeMap<K::Owned, V>,
45}
46
47impl<K, V> PackedMapBuilder<K, V>
48where
49 K: ?Sized + Ord + PackedItem + ToOwned,
50 K::Owned: Ord,
51{
52 pub fn new() -> Self {
54 Self { packed: PackedMap::new(), unpacked: BTreeMap::new() }
55 }
56
57 pub fn with_capacity(element_capacity: usize, buffer_capacity: usize) -> Self {
64 Self {
65 packed: PackedMap::with_capacity(element_capacity, buffer_capacity),
66 unpacked: BTreeMap::new(),
67 }
68 }
69
70 pub fn contains_key(&self, key: &K) -> bool {
72 self.packed.contains_key(key) || self.unpacked.contains_key(key)
73 }
74
75 pub fn get(&self, key: &K) -> Option<&V> {
77 self.packed.get(key).or_else(|| self.unpacked.get(key))
78 }
79
80 pub fn insert(&mut self, key: &K, value: V) -> Option<V> {
87 match self.packed.append_or_update(key, value) {
88 Ok(old_value) => old_value,
89 Err(value) => self.unpacked.insert(key.to_owned(), value),
90 }
91 }
92
93 pub fn build(self) -> PackedMap<K, V> {
95 if self.unpacked.is_empty() { self.packed } else { merge(self.packed, self.unpacked) }
96 }
97}
98
99impl<K, V> Default for PackedMapBuilder<K, V>
100where
101 K: ?Sized + Ord + PackedItem + ToOwned,
102 K::Owned: Ord,
103{
104 fn default() -> Self {
105 Self::new()
106 }
107}
108
109pub(crate) fn merge<K, V>(
110 mut packed: PackedMap<K, V>,
111 unpacked: BTreeMap<K::Owned, V>,
112) -> PackedMap<K, V>
113where
114 K: ?Sized + Ord + PackedItem + ToOwned,
115 K::Owned: Ord,
116{
117 let len1 = packed.keys.len();
118 let len2 = unpacked.len();
119
120 let buffer_len1 = packed.buffer_len();
121 let buffer_len2: usize = unpacked.keys().map(|key| key.borrow().as_bytes().len()).sum();
122
123 let mut out_keys = PackedVec::with_capacity(len1 + len2, buffer_len1 + buffer_len2);
124 let mut out_values = Vec::with_capacity(len1 + len2);
125
126 let mut drain1 = packed.drain();
127 let mut drain2 = unpacked.into_iter();
128
129 let mut next1 = drain1.next();
130 let mut next2 = drain2.next();
131
132 loop {
133 match (next1, next2) {
134 (Some((k1, v1)), Some((k2, v2))) => match k1.cmp(k2.borrow()) {
135 std::cmp::Ordering::Less => {
136 out_keys.push(k1);
137 out_values.push(v1);
138 next1 = drain1.next();
139 next2 = Some((k2, v2));
140 }
141 std::cmp::Ordering::Greater => {
142 out_keys.push(k2.borrow());
143 out_values.push(v2);
144 next1 = Some((k1, v1));
145 next2 = drain2.next();
146 }
147 std::cmp::Ordering::Equal => {
148 out_keys.push(k1);
149 out_values.push(v2);
150 next1 = drain1.next();
151 next2 = drain2.next();
152 }
153 },
154 (Some((k1, v1)), None) => {
155 out_keys.push(k1);
156 out_values.push(v1);
157 while let Some((k1, v1)) = drain1.next() {
158 out_keys.push(k1);
159 out_values.push(v1);
160 }
161 break;
162 }
163 (None, Some((k2, v2))) => {
164 out_keys.push(k2.borrow());
165 out_values.push(v2);
166 while let Some((k2, v2)) = drain2.next() {
167 out_keys.push(k2.borrow());
168 out_values.push(v2);
169 }
170 break;
171 }
172 (None, None) => break,
173 }
174 }
175
176 PackedMap::from_parts(out_keys, out_values)
177}
178
179#[cfg(test)]
180mod tests {
181 use super::*;
182
183 #[test]
184 fn test_packed_map_builder_lsm() {
185 let mut builder = PackedMapBuilder::new();
186 for i in 0..129 {
187 let s = format!("{:03}", 128 - i);
188 builder.insert(s.as_str(), i); }
190 let map = builder.build();
191 assert_eq!(map.element_len(), 129);
192 assert_eq!(map.get("000"), Some(&128));
193 assert_eq!(map.get(&format!("{}", 128)), Some(&0));
194 }
195
196 #[test]
197 fn test_builder_insert() {
198 let mut builder = PackedMapBuilder::new();
199
200 assert_eq!(builder.insert("a", 1), None);
201 assert_eq!(builder.insert("a", 2), Some(1));
202
203 builder.insert("b", 3);
204 builder.insert("c", 4);
205
206 assert_eq!(builder.insert("b", 3), Some(3));
207 assert_eq!(builder.insert("c", 4), Some(4));
208
209 assert!(!builder.contains_key("aa"));
210 assert_eq!(builder.get("a"), Some(&2));
211 assert_eq!(builder.get("aa"), None);
212
213 builder.insert("e", 6);
214 builder.insert("d", 5);
215
216 let map = builder.build();
217 assert_eq!(map.get("aa"), None);
218 assert_eq!(map.get("a"), Some(&2));
219 assert_eq!(map.get("b"), Some(&3));
220 assert_eq!(map.get("c"), Some(&4));
221 assert_eq!(map.get("d"), Some(&5));
222 assert_eq!(map.get("e"), Some(&6));
223 }
224
225 #[test]
226 fn test_insert_out_of_order_sorts() {
227 let mut builder = PackedMapBuilder::new();
228
229 builder.insert("a", 1);
230 builder.insert("c", 3);
231 builder.insert("b", 2);
232
233 let map = builder.build();
234 let collected: Vec<_> = map.iter().collect();
235 assert_eq!(collected, vec![("a", &1), ("b", &2), ("c", &3)]);
236
237 let keys: Vec<_> = map.keys.iter().collect();
238 assert_eq!(keys, vec!["a", "b", "c"]);
239 }
240
241 #[test]
242 fn test_insert_comprehensive() {
243 let mut builder = PackedMapBuilder::new();
244
245 builder.insert("a", 1);
246 builder.insert("c", 3);
247 builder.insert("b", 2);
248
249 assert_eq!(builder.insert("a", 10), Some(1));
250 assert_eq!(builder.insert("a", 10), Some(10));
251
252 assert_eq!(builder.insert("b", 20), Some(2));
253 assert_eq!(builder.insert("b", 20), Some(20));
254
255 assert!(!builder.contains_key("d"));
256
257 assert_eq!(builder.insert("d", 4), None);
258 assert_eq!(builder.insert("d", 4), Some(4));
259
260 let map = builder.build();
261 assert_eq!(map.get("a"), Some(&10));
262 assert_eq!(map.get("b"), Some(&20));
263 assert_eq!(map.get("c"), Some(&3));
264 assert_eq!(map.get("d"), Some(&4));
265 }
266
267 #[test]
268 fn test_packed_map_merge() {
269 let mut map1 = PackedMap::new();
270 map1.append_or_update("a", 1).unwrap();
271 map1.append_or_update("c", 3).unwrap();
272
273 let mut map2 = BTreeMap::new();
274 assert_eq!(map2.insert("b".into(), 2), None);
275 assert_eq!(map2.insert("c".into(), 30), None);
276
277 let merged = merge(map1, map2);
278 assert_eq!(merged.element_len(), 3);
279 assert_eq!(merged.get("a"), Some(&1));
280 assert_eq!(merged.get("b"), Some(&2));
281 assert_eq!(merged.get("c"), Some(&30));
282 }
283
284 #[derive(
285 Debug, Clone, Copy, Eq, zerocopy::IntoBytes, zerocopy::Immutable, zerocopy::Unaligned,
286 )]
287 #[repr(C)]
288 struct TestKey {
289 id: u8,
290 metadata: u8,
291 }
292
293 impl PartialEq for TestKey {
294 fn eq(&self, other: &Self) -> bool {
295 self.id == other.id
296 }
297 }
298
299 impl Ord for TestKey {
300 fn cmp(&self, other: &Self) -> std::cmp::Ordering {
301 self.id.cmp(&other.id)
302 }
303 }
304
305 impl PartialOrd for TestKey {
306 fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
307 Some(self.cmp(other))
308 }
309 }
310
311 impl PackedItem for TestKey {
312 unsafe fn from_bytes(bytes: &[u8]) -> &Self {
313 unsafe { &*(bytes.as_ptr() as *const TestKey) }
314 }
315 }
316
317 #[test]
318 fn test_insert_keeps_old_key_packed() {
319 let mut builder = PackedMapBuilder::new();
320 let key1 = TestKey { id: 1, metadata: 10 };
321 let key2 = TestKey { id: 1, metadata: 20 };
322
323 builder.insert(&key1, 1);
324 builder.insert(&key2, 2); let map = builder.build();
327 let (k, v) = map.iter().next().unwrap();
328 assert_eq!(k.metadata, 10);
329 assert_eq!(*v, 2);
330 }
331
332 #[test]
333 fn test_insert_keeps_old_key_unpacked() {
334 let mut builder = PackedMapBuilder::new();
335 let key1 = TestKey { id: 2, metadata: 10 };
336 let key2 = TestKey { id: 1, metadata: 20 };
337 let key3 = TestKey { id: 1, metadata: 30 };
338
339 builder.insert(&key1, 1);
340 builder.insert(&key2, 2); builder.insert(&key3, 3); let map = builder.build();
344 let mut iter = map.iter();
345 let (k1, v1) = iter.next().unwrap();
346 assert_eq!(k1.id, 1);
347 assert_eq!(k1.metadata, 20); assert_eq!(*v1, 3); }
350}