allocated_btree/
compressed.rs

1#![allow(dead_code)]
2
3use allocated::DropGuard;
4use allocated::FromIteratorIn;
5use core::borrow::Borrow;
6use core::mem::ManuallyDrop;
7use core::ops::Add;
8use core::ops::Mul;
9use core::ptr::NonNull;
10
11extern crate alloc;
12use alloc::vec;
13use alloc::vec::Vec;
14
15use allocator_api2::alloc::Allocator;
16use generic_array::ArrayLength;
17use typenum::Prod;
18use typenum::Sum;
19use typenum::U1;
20use typenum::U2;
21use typenum::U6;
22
23use allocated::AllocResult;
24use allocated::DropGuardResult;
25use allocated::DropIn;
26use allocated::IntoIteratorIn;
27use allocated::RecursiveDropIn;
28
29mod entry;
30mod node;
31mod wrapper;
32
33pub use entry::{Entry, OccupiedEntry, VacantEntry};
34use node::{
35    ChildPtr, IntoIter as NodeIntoIter, Iter as NodeIter, LeafNode, MutNodeRef, Node, NodeEntry,
36    NodeRef, NodeRefT, OccupiedNodeEntry,
37};
38pub use wrapper::CompressedBTreeMap;
39
40/// A compressed B-Tree map implementation using the allocated pattern.
41///
42/// This is the low-level "allocated" type that requires manual allocator passing.
43/// For most use cases, prefer the [`CompressedBTreeMap`] wrapper which owns its allocator
44/// and provides a safe, ergonomic API.
45///
46/// This implementation uses ~30% less memory than the naive implementation by using
47/// specialized node types:
48/// - **Leaf nodes**: Store only keys and values (no child pointers)
49/// - **Interior nodes**: Store keys, values, and child pointers
50///
51/// # Type Parameters
52///
53/// - `K`: Key type, must be `PartialOrd + Debug`
54/// - `V`: Value type
55/// - `B`: Branching factor (defaults to `U6` for 6-way tree). Controls the number
56///   of keys per node (2*B keys maximum).
57///
58/// # Examples
59///
60/// ```
61/// use allocated_btree::AllocatedCompressedBTreeMap;
62/// use allocated::CountingAllocator;
63///
64/// let alloc = CountingAllocator::default();
65/// let mut map = AllocatedCompressedBTreeMap::<u32, String>::new_in(&alloc)?;
66///
67/// unsafe {
68///     map.insert_in(&alloc, 1, "one".to_string())?;
69///     map.insert_in(&alloc, 2, "two".to_string())?;
70/// }
71///
72/// assert_eq!(map.len(), 2);
73/// # Ok::<(), allocated::AllocErrorWithLayout>(())
74/// ```
75pub struct AllocatedBTreeMap<K: core::cmp::PartialOrd + core::fmt::Debug, V, B: ArrayLength = U6>
76where
77    U2: Mul<B>,
78    Prod<U2, B>: ArrayLength,
79    U1: Add<Prod<U2, B>>,
80    Sum<U1, Prod<U2, B>>: ArrayLength,
81{
82    root: Option<ManuallyDrop<Node<K, V, B>>>,
83    n: usize,
84}
85
86impl<K: core::cmp::PartialOrd + core::fmt::Debug, V, B: ArrayLength> AllocatedBTreeMap<K, V, B>
87where
88    U2: Mul<B>,
89    Prod<U2, B>: ArrayLength,
90    U1: Add<Prod<U2, B>>,
91    Sum<U1, Prod<U2, B>>: ArrayLength,
92{
93    fn root_node_ref(&self) -> NodeRef<'_, K, V, B> {
94        self.root.as_ref().unwrap().as_ref()
95    }
96
97    fn root_mut_node_ref(&mut self) -> MutNodeRef<'_, K, V, B> {
98        self.root.as_mut().unwrap().as_mut()
99    }
100}
101
102impl<K: core::cmp::PartialOrd + core::fmt::Debug, V, B: ArrayLength> AllocatedBTreeMap<K, V, B>
103where
104    U2: Mul<B>,
105    Prod<U2, B>: ArrayLength,
106    U1: Add<Prod<U2, B>>,
107    Sum<U1, Prod<U2, B>>: ArrayLength,
108{
109    /// # Errors
110    ///
111    /// Will return `Err` if the allocation fails.
112    pub fn new_in<A: Allocator>(alloc: &A) -> DropGuardResult<Self, &A> {
113        Ok(LeafNode::new_in(alloc).map(|root| AllocatedBTreeMap {
114            root: Some(ManuallyDrop::new(Node::Leaf(root))),
115            n: 0,
116        }))
117    }
118
119    /// Returns `true` if the map contains no elements.
120    pub fn is_empty(&self) -> bool {
121        self.n == 0
122    }
123
124    /// Returns the number of elements in the map.
125    pub fn len(&self) -> usize {
126        self.n
127    }
128
129    /// Returns `true` if the map contains a value for the specified key.
130    pub fn contains_key<Q>(&self, key: &Q) -> bool
131    where
132        K: Borrow<Q> + core::cmp::PartialOrd + core::fmt::Debug,
133        Q: core::cmp::PartialOrd + core::fmt::Debug,
134    {
135        let root = self.root_node_ref();
136        match root.ref_entry(key, vec![]) {
137            NodeEntry::Vacant(_) => false,
138            NodeEntry::Occupied(_) => true,
139        }
140    }
141
142    /// Inserts a key-value pair into the map.
143    ///
144    /// If the map did not have this key present, `None` is returned.
145    /// If the map did have this key present, the value is updated and the old value is returned.
146    ///
147    /// # Safety
148    ///
149    /// `alloc` MUST be the allocator used to allocate this object.
150    ///
151    /// # Errors
152    ///
153    /// Will return `Err` if the allocation fails.
154    pub unsafe fn insert_in<A: Allocator>(
155        &mut self,
156        alloc: &A,
157        key: K,
158        value: V,
159    ) -> AllocResult<Option<V>> {
160        // SAFETY: Caller guarantees `alloc` is the allocator used for this tree
161        let entry = unsafe { self.entry_in(alloc, key) };
162        entry.insert(value)
163    }
164
165    /// Clears the map, removing all elements.
166    ///
167    /// # Safety
168    ///
169    /// `alloc` MUST be the allocator used to allocate this object.
170    ///
171    /// # Errors
172    ///
173    /// Will return `Err` if the allocation fails.
174    pub unsafe fn clear_in<A: Allocator>(&mut self, alloc: &A) -> AllocResult<()> {
175        if let Some(mut r) = self.root.take() {
176            // SAFETY: Caller guarantees `alloc` is the allocator used for this tree
177            unsafe { r.drop_in(alloc) };
178        }
179        self.n = 0;
180        self.root = Some(LeafNode::new_in(alloc).map(|n| Node::Leaf(n)).into_inner());
181        Ok(())
182    }
183
184    /// Returns a reference to the value corresponding to the key.
185    pub fn get<'s, Q>(&'s self, key: &Q) -> Option<&'s V>
186    where
187        K: Borrow<Q> + core::cmp::PartialOrd,
188        Q: core::cmp::PartialOrd + core::fmt::Debug,
189    {
190        let root = self.root_node_ref();
191        match root.ref_entry(key, vec![]) {
192            NodeEntry::Vacant(_) => None,
193            NodeEntry::Occupied(o) => Some(o.into_value()),
194        }
195    }
196
197    /// Returns the key-value pair corresponding to the supplied key.
198    pub fn get_key_value<'s, Q>(&'s self, key: &'s Q) -> Option<(&'s K, &'s V)>
199    where
200        K: Borrow<Q> + core::cmp::PartialOrd,
201        Q: core::cmp::PartialOrd + core::fmt::Debug,
202    {
203        let root = self.root_node_ref();
204        match root.ref_entry(key, vec![]) {
205            NodeEntry::Vacant(_) => None,
206            NodeEntry::Occupied(o) => Some(o.into_key_value()),
207        }
208    }
209
210    /// Returns a mutable reference to the value corresponding to the key.
211    pub fn get_mut<'s, Q>(&'s mut self, key: &'s Q) -> Option<&'s mut V>
212    where
213        K: Borrow<Q> + core::cmp::PartialOrd,
214        Q: core::cmp::PartialOrd + core::fmt::Debug,
215    {
216        let root = self.root_mut_node_ref();
217        match root.ref_entry(key, vec![]) {
218            NodeEntry::Vacant(_) => None,
219            NodeEntry::Occupied(o) => Some(o.into_mut()),
220        }
221    }
222
223    /// Returns an iterator over the key-value pairs of the map, in sorted order by key.
224    pub fn iter(&self) -> Iter<'_, K, V, B> {
225        let mut stack = Vec::new();
226        if self.n > 0 {
227            stack.push((ChildPtr::from_node_ref(self.root_node_ref()), 0));
228        }
229        Iter {
230            inner: NodeIter::<K, V, B, NodeRef<K, V, B>>::new(stack),
231        }
232    }
233
234    /// Returns an iterator over the keys of the map, in sorted order.
235    pub fn keys(&self) -> Keys<'_, K, V, B> {
236        let mut stack = Vec::new();
237        if self.n > 0 {
238            stack.push((ChildPtr::from_node_ref(self.root_node_ref()), 0));
239        }
240        Keys {
241            inner: NodeIter::<K, V, B, NodeRef<K, V, B>>::new(stack),
242        }
243    }
244
245    /// # Safety
246    ///
247    /// `alloc` MUST be the allocator used to allocate this object.
248    unsafe fn into_keys_in<A: Allocator>(self, alloc: &A) -> IntoKeys<'_, K, V, B, A> {
249        IntoKeys {
250            inner: NodeIntoIter::new(alloc, ManuallyDrop::into_inner(self.root.unwrap())),
251        }
252    }
253
254    /// Returns an iterator over the values of the map, in order by key.
255    pub fn values(&self) -> Values<'_, K, V, B> {
256        let mut stack = Vec::new();
257        if self.n > 0 {
258            stack.push((ChildPtr::from_node_ref(self.root_node_ref()), 0));
259        }
260        Values {
261            inner: NodeIter::<K, V, B, NodeRef<K, V, B>>::new(stack),
262        }
263    }
264
265    /// Returns a mutable iterator over the values of the map, in order by key.
266    pub fn values_mut(&mut self) -> ValuesMut<'_, K, V, B> {
267        let mut stack = Vec::new();
268        if self.n > 0 {
269            stack.push((
270                ChildPtr::from_mut_node_ref(&mut self.root_mut_node_ref()),
271                0,
272            ));
273        }
274        ValuesMut {
275            inner: NodeIter::<K, V, B, MutNodeRef<K, V, B>>::new(stack),
276        }
277    }
278
279    /// # Safety
280    ///
281    /// `alloc` MUST be the allocator used to allocate this object.
282    unsafe fn into_values_in<A: Allocator>(self, alloc: &A) -> IntoValues<'_, K, V, B, A> {
283        IntoValues {
284            inner: NodeIntoIter::new(alloc, ManuallyDrop::into_inner(self.root.unwrap())),
285        }
286    }
287
288    /// # Safety
289    ///
290    /// `alloc` MUST be the allocator used to allocate this object.
291    pub unsafe fn entry_in<'a, 's, A: Allocator>(
292        &'s mut self,
293        alloc: &'a A,
294        key: K,
295    ) -> Entry<'a, 's, A, K, V, B> {
296        let map: NonNull<AllocatedBTreeMap<K, V, B>> = NonNull::from_mut(self);
297        let node_ref: MutNodeRef<'s, K, V, B> = self.root_mut_node_ref();
298        // SAFETY: We have exclusive mutable access to self and are creating a NodeEntry
299        let inner_entry: NodeEntry<'s, K, K, V, B, MutNodeRef<'s, K, V, B>> =
300            unsafe { node_ref.entry_in(key, vec![]) };
301        Entry::new(alloc, inner_entry, map)
302    }
303
304    /// # Safety
305    ///
306    /// `alloc` MUST be the allocator used to allocate this object.
307    pub unsafe fn first_entry_in<'a, 's, A: Allocator>(
308        &'s mut self,
309        alloc: &'a A,
310    ) -> Option<OccupiedEntry<'a, 's, A, K, V, B>> {
311        if self.n == 0 {
312            return None;
313        }
314
315        let map = NonNull::new(self)?;
316        let root = self.root_mut_node_ref();
317        let inner_entry: OccupiedNodeEntry<'s, K, V, B, MutNodeRef<'s, K, V, B>> =
318            root.first_entry_in(vec![]);
319        // SAFETY: Requirements match function requirements
320        unsafe { Some(OccupiedEntry::new(alloc, inner_entry, map)) }
321    }
322
323    /// Returns a reference to the first key-value pair in the map.
324    /// The key in this pair is the minimum key in the map.
325    pub fn first_key_value<'s>(&'s self) -> Option<(&'s K, &'s V)> {
326        if self.n == 0 {
327            return None;
328        }
329
330        let root = self.root_node_ref();
331        let inner_entry: OccupiedNodeEntry<'s, K, V, B, NodeRef<'s, K, V, B>> =
332            root.first_entry_in(vec![]);
333        Some(inner_entry.into_key_value())
334    }
335}
336
337// impl<K: core::cmp::PartialOrd + Debug, V: Debug, B: ArrayLength> AllocatedBTreeMap<K, V, B>
338// where
339//     U2: Mul<B>,
340//     Prod<U2, B>: ArrayLength,
341//     U1: Add<Prod<U2, B>>,
342//     Sum<U1, Prod<U2, B>>: ArrayLength,
343// {
344//     fn to_dot(&self) -> Result<String, Box<dyn Error>> {
345//         let mut data = Vec::default();
346
347//         data.write_all(b"digraph G {\n")?;
348//         data.write_all(b"rankdir=\"LR\";\n")?;
349//         self.root.as_ref().unwrap().to_dot(&mut data)?;
350//         data.write_all(b"}\n")?;
351
352//         Ok(String::from_utf8(data)?)
353//     }
354// }
355
356impl<'a, K: core::cmp::PartialOrd + core::fmt::Debug, V, B: ArrayLength, A: Allocator>
357    FromIteratorIn<'a, (K, V), A> for AllocatedBTreeMap<K, V, B>
358where
359    U2: Mul<B>,
360    Prod<U2, B>: ArrayLength,
361    U1: Add<Prod<U2, B>>,
362    Sum<U1, Prod<U2, B>>: ArrayLength,
363{
364    fn from_iter_in<T>(alloc: &'a A, iter: T) -> DropGuardResult<Self, &'a A>
365    where
366        T: IntoIterator<Item = (K, V)>,
367    {
368        let mut btree: DropGuard<Self, &'a A> = Self::new_in(alloc)?;
369
370        for (k, v) in iter {
371            // SAFETY: alloc is the same allocator used to create btree.
372            unsafe {
373                btree.insert_in(alloc, k, v)?;
374            }
375        }
376
377        Ok(btree)
378    }
379}
380
381impl<K: core::cmp::PartialOrd + core::fmt::Debug, V, B: ArrayLength> DropIn
382    for AllocatedBTreeMap<K, V, B>
383where
384    U2: Mul<B>,
385    Prod<U2, B>: ArrayLength,
386    U1: Add<Prod<U2, B>>,
387    Sum<U1, Prod<U2, B>>: ArrayLength,
388{
389    /// # Safety
390    ///
391    /// `alloc` must be the allocator used to allocate this object.
392    unsafe fn drop_in<A: Allocator>(&mut self, alloc: &A) {
393        // SAFETY: requirements match function requirements
394        unsafe {
395            if let Some(r) = self.root.as_mut() {
396                r.drop_in(alloc);
397            }
398        }
399    }
400}
401
402impl<K, V, B> RecursiveDropIn for AllocatedBTreeMap<K, V, B>
403where
404    K: core::cmp::PartialOrd + core::fmt::Debug + DropIn,
405    V: DropIn,
406    B: ArrayLength,
407    U2: Mul<B>,
408    Prod<U2, B>: ArrayLength,
409    U1: Add<Prod<U2, B>>,
410    Sum<U1, Prod<U2, B>>: ArrayLength,
411{
412    /// # Safety
413    ///
414    /// `alloc` must be the allocator used to allocate this object.
415    unsafe fn recursive_drop_in<A: Allocator>(&mut self, alloc: &A) {
416        if let Some(root) = self.root.as_mut() {
417            // SAFETY: We're taking ownership of the root to iterate over it.
418            let root_owned = unsafe { ManuallyDrop::take(root) };
419            // Recursively drop all keys and values in the tree
420            for (mut k, mut v) in NodeIntoIter::new(alloc, root_owned) {
421                // SAFETY: alloc is the same allocator used for the B-tree.
422                unsafe { k.drop_in(alloc) };
423                // SAFETY: alloc is the same allocator used for the B-tree.
424                unsafe { v.drop_in(alloc) };
425            }
426        }
427        // SAFETY: alloc is the same allocator used for the B-tree. The tree structure itself is dropped.
428        unsafe { self.drop_in(alloc) };
429    }
430}
431
432impl<'a, K: core::cmp::PartialOrd + core::fmt::Debug, V, B: ArrayLength, A: Allocator + 'a>
433    IntoIteratorIn<'a, A> for AllocatedBTreeMap<K, V, B>
434where
435    U2: Mul<B>,
436    Prod<U2, B>: ArrayLength,
437    U1: Add<Prod<U2, B>>,
438    Sum<U1, Prod<U2, B>>: ArrayLength,
439{
440    type Item = (K, V);
441    type IntoIter = IntoIter<'a, K, V, B, A>;
442
443    unsafe fn into_iter_in(self, alloc: &'a A) -> Self::IntoIter {
444        IntoIter {
445            inner: NodeIntoIter::new(alloc, ManuallyDrop::into_inner(self.root.unwrap())),
446        }
447    }
448}
449
450impl<'s, K: core::cmp::PartialOrd + core::fmt::Debug, V, B: ArrayLength> IntoIterator
451    for &'s AllocatedBTreeMap<K, V, B>
452where
453    U2: Mul<B>,
454    Prod<U2, B>: ArrayLength,
455    U1: Add<Prod<U2, B>>,
456    Sum<U1, Prod<U2, B>>: ArrayLength,
457{
458    type IntoIter = Iter<'s, K, V, B>;
459    type Item = (&'s K, &'s V);
460
461    fn into_iter(self) -> Self::IntoIter {
462        self.iter()
463    }
464}
465
466mod iters;
467#[cfg(test)]
468mod tests;
469
470pub use iters::{IntoIter, IntoKeys, IntoValues, Iter, Keys, Values, ValuesMut};
allocated_btree/compressed.rs

allocated_btree/
compressed.rs
