Skip to main content

packed/
packed_map_builder.rs

1// Copyright 2026 The Fuchsia Authors. All rights reserved.
2// Use of this source code is governed by a BSD-style license that can be
3// found in the LICENSE file.
4
5use crate::packed_map::PackedMap;
6use crate::{PackedItem, PackedVec};
7use std::borrow::Borrow;
8use std::collections::BTreeMap;
9
10/// A builder for `PackedMap` that allows incrementally appending potentially
11/// unsorted key-value pairs.
12///
13/// # Algorithm
14///
15/// `PackedMapBuilder` uses an LSM-tree (Log-Structured Merge-tree) inspired approach
16/// to handle out-of-order insertions efficiently while maintaining a compact,
17/// immutable representation for the final map.
18///
19/// It maintains two structures:
20/// 1. `packed`: A `PackedMap` containing elements that are already sorted and packed.
21/// 2. `unpacked`: A `BTreeMap` serving as a buffer for new, potentially out-of-order insertions.
22///
23/// ## Operations
24///
25/// - **Insertion**:
26///   - We first attempt to insert the key-value pair directly into the `packed` map.
27///     `PackedMap::insert` succeeds if the key already exists (overwriting it) or
28///     if the key is greater than or equal to the last key in the map (maintaining sorted order).
29///     This fast path takes `O(log N)` or `O(1)` time.
30///   - If the key is out of order, the `packed` map rejects it. We then insert it into
31///     the `unpacked` buffer (`BTreeMap`).
32/// - **Lookup**: To find a key, we check both the `packed` map and the `unpacked` buffer.
33/// - **Build**: When `build()` is called, any remaining elements in the `unpacked` buffer
34///   are merged with the `packed` map to produce the final `PackedMap`.
35///
36/// This design optimizes for the common case where data is mostly sorted, while still
37/// supporting full unpacked insertion without severe performance degradation.
38pub 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    /// Creates a new empty `PackedMapBuilder`.
53    pub fn new() -> Self {
54        Self { packed: PackedMap::new(), unpacked: BTreeMap::new() }
55    }
56
57    /// Creates a new `PackedMapBuilder` with specified capacity.
58    ///
59    /// The `element_capacity` argument specifies the number of slices that can be
60    /// stored without reallocating the offsets vector. The `buffer_capacity`
61    /// argument specifies the cumulative length of slices that can be stored
62    /// without reallocating the data vector.
63    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    /// Returns `true` if the builder contains the given key.
71    pub fn contains_key(&self, key: &K) -> bool {
72        self.packed.contains_key(key) || self.unpacked.contains_key(key)
73    }
74
75    /// Returns a reference to the value corresponding to the key.
76    pub fn get(&self, key: &K) -> Option<&V> {
77        self.packed.get(key).or_else(|| self.unpacked.get(key))
78    }
79
80    /// Inserts a key-value pair into the builder.
81    ///
82    /// # Complexity
83    ///
84    /// - `O(1)` if the key is greater than or equal to the last key in the map.
85    /// - `O(log N)` otherwise, where `N` is the number of elements in the map.
86    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    /// Builds the `PackedMap`.
94    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); // Out of order!
189        }
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); // Equal, should go to packed
325
326        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); // Out of order, goes to unpacked
341        builder.insert(&key3, 3); // Out of order, goes to unpacked
342
343        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); // Keeps old key from unpacked
348        assert_eq!(*v1, 3); // Keeps new value
349    }
350}