be_tree/
lib.rs

1use std::mem;
2use std::ptr;
3
4const MAX_VALUES_PER_LEAF: usize = 4;
5
6/* A pivot is a key and a node of the subtree of values >= that key. */
7struct Pivot<K, V> {
8    min_key: K,
9    child: Box<Node<K, V>>,
10}
11
12struct LeafNode<K, V> {
13    elements: [(K, V); MAX_VALUES_PER_LEAF],
14    // must be <= MAX_VALUES_PER_LEAF
15    len: usize,
16}
17
18impl<K, V> LeafNode<K, V>
19where
20    K: Copy,
21    V: Clone,
22{
23    fn empty() -> Self {
24        unsafe { Self { elements: mem::MaybeUninit::uninit().assume_init(), len: 0 } }
25    }
26
27    fn from(items: &[(K, V)]) -> Self {
28        debug_assert!(items.len() <= MAX_VALUES_PER_LEAF);
29        let mut result = Self::empty();
30        result.elements.clone_from_slice(items);
31        result
32    }
33
34    fn valid_elements_mut(&mut self) -> &mut [(K, V)] {
35        &mut self.elements[0..self.len]
36    }
37
38    fn valid_elements(&self) -> &[(K, V)] {
39        &self.elements[0..self.len]
40    }
41}
42
43struct BranchNode<K, V> {
44    pivots: [Pivot<K, V>; MAX_VALUES_PER_LEAF],
45    // must be <= MAX_VALUES_PER_LEAF and > 1
46    len: usize,
47}
48
49impl<K, V> BranchNode<K, V>
50where
51    K: Copy,
52{
53    fn from(left: Pivot<K, V>, right: Pivot<K, V>) -> Self {
54        unsafe {
55            let mut result = Self { pivots: mem::MaybeUninit::uninit().assume_init(), len: 2 };
56            result.pivots[0] = left;
57            result.pivots[1] = right;
58            result
59        }
60    }
61    fn valid_pivots_mut(&mut self) -> &mut [Pivot<K, V>] {
62        &mut self.pivots[0..self.len]
63    }
64
65    fn valid_pivots(&self) -> &[Pivot<K, V>] {
66        &self.pivots[0..self.len]
67    }
68}
69
70enum Node<K, V> {
71    Branch(BranchNode<K, V>),
72    Leaf(LeafNode<K, V>),
73}
74
75impl<K, V> Node<K, V>
76where
77    K: Copy + Ord,
78    V: Clone,
79{
80    fn min_key(&self) -> K {
81        match *self {
82            Node::Branch(ref branch) => {
83                debug_assert!(branch.len > 1);
84                branch.pivots[0].min_key
85            }
86            Node::Leaf(ref leaf) => {
87                debug_assert_ne!(leaf.len, 0);
88                leaf.elements[0].0
89            }
90        }
91    }
92
93    fn insert(&mut self, key: K, value: V) {
94        let replace_node: Option<Self> = match *self {
95            Node::Branch(ref mut branch) => {
96                // Find a child node whose keys are not before the target key
97                match branch.valid_pivots().iter().position(|ref p| key <= p.min_key) {
98                    Some(i) => {
99                        // If there is one, insert into it and update the pivot key
100                        let pivot = &mut branch.pivots[i];
101                        pivot.min_key = key;
102                        pivot.child.insert(key, value)
103                    }
104                    // o/w, insert a new leaf at the end
105                    None => {
106                        branch.pivots[branch.len] =
107                            Pivot { min_key: key, child: Box::new(Node::Leaf(LeafNode::empty())) };
108                        branch.len += 1
109                        // XXX consider splitting branch
110                    }
111                };
112                None
113            }
114            Node::Leaf(ref mut leaf) => {
115                let index = leaf.valid_elements_mut().binary_search_by_key(&key, |&(k, _)| k);
116                match index {
117                    Err(i) => {
118                        // key is absent, true insert
119                        if leaf.len < MAX_VALUES_PER_LEAF {
120                            // there's space left, just insert
121                            unsafe { slice_insert(leaf.valid_elements_mut(), i, (key, value)) }
122                            leaf.len += 1;
123                            None
124                        } else {
125                            // must split the node: create the new node here
126                            let new_branch = {
127                                let (left, right) =
128                                    leaf.valid_elements_mut().split_at(MAX_VALUES_PER_LEAF / 2);
129                                let left_leaf = Box::new(Node::Leaf(LeafNode::from(left)));
130                                let right_leaf = Box::new(Node::Leaf(LeafNode::from(right)));
131                                Node::Branch(BranchNode::from(
132                                    Pivot { min_key: left_leaf.min_key(), child: left_leaf },
133                                    Pivot { min_key: right_leaf.min_key(), child: right_leaf },
134                                ))
135                            };
136                            Some(new_branch)
137                        }
138                    }
139                    // key is present, replace
140                    Ok(i) => {
141                        leaf.elements[i] = (key, value);
142                        None
143                    }
144                }
145            }
146        };
147        if let Some(new_branch) = replace_node {
148            *self = new_branch
149        }
150    }
151
152    fn delete(&mut self, key: K) {
153        match *self {
154            Node::Branch(ref mut branch) => {
155                // Find a child node whose keys are not before the target key
156                if let Some(ref mut pivot) =
157                    branch.valid_pivots_mut().iter_mut().find(|ref p| key <= p.min_key)
158                {
159                    // If there is one, delete from it and update the pivot key
160                    pivot.child.delete(key);
161                    pivot.min_key = pivot.child.min_key()
162                }
163            }
164            Node::Leaf(ref mut leaf) if leaf.len > 0 => {
165                let index = leaf.valid_elements_mut().binary_search_by_key(&key, |&(k, _)| k);
166                match index {
167                    Err(_) => (),
168                    Ok(i) => unsafe {
169                        slice_remove(leaf.valid_elements_mut(), i);
170                        leaf.len -= 1;
171                    },
172                }
173            }
174            _ => (),
175        }
176    }
177
178    fn get(&self, key: K) -> Option<&V> {
179        match *self {
180            Node::Branch(ref branch) => {
181                // Find a child node whose keys are not before the target key
182                match branch.valid_pivots().iter().find(|ref p| key <= p.min_key) {
183                    Some(ref pivot) => {
184                        // If there is one, query it
185                        pivot.child.get(key)
186                    }
187                    // o/w, the key doesn't exist
188                    None => None,
189                }
190            }
191            Node::Leaf(ref leaf) if leaf.len > 0 => {
192                let index = leaf.valid_elements().binary_search_by_key(&key, |&(k, _)| k);
193                match index {
194                    Err(_) => None,
195                    Ok(i) => Some(&leaf.elements[i].1),
196                }
197            }
198            _ => None,
199        }
200    }
201}
202
203unsafe fn slice_insert<T>(slice: &mut [T], idx: usize, val: T) {
204    ptr::copy(
205        slice.as_mut_ptr().add(idx),
206        slice.as_mut_ptr().offset(idx as isize + 1),
207        slice.len() - idx,
208    );
209    ptr::write(slice.get_unchecked_mut(idx), val);
210}
211
212unsafe fn slice_remove<T>(slice: &mut [T], idx: usize) -> T {
213    let ret = ptr::read(slice.get_unchecked(idx));
214    ptr::copy(
215        slice.as_ptr().offset(idx as isize + 1),
216        slice.as_mut_ptr().add(idx),
217        slice.len() - idx - 1,
218    );
219    ret
220}
221
222/// A map based on a B𝛆-tree
223pub struct BeTree<K, V> {
224    root: Node<K, V>,
225}
226
227impl<K, V> BeTree<K, V>
228where
229    K: Copy + Ord,
230    V: Clone,
231{
232    /// Create an empty B𝛆-tree.
233    pub fn new() -> Self {
234        BeTree { root: Node::Leaf(LeafNode::empty()) }
235    }
236
237    /// Clear the tree, removing all entries.
238    pub fn clear(&mut self) {
239        match self.root {
240            Node::Leaf(ref mut leaf) => leaf.len = 0,
241            _ => self.root = Node::Leaf(LeafNode::empty()),
242        }
243    }
244
245    /// Insert a key-value pair into the tree.
246    ///
247    /// If the key is already present in the tree, the value is replaced. The key is not updated, though; this matters for
248    /// types that can be `==` without being identical.
249    pub fn insert(&mut self, key: K, value: V) {
250        self.root.insert(key, value)
251    }
252
253    /// Remove a key (and its value) from the tree.
254    ///
255    /// If the key is not present, silently does nothing.
256    pub fn delete(&mut self, key: K) {
257        self.root.delete(key)
258    }
259
260    /// Retrieve a reference to the value corresponding to the key.
261    pub fn get(&self, key: K) -> Option<&V> {
262        self.root.get(key)
263    }
264}
265
266impl<K, V> Default for BeTree<K, V>
267where
268    K: Copy + Ord,
269    V: Clone,
270{
271    fn default() -> Self {
272        Self::new()
273    }
274}
275
276#[cfg(test)]
277mod tests {
278    use super::{BeTree, MAX_VALUES_PER_LEAF};
279
280    #[test]
281    fn can_construct() {
282        BeTree::<i32, char>::new();
283    }
284
285    #[test]
286    fn can_insert_single() {
287        let mut b = BeTree::new();
288        b.insert(0, 'x');
289        let result = b.get(0);
290        assert_eq!(Some(&'x'), result);
291    }
292
293    #[test]
294    fn can_insert_two() {
295        let mut b = BeTree::new();
296        b.insert(0, 'x');
297        b.insert(-1, 'y');
298        assert_eq!(Some(&'x'), b.get(0));
299        assert_eq!(Some(&'y'), b.get(-1));
300    }
301
302    #[test]
303    fn can_split() {
304        let mut b = BeTree::new();
305        // insert MAX_VALUES_PER_LEAF + 1 items
306        for i in 0..MAX_VALUES_PER_LEAF {
307            b.insert(i, i);
308        }
309        // are they all there?
310        for i in 0..MAX_VALUES_PER_LEAF {
311            assert_eq!(Some(&i), b.get(i));
312        }
313    }
314
315    #[test]
316    fn can_clear() {
317        let mut b = BeTree::new();
318        b.insert(0, 'x');
319        b.insert(-1, 'y');
320        b.clear();
321        assert_eq!(None, b.get(0));
322    }
323
324    #[test]
325    fn insert_replaces_existing() {
326        let mut b = BeTree::new();
327        b.insert(0, 'x');
328        b.insert(0, 'y');
329        assert_eq!(Some(&'y'), b.get(0));
330    }
331
332    #[test]
333    fn can_delete_existing() {
334        let mut b = BeTree::new();
335        b.insert(0, 'x');
336        b.delete(0);
337        assert_eq!(b.get(0), None)
338    }
339
340    #[test]
341    fn can_delete_only_existing() {
342        let mut b = BeTree::new();
343        b.insert(0, 'x');
344        b.insert(2, 'y');
345        b.delete(0);
346        assert_eq!(b.get(0), None);
347        assert_eq!(Some(&'y'), b.get(2));
348    }
349
350    #[test]
351    fn can_delete_nothing() {
352        let mut b = BeTree::<i32, char>::new();
353        b.delete(0);
354    }
355}