allocated_btree/compressed/
wrapper.rs

1//! Ergonomic wrapper for the compressed B-Tree implementation.
2//!
3//! This module provides [`CompressedBTreeMap<K, V, B, A>`], a wrapper around
4//! the allocated B-Tree that owns an allocator, making it safe and
5//! ergonomic to use.
6
7use core::borrow::Borrow;
8use core::fmt::Debug;
9use core::mem::ManuallyDrop;
10use core::ops::Add;
11use core::ops::Mul;
12
13use allocator_api2::alloc::{Allocator, Global};
14use generic_array::ArrayLength;
15use typenum::{Prod, Sum, U1, U2, U6};
16
17use allocated::{AllocResult, AllocResultExt, DropIn};
18
19use super::{AllocatedBTreeMap, Entry, Iter, Keys, OccupiedEntry, Values, ValuesMut};
20
21/// An ergonomic compressed B-Tree map that owns its allocator.
22///
23/// This is a memory-optimized B-Tree that uses specialized leaf nodes
24/// without child pointers, reducing memory usage by approximately 30%
25/// compared to the naive implementation.
26///
27/// This is the recommended type for most use cases. It wraps the allocated
28/// B-Tree and provides safe methods without requiring `unsafe` blocks or
29/// passing allocators manually.
30///
31/// # Example
32///
33/// ```
34/// use allocated_btree::CompressedBTreeMap;
35///
36/// let mut map = CompressedBTreeMap::new();
37/// map.insert(1, "one").unwrap();
38/// map.insert(2, "two").unwrap();
39///
40/// assert_eq!(map.get(&1), Some(&"one"));
41/// assert_eq!(map.len(), 2);
42/// ```
43pub struct CompressedBTreeMap<
44    K: core::cmp::PartialOrd + core::fmt::Debug,
45    V,
46    B: ArrayLength = U6,
47    A: Allocator = Global,
48> where
49    U2: Mul<B>,
50    Prod<U2, B>: ArrayLength,
51    U1: Add<Prod<U2, B>>,
52    Sum<U1, Prod<U2, B>>: ArrayLength,
53{
54    alloc: A,
55    raw: ManuallyDrop<AllocatedBTreeMap<K, V, B>>,
56}
57
58impl<K: core::cmp::PartialOrd + Debug, V> CompressedBTreeMap<K, V> {
59    /// Create a new empty compressed B-Tree using the global allocator.
60    ///
61    /// # Panics
62    ///
63    /// Panics if allocation fails.
64    #[inline]
65    pub fn new() -> Self {
66        let raw = AllocatedBTreeMap::new_in(&Global)
67            .handle_alloc_error()
68            .into_inner();
69
70        Self { alloc: Global, raw }
71    }
72}
73
74impl<K: core::cmp::PartialOrd + Debug, V> Default for CompressedBTreeMap<K, V> {
75    fn default() -> Self {
76        Self::new()
77    }
78}
79
80impl<K: core::cmp::PartialOrd + Debug, V, B: ArrayLength, A: Allocator> Drop
81    for CompressedBTreeMap<K, V, B, A>
82where
83    U2: Mul<B>,
84    Prod<U2, B>: ArrayLength,
85    U1: Add<Prod<U2, B>>,
86    Sum<U1, Prod<U2, B>>: ArrayLength,
87{
88    fn drop(&mut self) {
89        // SAFETY: `self.raw` was allocated by `self.alloc`
90        unsafe {
91            self.raw.drop_in(&self.alloc);
92        }
93    }
94}
95
96impl<K: core::cmp::PartialOrd + Debug, V, B: ArrayLength, A: Allocator>
97    CompressedBTreeMap<K, V, B, A>
98where
99    U2: Mul<B>,
100    Prod<U2, B>: ArrayLength,
101    U1: Add<Prod<U2, B>>,
102    Sum<U1, Prod<U2, B>>: ArrayLength,
103{
104    /// Create a new empty compressed B-Tree using the provided allocator.
105    ///
106    /// # Errors
107    ///
108    /// Returns an error if allocation fails.
109    pub fn new_in(alloc: A) -> AllocResult<Self> {
110        let raw = AllocatedBTreeMap::new_in(&alloc)?.into_inner();
111        Ok(Self { alloc, raw })
112    }
113
114    /// Returns the number of elements in the map.
115    #[inline]
116    pub fn len(&self) -> usize {
117        self.raw.len()
118    }
119
120    /// Returns `true` if the map contains no elements.
121    #[inline]
122    pub fn is_empty(&self) -> bool {
123        self.raw.is_empty()
124    }
125
126    /// Returns `true` if the map contains a value for the specified key.
127    pub fn contains_key<Q>(&self, key: &Q) -> bool
128    where
129        K: Borrow<Q>,
130        Q: core::cmp::PartialOrd + Debug,
131    {
132        self.raw.contains_key(key)
133    }
134
135    /// Returns a reference to the value corresponding to the key.
136    pub fn get<Q>(&self, key: &Q) -> Option<&V>
137    where
138        K: Borrow<Q>,
139        Q: core::cmp::PartialOrd + Debug,
140    {
141        self.raw.get(key)
142    }
143
144    /// Returns the key-value pair corresponding to the supplied key.
145    pub fn get_key_value<'s, Q>(&'s self, key: &'s Q) -> Option<(&'s K, &'s V)>
146    where
147        K: Borrow<Q>,
148        Q: core::cmp::PartialOrd + Debug,
149    {
150        self.raw.get_key_value(key)
151    }
152
153    /// Returns a mutable reference to the value corresponding to the key.
154    pub fn get_mut<'s, Q>(&'s mut self, key: &'s Q) -> Option<&'s mut V>
155    where
156        K: Borrow<Q>,
157        Q: core::cmp::PartialOrd + Debug,
158    {
159        self.raw.get_mut(key)
160    }
161
162    /// Inserts a key-value pair into the map.
163    ///
164    /// If the map did not have this key present, `None` is returned.
165    /// If the map did have this key present, the value is updated,
166    /// and the old value is returned.
167    ///
168    /// # Errors
169    ///
170    /// Returns an error if allocation fails during tree rebalancing.
171    pub fn insert(&mut self, key: K, value: V) -> AllocResult<Option<V>> {
172        // SAFETY: `self.alloc` was used to allocate `self.raw`
173        unsafe { self.raw.insert_in(&self.alloc, key, value) }
174    }
175
176    /// Clears the map, removing all key-value pairs.
177    ///
178    /// # Errors
179    ///
180    /// Returns an error if allocation fails.
181    pub fn clear(&mut self) -> AllocResult<()> {
182        // SAFETY: `self.alloc` was used to allocate `self.raw`
183        unsafe { self.raw.clear_in(&self.alloc) }
184    }
185
186    /// Gets the given key's corresponding entry in the map for in-place manipulation.
187    pub fn entry(&mut self, key: K) -> Entry<'_, '_, A, K, V, B> {
188        // SAFETY: `self.alloc` was used to allocate `self.raw`
189        unsafe { self.raw.entry_in(&self.alloc, key) }
190    }
191
192    /// Gets the first entry in the map for in-place manipulation.
193    pub fn first_entry(&mut self) -> Option<OccupiedEntry<'_, '_, A, K, V, B>> {
194        // SAFETY: `self.alloc` was used to allocate `self.raw`
195        unsafe { self.raw.first_entry_in(&self.alloc) }
196    }
197
198    /// Returns the first key-value pair in the map.
199    pub fn first_key_value(&self) -> Option<(&K, &V)> {
200        self.raw.first_key_value()
201    }
202
203    /// Gets an iterator over the entries of the map, sorted by key.
204    pub fn iter(&self) -> Iter<'_, K, V, B> {
205        self.raw.iter()
206    }
207
208    /// Gets an iterator over the keys of the map, in sorted order.
209    pub fn keys(&self) -> Keys<'_, K, V, B> {
210        self.raw.keys()
211    }
212
213    /// Gets an iterator over the values of the map, in order by key.
214    pub fn values(&self) -> Values<'_, K, V, B> {
215        self.raw.values()
216    }
217
218    /// Gets a mutable iterator over the values of the map, in order by key.
219    pub fn values_mut(&mut self) -> ValuesMut<'_, K, V, B> {
220        self.raw.values_mut()
221    }
222}
223
224impl<'s, K: core::cmp::PartialOrd + Debug, V, B: ArrayLength, A: Allocator> IntoIterator
225    for &'s CompressedBTreeMap<K, V, B, A>
226where
227    U2: Mul<B>,
228    Prod<U2, B>: ArrayLength,
229    U1: Add<Prod<U2, B>>,
230    Sum<U1, Prod<U2, B>>: ArrayLength,
231{
232    type IntoIter = Iter<'s, K, V, B>;
233    type Item = (&'s K, &'s V);
234
235    fn into_iter(self) -> Self::IntoIter {
236        self.iter()
237    }
238}
allocated_btree/compressed/wrapper.rs

allocated_btree/compressed/
wrapper.rs
