re_byte_size/
bookkeeping_btreemap.rs

1//! A `BTreeMap` wrapper with continuous size bookkeeping for instant `SizeBytes` queries.
2
3use std::collections::BTreeMap;
4
5use crate::SizeBytes;
6use crate::std_sizes::btree_heap_size;
7
8/// A [`BTreeMap`] wrapper with O(1) size queries via continuous bookkeeping.
9///
10/// Tracks the total size of keys and values, making [`SizeBytes::heap_size_bytes`] instant.
11#[derive(Clone, Debug, PartialEq, Eq)]
12pub struct BookkeepingBTreeMap<K, V> {
13    map: BTreeMap<K, V>,
14
15    /// The total heap uses internally by (non-POD) keys and values.
16    ///
17    /// For Plain Old Data types, this will always be zero.
18    ///
19    /// This does NOT include the size of the actual btree data-structure itself.
20    /// We estimate that with [`btree_heap_size`].
21    kv_heap_size_bytes: u64,
22}
23
24impl<K: Ord + SizeBytes, V: SizeBytes> SizeBytes for BookkeepingBTreeMap<K, V> {
25    #[inline]
26    fn heap_size_bytes(&self) -> u64 {
27        btree_heap_size(self.len(), size_of::<K>() + size_of::<V>()) + self.kv_heap_size_bytes
28    }
29}
30
31impl<K, V> Default for BookkeepingBTreeMap<K, V>
32where
33    K: Ord + SizeBytes,
34    V: SizeBytes,
35{
36    fn default() -> Self {
37        Self::new()
38    }
39}
40
41impl<K, V> BookkeepingBTreeMap<K, V>
42where
43    K: Ord + SizeBytes,
44    V: SizeBytes,
45{
46    /// Creates an empty `BookkeepingBTreeMap`.
47    #[inline]
48    pub fn new() -> Self {
49        Self {
50            map: BTreeMap::new(),
51            kv_heap_size_bytes: 0,
52        }
53    }
54
55    /// Returns `true` if the map contains no elements.
56    #[inline]
57    pub fn is_empty(&self) -> bool {
58        self.map.is_empty()
59    }
60
61    /// Returns the number of elements in the map.
62    #[inline]
63    pub fn len(&self) -> usize {
64        self.map.len()
65    }
66
67    /// Returns an iterator over the key-value pairs of the map, in sorted order by key.
68    #[inline]
69    pub fn iter(&self) -> std::collections::btree_map::Iter<'_, K, V> {
70        self.map.iter()
71    }
72
73    /// Mutates an entry, inserting `default_value` if the key doesn't exist.
74    ///
75    /// Size changes are tracked automatically.
76    pub fn mutate_entry(&mut self, key: K, default_value: V, mut mutator: impl FnMut(&mut V)) {
77        use std::collections::btree_map::Entry;
78
79        match self.map.entry(key) {
80            Entry::Vacant(vacant) => {
81                self.kv_heap_size_bytes += vacant.key().heap_size_bytes();
82                let value_ref = vacant.insert(default_value);
83                mutator(value_ref);
84                self.kv_heap_size_bytes += value_ref.heap_size_bytes();
85            }
86            Entry::Occupied(mut occupied) => {
87                self.kv_heap_size_bytes -= occupied.get().heap_size_bytes();
88                mutator(occupied.get_mut());
89                self.kv_heap_size_bytes += occupied.get().heap_size_bytes();
90            }
91        }
92    }
93
94    /// Finds and mutates the last entry smaller or equal to the given `key`.
95    ///
96    /// Equivalent to `.range_mut(..=key).next_back()` but with automatic size tracking.
97    /// Returns the mutator's return value, or `None` if no entry exists <= `key`.
98    pub fn mutate_latest_at<Ret>(
99        &mut self,
100        key: &K,
101        mut mutator: impl FnMut(&K, &mut V) -> Ret,
102    ) -> Option<Ret>
103    where
104        K: Clone,
105    {
106        let (key, value) = self.map.range_mut(..=key).next_back()?;
107        self.kv_heap_size_bytes -= value.heap_size_bytes();
108        let ret = mutator(key, value);
109        self.kv_heap_size_bytes += value.heap_size_bytes();
110        Some(ret)
111    }
112
113    /// Returns a reference to the inner [`BTreeMap`].
114    #[inline]
115    pub fn as_map(&self) -> &BTreeMap<K, V> {
116        &self.map
117    }
118
119    /// Inserts a key-value pair, returning the old value if the key was present.
120    pub fn insert(&mut self, key: K, value: V) -> Option<V> {
121        let new_key_size = key.heap_size_bytes();
122        let new_value_size = value.heap_size_bytes();
123
124        let old_value = self.map.insert(key, value);
125
126        if let Some(old_value) = &old_value {
127            // We're replacing an existing value, but the key remains the same:
128            self.kv_heap_size_bytes -= old_value.heap_size_bytes();
129            self.kv_heap_size_bytes += new_value_size;
130        } else {
131            // New key-value pair - add both sizes:
132            self.kv_heap_size_bytes += new_key_size + new_value_size;
133        }
134
135        old_value
136    }
137
138    /// Removes a key, returning its value if it was present.
139    pub fn remove(&mut self, key: &K) -> Option<V> {
140        if let Some(value) = self.map.remove(key) {
141            self.kv_heap_size_bytes -= key.heap_size_bytes();
142            self.kv_heap_size_bytes -= value.heap_size_bytes();
143            Some(value)
144        } else {
145            None
146        }
147    }
148
149    /// Extends the map with key-value pairs from an iterator.
150    pub fn extend<I>(&mut self, iter: I)
151    where
152        I: IntoIterator<Item = (K, V)>,
153    {
154        for (key, value) in iter {
155            self.insert(key, value);
156        }
157    }
158}
159
160impl<'a, K, V> IntoIterator for &'a BookkeepingBTreeMap<K, V>
161where
162    K: Ord + SizeBytes,
163    V: SizeBytes,
164{
165    type Item = (&'a K, &'a V);
166    type IntoIter = std::collections::btree_map::Iter<'a, K, V>;
167
168    fn into_iter(self) -> Self::IntoIter {
169        self.iter()
170    }
171}
172
173#[cfg(test)]
174mod tests {
175    use super::*;
176
177    fn heap_size_of_map<K: Ord + SizeBytes, V: SizeBytes>(map: &BookkeepingBTreeMap<K, V>) -> u64 {
178        let entry_heap_size: u64 = map
179            .iter()
180            .map(|(k, v)| k.heap_size_bytes() + v.heap_size_bytes())
181            .sum();
182        crate::std_sizes::btree_heap_size(map.len(), size_of::<K>() + size_of::<V>())
183            + entry_heap_size
184    }
185
186    #[test]
187    fn test_multiple_entries() {
188        let mut map: BookkeepingBTreeMap<u64, String> = BookkeepingBTreeMap::new();
189        assert_eq!(map.heap_size_bytes(), 0);
190
191        map.insert(1, "one".to_owned());
192        assert_eq!(map.heap_size_bytes(), heap_size_of_map(&map));
193
194        map.insert(2, "two".to_owned());
195        assert_eq!(map.heap_size_bytes(), heap_size_of_map(&map));
196
197        map.insert(2, "two, but now it is different".to_owned());
198        assert_eq!(map.heap_size_bytes(), heap_size_of_map(&map));
199
200        map.remove(&1);
201        assert_eq!(map.heap_size_bytes(), heap_size_of_map(&map));
202
203        map.remove(&2);
204        assert_eq!(map.heap_size_bytes(), heap_size_of_map(&map));
205        assert_eq!(map.heap_size_bytes(), 0);
206    }
207
208    #[test]
209    fn test_new_and_default() {
210        let map1: BookkeepingBTreeMap<u64, String> = BookkeepingBTreeMap::new();
211        let map2: BookkeepingBTreeMap<u64, String> = BookkeepingBTreeMap::default();
212
213        assert_eq!(map1.heap_size_bytes(), 0);
214        assert_eq!(map2.heap_size_bytes(), 0);
215        assert!(map1.is_empty());
216        assert!(map2.is_empty());
217        assert_eq!(map1.len(), 0);
218        assert_eq!(map2.len(), 0);
219    }
220
221    #[test]
222    fn test_insert_bookkeeping() {
223        let mut map: BookkeepingBTreeMap<u64, String> = BookkeepingBTreeMap::new();
224
225        let old = map.insert(1, "hello".to_owned());
226        assert_eq!(old, None);
227        assert_eq!(map.len(), 1);
228        assert_eq!(map.heap_size_bytes(), heap_size_of_map(&map));
229
230        let old = map.insert(2, "world".to_owned());
231        assert_eq!(old, None);
232        assert_eq!(map.len(), 2);
233        assert_eq!(map.heap_size_bytes(), heap_size_of_map(&map));
234
235        let old = map.insert(1, "hello, this is much longer!".to_owned());
236        assert_eq!(old, Some("hello".to_owned()));
237        assert_eq!(map.len(), 2);
238        assert_eq!(map.heap_size_bytes(), heap_size_of_map(&map));
239
240        let old = map.insert(1, "hi".to_owned());
241        assert_eq!(old, Some("hello, this is much longer!".to_owned()));
242        assert_eq!(map.len(), 2);
243        assert_eq!(map.heap_size_bytes(), heap_size_of_map(&map));
244    }
245
246    #[test]
247    fn test_remove_bookkeeping() {
248        let mut map: BookkeepingBTreeMap<u64, String> = BookkeepingBTreeMap::new();
249
250        map.insert(1, "one".to_owned());
251        map.insert(2, "two".to_owned());
252        map.insert(3, "three".to_owned());
253        assert_eq!(map.heap_size_bytes(), heap_size_of_map(&map));
254
255        let removed = map.remove(&2);
256        assert_eq!(removed, Some("two".to_owned()));
257        assert_eq!(map.len(), 2);
258        assert_eq!(map.heap_size_bytes(), heap_size_of_map(&map));
259
260        let removed = map.remove(&99);
261        assert_eq!(removed, None);
262        assert_eq!(map.len(), 2);
263        assert_eq!(map.heap_size_bytes(), heap_size_of_map(&map));
264
265        map.remove(&1);
266        map.remove(&3);
267        assert_eq!(map.heap_size_bytes(), 0);
268        assert!(map.is_empty());
269    }
270
271    #[test]
272    fn test_extend_bookkeeping() {
273        let mut map: BookkeepingBTreeMap<u64, String> = BookkeepingBTreeMap::new();
274
275        map.extend(vec![
276            (1, "one".to_owned()),
277            (2, "two".to_owned()),
278            (3, "three".to_owned()),
279        ]);
280        assert_eq!(map.len(), 3);
281        assert_eq!(map.heap_size_bytes(), heap_size_of_map(&map));
282
283        map.extend(vec![(2, "TWO".to_owned()), (4, "four".to_owned())]);
284        assert_eq!(map.len(), 4);
285        assert_eq!(map.heap_size_bytes(), heap_size_of_map(&map));
286    }
287
288    #[test]
289    fn test_mutate_entry_bookkeeping() {
290        let mut map: BookkeepingBTreeMap<u64, Vec<String>> = BookkeepingBTreeMap::new();
291
292        map.mutate_entry(1, Vec::new(), |vec| {
293            vec.push("hello".to_owned());
294        });
295        assert_eq!(map.len(), 1);
296        assert_eq!(map.heap_size_bytes(), heap_size_of_map(&map));
297
298        map.mutate_entry(1, Vec::new(), |vec| {
299            vec.push("world".to_owned());
300        });
301        assert_eq!(map.len(), 1);
302        assert_eq!(map.heap_size_bytes(), heap_size_of_map(&map));
303
304        map.mutate_entry(1, Vec::new(), |vec| {
305            vec.pop();
306        });
307        assert_eq!(map.len(), 1);
308        assert_eq!(map.heap_size_bytes(), heap_size_of_map(&map));
309
310        map.mutate_entry(1, Vec::new(), |vec| {
311            vec.clear();
312        });
313        assert_eq!(map.len(), 1);
314        assert_eq!(map.heap_size_bytes(), heap_size_of_map(&map));
315    }
316
317    #[test]
318    fn test_mutate_entry_before_bookkeeping() {
319        let mut map: BookkeepingBTreeMap<u64, Vec<String>> = BookkeepingBTreeMap::new();
320
321        map.insert(10, vec!["ten".to_owned()]);
322        map.insert(20, vec!["twenty".to_owned()]);
323        map.insert(30, vec!["thirty".to_owned()]);
324        assert_eq!(map.heap_size_bytes(), heap_size_of_map(&map));
325
326        let result = map.mutate_latest_at(&20, |key, vec| {
327            assert_eq!(*key, 20);
328            vec.push("added".to_owned());
329            *key
330        });
331        assert_eq!(result, Some(20));
332        assert_eq!(map.heap_size_bytes(), heap_size_of_map(&map));
333
334        let result = map.mutate_latest_at(&100, |key, vec| {
335            assert_eq!(*key, 30);
336            vec.clear();
337            *key
338        });
339        assert_eq!(result, Some(30));
340        assert_eq!(map.heap_size_bytes(), heap_size_of_map(&map));
341
342        let result = map.mutate_latest_at(&5, |_key, vec| {
343            vec.push("should not happen".to_owned());
344        });
345        assert_eq!(result, None);
346        assert_eq!(map.heap_size_bytes(), heap_size_of_map(&map));
347    }
348
349    #[test]
350    fn test_iter() {
351        let mut map: BookkeepingBTreeMap<u64, String> = BookkeepingBTreeMap::new();
352
353        map.insert(3, "three".to_owned());
354        map.insert(1, "one".to_owned());
355        map.insert(2, "two".to_owned());
356
357        let items: Vec<_> = map.iter().collect();
358        assert_eq!(items.len(), 3);
359        assert_eq!(*items[0].0, 1);
360        assert_eq!(*items[1].0, 2);
361        assert_eq!(*items[2].0, 3);
362    }
363
364    #[test]
365    fn test_into_iterator() {
366        let mut map: BookkeepingBTreeMap<u64, String> = BookkeepingBTreeMap::new();
367
368        map.insert(2, "two".to_owned());
369        map.insert(1, "one".to_owned());
370
371        let items: Vec<_> = (&map).into_iter().collect();
372        assert_eq!(items.len(), 2);
373        assert_eq!(*items[0].0, 1);
374        assert_eq!(*items[1].0, 2);
375    }
376
377    #[test]
378    fn test_clone() {
379        let mut map1: BookkeepingBTreeMap<u64, String> = BookkeepingBTreeMap::new();
380
381        map1.insert(1, "one".to_owned());
382        map1.insert(2, "two".to_owned());
383
384        let map2 = map1.clone();
385
386        assert_eq!(map1.len(), map2.len());
387        assert_eq!(map1.heap_size_bytes(), map2.heap_size_bytes());
388        assert_eq!(map1, map2);
389        assert_eq!(map2.heap_size_bytes(), heap_size_of_map(&map2));
390    }
391
392    #[test]
393    fn test_partial_eq() {
394        let mut map1: BookkeepingBTreeMap<u64, String> = BookkeepingBTreeMap::new();
395        let mut map2: BookkeepingBTreeMap<u64, String> = BookkeepingBTreeMap::new();
396
397        map1.insert(1, "one".to_owned());
398        map2.insert(1, "one".to_owned());
399        assert_eq!(map1, map2);
400
401        map1.insert(2, "two".to_owned());
402        assert_ne!(map1, map2);
403
404        map2.insert(2, "TWO".to_owned());
405        assert_ne!(map1, map2);
406
407        map2.insert(2, "two".to_owned());
408        assert_eq!(map1, map2);
409    }
410
411    #[test]
412    fn test_as_map() {
413        let mut map: BookkeepingBTreeMap<u64, String> = BookkeepingBTreeMap::new();
414
415        map.insert(1, "one".to_owned());
416        map.insert(2, "two".to_owned());
417
418        let inner_map = map.as_map();
419        assert_eq!(inner_map.len(), 2);
420        assert_eq!(inner_map.get(&1), Some(&"one".to_owned()));
421        assert_eq!(inner_map.get(&2), Some(&"two".to_owned()));
422    }
423
424    #[test]
425    fn test_bookkeeping_stress() {
426        let mut map: BookkeepingBTreeMap<u64, Vec<String>> = BookkeepingBTreeMap::new();
427
428        for i in 0..100 {
429            map.insert(i, vec![format!("value_{i}")]);
430            assert_eq!(map.heap_size_bytes(), heap_size_of_map(&map));
431        }
432
433        for i in (0..100).step_by(5) {
434            map.mutate_entry(i, Vec::new(), |vec| {
435                vec.push(format!("extra_{i}"));
436            });
437            assert_eq!(map.heap_size_bytes(), heap_size_of_map(&map));
438        }
439
440        for i in (0..100).step_by(3) {
441            map.remove(&i);
442            assert_eq!(map.heap_size_bytes(), heap_size_of_map(&map));
443        }
444
445        assert_eq!(map.heap_size_bytes(), heap_size_of_map(&map));
446    }
447}
re_byte_size/bookkeeping_btreemap.rs

re_byte_size/
bookkeeping_btreemap.rs
