brie_tree/
lib.rs

1//! This crate provides [`BTree`], a fast B+ Tree implementation using integer
2//! keys.
3
4#![cfg_attr(not(any(test, feature = "internal_benches")), no_std)]
5#![warn(missing_docs)]
6
7use core::mem;
8
9use allocator_api2::alloc::{Allocator, Global};
10use int::BTreeInteger;
11use node::{NodePool, NodeRef, UninitNodeRef};
12use stack::Height;
13
14#[macro_use]
15mod node;
16
17mod cursor;
18mod int;
19mod iter;
20mod simd;
21mod stack;
22#[cfg(test)]
23mod tests;
24
25pub use nonmax;
26
27pub use cursor::*;
28pub use iter::*;
29
30use crate::int::int_from_key;
31
32/// Trait which must be implemented for all keys inserted into a [`BTree`].
33///
34/// [`BTree`] requires that keys be integers and reserves the maximum integer
35/// value for internal use. This trait is already implementated for all integers
36/// from the [`nonmax`] crate, but this crate allows for custom key types that
37/// are convertible to/from an integer.
38///
39/// Note that keys in the [`BTree`] are ordered by their integer value and not
40/// the [`Ord`] implementation of the key type.
41pub trait BTreeKey: Copy {
42    /// Non-max integer type that this key
43    ///
44    /// This must be one of the integer types from the [`nonmax`] crate:
45    /// - [`nonmax::NonMaxU8`]
46    /// - [`nonmax::NonMaxU16`]
47    /// - [`nonmax::NonMaxU32`]
48    /// - [`nonmax::NonMaxU64`]
49    /// - [`nonmax::NonMaxU128`]
50    /// - [`nonmax::NonMaxI8`]
51    /// - [`nonmax::NonMaxI16`]
52    /// - [`nonmax::NonMaxI32`]
53    /// - [`nonmax::NonMaxI64`]
54    /// - [`nonmax::NonMaxI128`]
55    #[allow(private_bounds)]
56    type Int: BTreeInteger;
57
58    /// Converts the key to an integer.
59    fn to_int(self) -> Self::Int;
60
61    /// Recovers the key from an integer.
62    fn from_int(int: Self::Int) -> Self;
63}
64
65/// An ordered map based on a [B+ Tree].
66///
67/// This is similar to the standard library's `BTreeMap` but differs in several
68/// ways:
69/// - Lookups and insertions are 2-4x faster than `BTreeMap`.
70/// - `BTree` can optionally be used as a multi-map and hold duplicate keys.
71/// - Keys must be `Copy` and convertible to and from integers via the
72///   [`BTreeKey`] trait.
73/// - The maximum integer value is reserved for internal use and cannot be used
74///   by keys.
75/// - Elements in the tree are ordered by the integer value of the key instead
76///   of the [`Ord`] implementation of the keys.
77/// - [`Cursor`] and [`CursorMut`] can be used to seek back-and-forth in the
78///   tree while inserting or removing elements.
79/// - Iterators only support forward iteration.
80///
81/// The data structure design is based on the [B- Tree] by Sergey Slotin, but
82/// has been significantly extended.
83///
84/// [B+ Tree]: https://en.wikipedia.org/wiki/B%2B_tree
85/// [B- Tree]: https://en.algorithmica.org/hpc/data-structures/b-tree/
86pub struct BTree<K: BTreeKey, V, A: Allocator = Global> {
87    internal: NodePool<K::Int, NodeRef>,
88    leaf: NodePool<K::Int, V>,
89    height: Height<K::Int>,
90    root: NodeRef,
91    alloc: A,
92}
93
94impl<K: BTreeKey, V> BTree<K, V, Global> {
95    /// Creates a new, empty [`BTree`].
96    ///
97    /// This requires an initial memory allocation on creation.
98    #[inline]
99    pub fn new() -> Self {
100        Self::new_in(Global)
101    }
102}
103
104impl<K: BTreeKey, V, A: Allocator> BTree<K, V, A> {
105    /// Creates a new, empty [`BTree`] with the given allocator.
106    ///
107    /// This requires an initial memory allocation on creation.
108    #[inline]
109    pub fn new_in(alloc: A) -> Self {
110        let mut out = Self {
111            internal: NodePool::new(),
112            leaf: NodePool::new(),
113            height: Height::leaf(),
114            root: NodeRef::zero(),
115            alloc,
116        };
117        let root = unsafe { out.leaf.alloc_node(&out.alloc) };
118        out.init_root(root);
119        out
120    }
121
122    /// Initializes the root node to the leaf node at offset zero.
123    #[inline]
124    fn init_root(&mut self, root: UninitNodeRef) {
125        let root = unsafe { root.init_keys(&mut self.leaf) };
126        unsafe {
127            root.set_next_leaf(None, &mut self.leaf);
128        }
129        debug_assert_eq!(root, NodeRef::zero());
130        self.root = NodeRef::zero();
131    }
132
133    /// Clears the map, removing all elements.
134    #[inline]
135    pub fn clear(&mut self) {
136        // Drop values. We don't need to modify the keys since we're about to
137        // free the nodes anyways.
138        if mem::needs_drop::<V>() {
139            let mut iter = self.raw_iter();
140            while let Some((_key, value_ptr)) = unsafe { iter.next(&self.leaf) } {
141                unsafe {
142                    value_ptr.drop_in_place();
143                }
144            }
145        }
146
147        // Free all nodes without freeing the underlying allocations.
148        let root = self.leaf.clear_and_alloc_node();
149        self.internal.clear();
150
151        // Re-initialize the root node.
152        self.height = Height::leaf();
153        self.init_root(root);
154    }
155
156    /// Returns `true` if the map contains no elements.
157    #[inline]
158    pub fn is_empty(&self) -> bool {
159        if self.height != Height::leaf() {
160            return false;
161        }
162        let first_key = unsafe { self.root.key(pos!(0), &self.leaf) };
163        first_key == K::Int::MAX
164    }
165
166    /// Returns a reference to the value corresponding to the key.
167    #[inline]
168    pub fn get(&self, key: K) -> Option<&V> {
169        self.range(key..=key).next().map(|(_k, v)| v)
170    }
171
172    /// Returns a mutable reference to the value corresponding to the key.
173    #[inline]
174    pub fn get_mut(&mut self, key: K) -> Option<&mut V> {
175        self.range_mut(key..=key).next().map(|(_k, v)| v)
176    }
177
178    /// Inserts a key-value pair into the map.
179    ///
180    /// If the map did not have this key present, `None` is returned.
181    ///
182    /// If the map did have this key present, the value is updated, and the old
183    /// value is returned.
184    #[inline]
185    pub fn insert(&mut self, key: K, value: V) -> Option<V> {
186        let mut cursor = unsafe { CursorMut::uninit(self) };
187        cursor.seek(int_from_key(key));
188        if let Some((k, v)) = cursor.entry_mut()
189            && k.to_int() == key.to_int()
190        {
191            return Some(mem::replace(v, value));
192        }
193        cursor.insert_before(key, value);
194        None
195    }
196
197    /// Inserts a key-value pair into the map while allowing for multiple
198    /// identical keys.
199    ///
200    /// This allows the `BTree` to be used as a multi-map where each key can
201    /// have multiple values. In this case [`BTree::get`], [`BTree::get_mut`]
202    /// and [`BTree::remove`] will only operate on one of the associated values
203    /// (arbitrarily chosen).
204    #[inline]
205    pub fn insert_multi(&mut self, key: K, value: V) {
206        let mut cursor = unsafe { CursorMut::uninit(self) };
207        cursor.seek(int_from_key(key));
208        cursor.insert_before(key, value);
209    }
210
211    /// Removes a key from the map, returning the value at the key if the key
212    /// was previously in the map.
213    #[inline]
214    pub fn remove(&mut self, key: K) -> Option<V> {
215        let mut cursor = unsafe { CursorMut::uninit(self) };
216        cursor.seek(int_from_key(key));
217        if cursor.key()?.to_int() == key.to_int() {
218            return Some(cursor.remove().1);
219        }
220        None
221    }
222}
223
224impl<K: BTreeKey, V, A: Allocator> Drop for BTree<K, V, A> {
225    #[inline]
226    fn drop(&mut self) {
227        // Drop values. We don't need to modify the keys since we're about to
228        // free the nodes anyways.
229        if mem::needs_drop::<V>() {
230            let mut iter = self.raw_iter();
231            while let Some((_key, value_ptr)) = unsafe { iter.next(&self.leaf) } {
232                unsafe {
233                    value_ptr.drop_in_place();
234                }
235            }
236        }
237
238        // Release all allocated memory
239        unsafe {
240            self.internal.clear_and_free(&self.alloc);
241            self.leaf.clear_and_free(&self.alloc);
242        }
243    }
244}
245
246impl<K: BTreeKey, V, A: Default + Allocator> Default for BTree<K, V, A> {
247    #[inline]
248    fn default() -> Self {
249        Self::new_in(Default::default())
250    }
251}
252
253impl<K: BTreeKey, V> FromIterator<(K, V)> for BTree<K, V> {
254    #[inline]
255    fn from_iter<T: IntoIterator<Item = (K, V)>>(iter: T) -> Self {
256        let mut btree = BTree::new();
257        btree.extend(iter);
258        btree
259    }
260}
261
262impl<K: BTreeKey, V, A: Allocator> Extend<(K, V)> for BTree<K, V, A> {
263    #[inline]
264    fn extend<T: IntoIterator<Item = (K, V)>>(&mut self, iter: T) {
265        iter.into_iter().for_each(|(k, v)| {
266            self.insert(k, v);
267        });
268    }
269}
270
271impl<K: BTreeKey, V: Clone, A: Allocator + Clone> Clone for BTree<K, V, A> {
272    #[inline]
273    fn clone(&self) -> Self {
274        let mut btree = BTree::new_in(self.alloc.clone());
275        btree.extend(self.iter());
276        btree
277    }
278}
279
280impl<'a, K: BTreeKey, V: Clone, A: Allocator> Extend<(K, &'a V)> for BTree<K, V, A> {
281    #[inline]
282    fn extend<T: IntoIterator<Item = (K, &'a V)>>(&mut self, iter: T) {
283        iter.into_iter().for_each(|(k, v)| {
284            self.insert(k, v.clone());
285        });
286    }
287}
brie_tree/lib.rs

brie_tree/
lib.rs
