sweep_bptree/tree/
mod.rs

1mod inner_node;
2mod slice_utils;
3use std::{borrow::Borrow, mem::ManuallyDrop};
4
5mod consts;
6
7pub use inner_node::*;
8mod leaf_node;
9pub use leaf_node::*;
10mod node_id;
11pub use node_id::*;
12mod cursor;
13pub use cursor::*;
14mod iterator;
15pub use iterator::*;
16mod node_stores;
17pub use node_stores::*;
18
19pub mod visit;
20
21mod bulk_load;
22pub use crate::argument::*;
23
24use self::entry_ref::{EntryRef, VisitStack};
25mod entry_ref;
26
27mod tree_remove;
28
29/// B plus tree implementation, with following considerations:
30///
31/// 1. Performance critical, for sweep like algo, the sweepline's searching and updating is on hot path
32/// 2. Cursor support, after one search, it should be fairly easy to move next or prev without another search
33/// 3. Memory efficient, reduce mem alloc
34///
35/// # Example
36/// ```rust
37/// use sweep_bptree::{BPlusTree, NodeStoreVec};
38///
39/// // create a node_store with `u64` as key, `(f64, f64)` as value
40/// let mut node_store = NodeStoreVec::<u64, (f64, f64)>::new();
41/// let mut tree = BPlusTree::new(node_store);
42///
43/// // insert new value
44/// assert!(tree.insert(3, (0., 0.)).is_none());
45///
46/// // update by insert again
47/// assert_eq!(tree.insert(3, (1., 1.)).unwrap(), (0., 0.));
48///
49/// // remove the value
50/// assert_eq!(tree.remove(&3).unwrap(), (1.0, 1.0));
51///
52/// assert!(tree.is_empty());
53/// ```
54///
55/// # Example
56/// Create multiple owned cursors
57///
58/// ``` rust
59/// use sweep_bptree::{BPlusTree, NodeStoreVec};
60/// let mut node_store = NodeStoreVec::<u64, (f64, f64)>::new();
61/// let mut tree = BPlusTree::new(node_store);
62///
63/// for i in 0..100 {
64///     tree.insert(i, (i as f64, i as f64));
65/// }
66///
67/// let cursor_0 = tree.cursor_first().unwrap();
68/// let cursor_1 = tree.cursor_first().unwrap();
69///
70/// // remove the key 0
71/// tree.remove(&0);
72///
73/// // cursor's value should be empty now
74/// assert!(cursor_0.value(&tree).is_none());
75///
76/// // move to next
77/// let cursor_next = cursor_0.next(&tree).unwrap();
78/// assert_eq!(*cursor_next.key(), 1);
79/// assert_eq!(cursor_next.value(&tree).unwrap().0, 1.);
80///
81/// // insert a new value to key 0
82/// tree.insert(0, (100., 100.));
83/// // now cursor_1 should retrieve the new value
84/// assert_eq!(cursor_1.value(&tree).unwrap().0, 100.);
85/// ```
86#[derive(Clone)]
87pub struct BPlusTree<S: NodeStore> {
88    root: NodeId,
89    root_argument: S::Argument,
90    len: usize,
91    node_store: ManuallyDrop<S>,
92    st: Statistic,
93}
94
95impl<S> BPlusTree<S>
96where
97    S: NodeStore,
98{
99    /// Create a new `BPlusTree` with the given `NodeStore`.
100    pub fn new(mut node_store: S) -> Self {
101        let (root_id, _) = node_store.new_empty_leaf();
102        Self {
103            root: NodeId::Leaf(root_id),
104            root_argument: S::Argument::default(),
105            node_store: ManuallyDrop::new(node_store),
106            len: 0,
107
108            st: Statistic::default(),
109        }
110    }
111
112    /// Create a new `BPlusTree` from existing parts
113    fn new_from_parts(mut node_store: S, root: NodeId, len: usize) -> Self {
114        let argument = if len > 0 {
115            Self::new_argument_for_id(&mut node_store, root)
116        } else {
117            S::Argument::default()
118        };
119
120        let me = Self {
121            root,
122            root_argument: argument,
123            node_store: ManuallyDrop::new(node_store),
124            len,
125
126            st: Statistic::default(),
127        };
128
129        #[cfg(test)]
130        me.validate();
131
132        me
133    }
134
135    /// Gets a reference to the `NodeStore` that this `BPlusTree` was created with.
136    pub fn node_store(&self) -> &S {
137        &self.node_store
138    }
139
140    /// Returns the number of elements in the tree.
141    pub fn len(&self) -> usize {
142        self.len
143    }
144
145    /// Returns true if the tree contains no elements.
146    pub fn is_empty(&self) -> bool {
147        self.len() == 0
148    }
149
150    /// Returns a reference to root argument
151    pub fn root_argument(&self) -> &S::Argument {
152        &self.root_argument
153    }
154
155    /// Insert a new key-value pair into the tree.
156    pub fn insert(&mut self, k: S::K, v: S::V) -> Option<S::V> {
157        let node_id = self.root;
158
159        let result = match self.descend_insert(node_id, k, v) {
160            DescendInsertResult::Inserted => {
161                self.root_argument = Self::new_argument_for_id(&self.node_store, node_id);
162                None
163            }
164            DescendInsertResult::Updated(prev_v) => Some(prev_v),
165            DescendInsertResult::Split(k, new_child_id) => {
166                let prev_root_argument = Self::new_argument_for_id(&self.node_store, node_id);
167                let new_argument = Self::new_argument_for_id(&self.node_store, new_child_id);
168
169                let new_root = InnerNode::<S::K, S::Argument>::new(
170                    [k],
171                    [node_id, new_child_id],
172                    [prev_root_argument, new_argument],
173                );
174                let new_root_argument =
175                    S::Argument::from_inner(new_root.keys(), new_root.arguments());
176
177                let new_root_id = self.node_store.add_inner(new_root);
178                self.root = new_root_id.into();
179                self.root_argument = new_root_argument;
180                None
181            }
182        };
183
184        if result.is_none() {
185            self.len += 1;
186        }
187
188        #[cfg(test)]
189        self.validate();
190
191        result
192    }
193
194    /// consume self and return the parts. This is useful when implementing `IntoIter`
195    fn into_parts(self) -> (S, NodeId, usize) {
196        let mut me = ManuallyDrop::new(self);
197        (
198            unsafe { ManuallyDrop::take(&mut me.node_store) },
199            me.root,
200            me.len,
201        )
202    }
203
204    pub fn statistic(&self) -> &Statistic {
205        &self.st
206    }
207
208    fn descend_insert(
209        &mut self,
210        node_id: NodeId,
211        k: S::K,
212        v: S::V,
213    ) -> DescendInsertResult<S::K, S::V> {
214        match node_id {
215            NodeId::Inner(id) => self.insert_inner(id, k, v),
216            NodeId::Leaf(leaf_id) => self.insert_leaf(leaf_id, k, v),
217        }
218    }
219
220    fn insert_inner(
221        &mut self,
222        mut id: InnerNodeId,
223        k: S::K,
224        v: S::V,
225    ) -> DescendInsertResult<S::K, S::V> {
226        let mut stack = VisitStack::new();
227        let mut r = loop {
228            let node = self.node_store.get_inner(id);
229            let (child_idx, child_id) = node.locate_child(&k);
230            stack.push(id, child_idx, child_id);
231
232            match child_id {
233                NodeId::Inner(child_inner) => {
234                    id = child_inner;
235                    continue;
236                }
237                NodeId::Leaf(leaf_id) => break self.insert_leaf(leaf_id, k, v),
238            }
239        };
240
241        loop {
242            // ascend process. Need to process split and update argument
243            match stack.pop() {
244                Some((id, child_idx, child_id)) => match r {
245                    DescendInsertResult::Split(key, right_child) => {
246                        let right_child_argument =
247                            Self::new_argument_for_id(&mut self.node_store, right_child);
248                        let child_argument =
249                            Self::new_argument_for_id(&mut self.node_store, child_id);
250
251                        let inner_node = self.node_store.get_mut_inner(id);
252                        // it's easier to update argument here, the split logic handles arguments split
253                        // too.
254                        inner_node.set_argument(child_idx, child_argument);
255
256                        if !inner_node.is_full() {
257                            let slot = child_idx;
258                            inner_node.insert_at(slot, key, right_child, right_child_argument);
259                            r = DescendInsertResult::Inserted;
260                        } else {
261                            let (prompt_k, new_node) =
262                                inner_node.split(child_idx, key, right_child, right_child_argument);
263                            let new_node_id = self.node_store.add_inner(new_node);
264                            r = DescendInsertResult::Split(prompt_k, NodeId::Inner(new_node_id));
265                        }
266                    }
267                    DescendInsertResult::Inserted => {
268                        let child_argument =
269                            Self::new_argument_for_id(&mut self.node_store, child_id);
270                        let inner_node = self.node_store.get_mut_inner(id);
271                        inner_node.set_argument(child_idx, child_argument);
272
273                        continue;
274                    }
275                    DescendInsertResult::Updated(_) => {
276                        // the key didn't change, so does the argument
277                        continue;
278                    }
279                },
280                None => return r,
281            }
282        }
283    }
284
285    fn new_argument_for_id(node_store: &S, id: NodeId) -> S::Argument {
286        match id {
287            NodeId::Inner(inner) => {
288                let inner = node_store.get_inner(inner);
289                S::Argument::from_inner(inner.keys(), inner.arguments())
290            }
291            NodeId::Leaf(leaf) => {
292                let leaf = node_store.get_leaf(leaf);
293                S::Argument::from_leaf(leaf.keys())
294            }
295        }
296    }
297
298    fn insert_leaf(&mut self, id: LeafNodeId, k: S::K, v: S::V) -> DescendInsertResult<S::K, S::V> {
299        let leaf_node = self.node_store.get_mut_leaf(id);
300        match leaf_node.try_upsert(k, v) {
301            LeafUpsertResult::Inserted => {
302                self.node_store.cache_leaf(id);
303                DescendInsertResult::Inserted
304            }
305            LeafUpsertResult::Updated(v) => DescendInsertResult::Updated(v),
306            LeafUpsertResult::IsFull(idx, k, v) => {
307                let right_id = self.node_store.reserve_leaf();
308
309                let l_leaf = self.node_store.get_mut_leaf(id);
310                let r_leaf = l_leaf.split_new_leaf(idx, (k, v), right_id, id);
311                let slot_key: S::K = r_leaf.data_at(0).0.clone();
312
313                // fix r_leaf's next's prev
314                if let Some(next) = r_leaf.next() {
315                    self.node_store.get_mut_leaf(next).set_prev(Some(right_id));
316                }
317                self.node_store.assign_leaf(right_id, r_leaf);
318
319                let updated_id = if idx >= S::leaf_n() as usize / 2 {
320                    right_id
321                } else {
322                    id
323                };
324                self.node_store.cache_leaf(updated_id);
325                DescendInsertResult::Split(slot_key, NodeId::Leaf(right_id))
326            }
327        }
328    }
329
330    /// Get reference to value identified by key.
331    pub fn get<Q: ?Sized + Ord>(&self, k: &Q) -> Option<&S::V>
332    where
333        S::K: Borrow<Q>,
334    {
335        if let Some(leaf_id) = self.node_store.try_cache(k) {
336            // cache hit
337            return self.find_in_leaf(leaf_id, k);
338        }
339
340        self.find_descend(self.root, k)
341    }
342
343    /// Get mutable reference to value identified by key.
344    pub fn get_mut<Q: ?Sized + Ord>(&mut self, k: &Q) -> Option<&mut S::V>
345    where
346        S::K: Borrow<Q>,
347    {
348        let mut cache_leaf_id: Option<LeafNodeId> = None;
349        if let Some(leaf_id) = self.node_store.try_cache(k) {
350            // cache hit
351            cache_leaf_id = Some(leaf_id);
352        }
353
354        if let Some(leaf_id) = cache_leaf_id {
355            return self.find_in_leaf_mut(leaf_id, k);
356        }
357
358        self.find_descend_mut(self.root, k)
359    }
360
361    /// Returns first key-value pair in the map.
362    pub fn first(&self) -> Option<(&S::K, &S::V)> {
363        if self.is_empty() {
364            return None;
365        }
366
367        let mut node_id = self.root;
368        loop {
369            match node_id {
370                NodeId::Inner(inner_id) => {
371                    let inner_node = self.node_store.get_inner(inner_id);
372                    node_id = inner_node.child_id(0);
373                    continue;
374                }
375                NodeId::Leaf(leaf_id) => {
376                    let leaf_node = self.node_store.get_leaf(leaf_id);
377                    return Some(leaf_node.data_at(0));
378                }
379            }
380        }
381    }
382
383    /// Returns last key-value pair in the map.
384    pub fn last(&self) -> Option<(&S::K, &S::V)> {
385        if self.is_empty() {
386            return None;
387        }
388
389        let mut node_id = self.root;
390        loop {
391            match node_id {
392                NodeId::Inner(inner_id) => {
393                    let inner_node = self.node_store.get_inner(inner_id);
394                    node_id = inner_node.child_id(inner_node.len());
395                    continue;
396                }
397                NodeId::Leaf(leaf_id) => {
398                    let leaf_node = self.node_store.get_leaf(leaf_id);
399                    return Some(leaf_node.data_at(leaf_node.len() - 1));
400                }
401            }
402        }
403    }
404
405    fn find_descend<Q: ?Sized + Ord>(&self, mut node_id: NodeId, k: &Q) -> Option<&S::V>
406    where
407        S::K: Borrow<Q>,
408    {
409        loop {
410            match node_id {
411                NodeId::Inner(inner_id) => {
412                    let inner_node = self.node_store.get_inner(inner_id);
413                    let (_, child_id) = inner_node.locate_child(k);
414                    node_id = child_id;
415                    continue;
416                }
417                NodeId::Leaf(leaf_id) => return self.find_in_leaf_and_cache_it(leaf_id, k),
418            }
419        }
420    }
421
422    fn find_in_leaf<Q: ?Sized + Ord>(&self, leaf_id: LeafNodeId, k: &Q) -> Option<&S::V>
423    where
424        S::K: Borrow<Q>,
425    {
426        let leaf_node = self.node_store.get_leaf(leaf_id);
427        let (_, kv) = leaf_node.locate_slot_with_value(k);
428        kv
429    }
430
431    fn find_descend_mut<Q: ?Sized + Ord>(&mut self, node_id: NodeId, k: &Q) -> Option<&mut S::V>
432    where
433        S::K: Borrow<Q>,
434    {
435        match node_id {
436            NodeId::Inner(inner_id) => {
437                let inner_node = self.node_store.get_inner(inner_id);
438                let (_, child_id) = inner_node.locate_child(k);
439                self.find_descend_mut(child_id, k)
440            }
441            NodeId::Leaf(leaf_id) => self.find_in_leaf_mut_and_cache_it(leaf_id, k),
442        }
443    }
444
445    fn find_in_leaf_mut<Q: ?Sized + Ord>(&mut self, leaf_id: LeafNodeId, k: &Q) -> Option<&mut S::V>
446    where
447        S::K: Borrow<Q>,
448    {
449        let leaf_node = self.node_store.get_mut_leaf(leaf_id);
450        let (_, v) = leaf_node.locate_slot_mut(k);
451        v
452    }
453
454    fn find_in_leaf_mut_and_cache_it<Q: ?Sized + Ord>(
455        &mut self,
456        leaf_id: LeafNodeId,
457        k: &Q,
458    ) -> Option<&mut S::V>
459    where
460        S::K: Borrow<Q>,
461    {
462        self.node_store.cache_leaf(leaf_id);
463        let leaf_node = self.node_store.get_mut_leaf(leaf_id);
464        let (_, kv) = leaf_node.locate_slot_mut(k);
465        kv
466    }
467
468    /// Find the key in leaf, and cache leaf, this method only called when
469    /// cache miss
470    fn find_in_leaf_and_cache_it<Q: ?Sized + Ord>(
471        &self,
472        leaf_id: LeafNodeId,
473        k: &Q,
474    ) -> Option<&S::V>
475    where
476        S::K: Borrow<Q>,
477    {
478        self.node_store.cache_leaf(leaf_id);
479        let leaf = self.node_store.get_leaf(leaf_id);
480        let (_, v) = leaf.locate_slot_with_value(k);
481        v
482    }
483
484    /// delete element identified by K
485    pub fn remove<Q: ?Sized>(&mut self, k: &Q) -> Option<S::V>
486    where
487        Q: Ord,
488        S::K: Borrow<Q>,
489    {
490        self.remove_impl(k).map(|kv| kv.1)
491    }
492
493    fn key_to_ref<'a, Q: ?Sized>(&'a self, k: &Q) -> Option<EntryRef<&Self>>
494    where
495        Q: Ord,
496        S::K: Borrow<Q>,
497    {
498        let mut inner_stack = VisitStack::new();
499        let mut node_id = self.root;
500        loop {
501            match node_id {
502                NodeId::Inner(inner_id) => {
503                    let inner_node = self.node_store.get_inner(inner_id);
504                    let (child_offset, child_id) = inner_node.locate_child(k);
505                    inner_stack.push(inner_id, child_offset, child_id);
506
507                    node_id = child_id;
508                }
509                NodeId::Leaf(leaf_id) => {
510                    let leaf = self.node_store.get_leaf(leaf_id);
511
512                    match leaf.locate_slot(k) {
513                        Ok(idx) => return Some(EntryRef::new(self, inner_stack, leaf_id, idx)),
514                        Err(_) => return None,
515                    };
516                }
517            }
518        }
519    }
520
521    /// get the first leaf_id if exists
522    pub(crate) fn first_leaf(&self) -> Option<LeafNodeId> {
523        match self.root {
524            NodeId::Inner(inner_id) => {
525                let mut result = None;
526
527                self.descend_visit_inner(inner_id, |inner_node| {
528                    let first_child_id = inner_node.child_id(0);
529                    match first_child_id {
530                        NodeId::Inner(inner_id) => Some(inner_id),
531                        NodeId::Leaf(leaf_id) => {
532                            result = Some(leaf_id);
533                            None
534                        }
535                    }
536                });
537
538                result
539            }
540            NodeId::Leaf(leaf_id) => Some(leaf_id),
541        }
542    }
543
544    /// get the last leaf_id if exists
545    pub(crate) fn last_leaf(&self) -> Option<LeafNodeId> {
546        match self.root {
547            NodeId::Inner(inner_id) => {
548                let mut result = None;
549
550                self.descend_visit_inner(inner_id, |inner_node| {
551                    let child_id = inner_node.child_id(inner_node.len());
552                    match child_id {
553                        NodeId::Inner(inner_id) => Some(inner_id),
554                        NodeId::Leaf(leaf_id) => {
555                            result = Some(leaf_id);
556                            None
557                        }
558                    }
559                });
560
561                result
562            }
563            NodeId::Leaf(leaf_id) => Some(leaf_id),
564        }
565    }
566
567    /// Locate the leaf node for `k`.
568    /// Returns the leaf whose range contains `k`.
569    /// User should query the leaf and check key existance.
570    pub(crate) fn locate_leaf(&self, k: &S::K) -> Option<LeafNodeId> {
571        if let Some(leaf_id) = self.node_store.try_cache(k) {
572            return Some(leaf_id);
573        }
574
575        let leaf_id = match self.root {
576            NodeId::Inner(inner_id) => {
577                let mut result = None;
578                self.descend_visit_inner(inner_id, |inner_node| {
579                    let (_idx, node_id) = inner_node.locate_child(k);
580                    match node_id {
581                        NodeId::Inner(inner_node) => Some(inner_node),
582                        NodeId::Leaf(leaf_id) => {
583                            result = Some(leaf_id);
584                            None
585                        }
586                    }
587                });
588                result
589            }
590            NodeId::Leaf(leaf_id) => Some(leaf_id),
591        }?;
592
593        Some(leaf_id)
594    }
595
596    fn descend_visit_inner(
597        &self,
598        mut node_id: InnerNodeId,
599        mut f: impl FnMut(&InnerNode<S::K, S::Argument>) -> Option<InnerNodeId>,
600    ) -> Option<()> {
601        loop {
602            let inner = self.node_store.get_inner(node_id);
603            match f(inner) {
604                None => {
605                    return None;
606                }
607                Some(id_to_visit) => node_id = id_to_visit,
608            }
609        }
610    }
611
612    /// Create an iterator on (&K, &V) pairs
613    pub fn iter(&self) -> iterator::Iter<S> {
614        iterator::Iter::new(self)
615    }
616
617    /// Consume the tree and create an iterator on (K, &V) pairs
618    pub fn into_iter(self) -> iterator::IntoIter<S> {
619        iterator::IntoIter::new(self)
620    }
621
622    /// Create an cursor from first elem
623    pub fn cursor_first(&self) -> Option<Cursor<S::K>> {
624        Cursor::first(self).map(|c| c.0)
625    }
626
627    /// Create an cursor for k
628    pub fn get_cursor(&self, k: &S::K) -> Option<(Cursor<S::K>, Option<&S::V>)> {
629        let node_id = self.root;
630        let leaf_id = match node_id {
631            NodeId::Inner(inner_id) => {
632                let mut result = None;
633                self.descend_visit_inner(inner_id, |inner_node| {
634                    let (_idx, node_id) = inner_node.locate_child(k);
635                    match node_id {
636                        NodeId::Inner(inner_node) => Some(inner_node),
637                        NodeId::Leaf(leaf_id) => {
638                            result = Some(leaf_id);
639                            None
640                        }
641                    }
642                });
643                result
644            }
645            NodeId::Leaf(leaf_id) => Some(leaf_id),
646        }?;
647
648        let leaf = self.node_store.get_leaf(leaf_id);
649        let (idx, v) = leaf.locate_slot_with_value(k);
650        Some((Cursor::new(k.clone(), leaf_id, idx), v))
651    }
652
653    /// Clear the tree
654    pub fn clear(&mut self) {
655        // todo: should we keep the node_store's capacity?
656        std::mem::drop(std::mem::replace(self, Self::new(S::default())));
657    }
658
659    /// get by argument
660    fn get_ref_by_argument<Q>(&self, mut query: Q) -> Option<EntryRef<&Self>>
661    where
662        S::Argument: SearchArgument<S::K, Query = Q>,
663    {
664        let mut node_id = self.root;
665        let mut stack = VisitStack::new();
666
667        loop {
668            match node_id {
669                NodeId::Inner(inner_id) => {
670                    let inner = self.node_store.get_inner(inner_id);
671                    let (offset, new_query) = <S::Argument as SearchArgument<_>>::locate_in_inner(
672                        query,
673                        inner.keys(),
674                        inner.arguments(),
675                    )?;
676                    node_id = inner.child_id(offset);
677
678                    stack.push(inner_id, offset, node_id);
679                    query = new_query;
680                }
681                NodeId::Leaf(leaf_id) => {
682                    let leaf = self.node_store.get_leaf(leaf_id);
683                    let slot =
684                        <S::Argument as SearchArgument<_>>::locate_in_leaf(query, leaf.keys())?;
685
686                    return Some(EntryRef::new(self, stack, leaf_id, slot));
687                }
688            }
689        }
690    }
691
692    /// get by argument
693    pub fn get_by_argument<Q>(&self, query: Q) -> Option<(&S::K, &S::V)>
694    where
695        S::Argument: SearchArgument<S::K, Query = Q>,
696    {
697        let entry_ref = self.get_ref_by_argument(query)?;
698        Self::get_by_ref(entry_ref)
699    }
700
701    /// get mut reference to value by argument's Query
702    pub fn get_mut_by_argument<Q>(&mut self, query: Q) -> Option<&mut S::V>
703    where
704        S::Argument: SearchArgument<S::K, Query = Q>,
705    {
706        let entry_ref = self.get_ref_by_argument(query)?;
707        Some(Self::get_mut_by_ref(entry_ref.to_detached().to_ref(self)))
708    }
709
710    /// Get rank for argument
711    pub fn rank_by_argument<R>(&self, k: &S::K) -> Result<R, R>
712    where
713        S::Argument: RankArgument<S::K, Rank = R>,
714    {
715        let mut node_id = self.root;
716        let mut rank = <S::Argument as RankArgument<S::K>>::initial_value();
717
718        loop {
719            match node_id {
720                NodeId::Inner(inner_id) => {
721                    let inner = self.node_store.get_inner(inner_id);
722                    let (child_idx, child_id) = inner.locate_child(k);
723                    node_id = child_id;
724                    let arguments = &inner.arguments()[0..child_idx];
725                    rank = <S::Argument as RankArgument<S::K>>::fold_inner(k, rank, arguments);
726                }
727                NodeId::Leaf(leaf_id) => {
728                    let leaf = self.node_store.get_leaf(leaf_id);
729                    let slot = leaf.locate_slot(k);
730                    return <S::Argument as RankArgument<_>>::fold_leaf(k, rank, slot, leaf.keys());
731                }
732            }
733        }
734    }
735
736    /// remove by argument
737    pub fn remove_by_argument<Q>(&mut self, query: Q) -> Option<(S::K, S::V)>
738    where
739        S::Argument: SearchArgument<S::K, Query = Q>,
740    {
741        let entry_ref = self.get_ref_by_argument(query)?;
742        Self::remove_by_ref(entry_ref.to_detached().to_ref(self))
743    }
744
745    /// Get the (&K, &V) pair for referece
746    fn get_by_ref(entry_ref: EntryRef<&Self>) -> Option<(&S::K, &S::V)> {
747        let leaf = entry_ref.tree.node_store.get_leaf(entry_ref.leaf_id);
748        let slot = entry_ref.offset;
749        Some(leaf.data_at(slot))
750    }
751
752    /// Get the &mut V for referece
753    fn get_mut_by_ref(entry_ref: EntryRef<&mut Self>) -> &mut S::V {
754        let EntryRef {
755            tree,
756            leaf_id,
757            offset,
758            ..
759        } = entry_ref;
760
761        let leaf = tree.node_store.get_mut_leaf(leaf_id);
762        let slot = offset;
763        leaf.value_at_mut(slot)
764    }
765
766    #[cfg(test)]
767    fn validate(&self) {
768        let Some(mut leaf_id) = self.first_leaf() else { return; };
769        let mut last_leaf_id: Option<LeafNodeId> = None;
770
771        // ensures all prev and next are correct
772        loop {
773            let leaf = self.node_store.get_leaf(leaf_id);
774
775            let p = leaf.prev();
776            let n = leaf.next();
777
778            if let Some(last_leaf_id) = last_leaf_id {
779                assert_eq!(last_leaf_id, p.unwrap());
780            }
781
782            if n.is_none() {
783                break;
784            }
785
786            last_leaf_id = Some(leaf_id);
787            leaf_id = n.unwrap();
788        }
789    }
790}
791
792impl<S: NodeStore> Drop for BPlusTree<S> {
793    fn drop(&mut self) {
794        unsafe { drop(std::ptr::read(self).into_iter()) }
795    }
796}
797
798/// Statistic data used to guide the perf tuning
799#[derive(Default, Debug, Clone)]
800pub struct Statistic {
801    pub rotate_right_inner: u64,
802    pub rotate_left_inner: u64,
803
804    pub merge_with_left_inner: u64,
805    pub merge_with_right_inner: u64,
806
807    pub rotate_right_leaf: u64,
808    pub rotate_left_leaf: u64,
809
810    pub merge_with_left_leaf: u64,
811    pub merge_with_right_leaf: u64,
812}
813
814enum DescendInsertResult<K, V> {
815    /// Update existing value, V is the previous value
816    Updated(V),
817    /// Inserted, and not split
818    Inserted,
819    /// need to split
820    Split(K, NodeId),
821}
822
823/// NodeStore is the node storage for tree, responsible for
824/// define node types, manage node memory, and provide node access
825pub trait NodeStore: Default {
826    /// Key type for the tree
827    type K: Key;
828    /// Value type for the tree
829    type V;
830
831    /// The Argument type
832    type Argument: Argument<Self::K>;
833
834    /// Get the max number of keys inner node can hold
835    fn inner_n() -> u16;
836    /// Get the max number of elements leaf node can hold
837    fn leaf_n() -> u16;
838
839    /// Create an empty inner node and returns its id
840    #[cfg(test)]
841    fn new_empty_inner(&mut self) -> InnerNodeId;
842
843    /// Add the inner node to the store and returns its id
844    fn add_inner(&mut self, node: Box<InnerNode<Self::K, Self::Argument>>) -> InnerNodeId;
845
846    /// Get the inner node
847    /// # Panics
848    /// if id is invalid or the node is already removed, panic
849    fn get_inner(&self, id: InnerNodeId) -> &InnerNode<Self::K, Self::Argument>;
850
851    /// Get the inner node
852    /// if id is invalid or the node already removed, remove None
853    fn try_get_inner(&self, id: InnerNodeId) -> Option<&InnerNode<Self::K, Self::Argument>>;
854
855    /// Get a mut reference to the `InnerNode`
856    fn get_mut_inner(&mut self, id: InnerNodeId) -> &mut InnerNode<Self::K, Self::Argument>;
857
858    /// Get a mut pointer to inner node.
859    /// User must ensure there is non shared reference to the node co-exists
860    unsafe fn get_mut_inner_ptr(
861        &mut self,
862        id: InnerNodeId,
863    ) -> *mut InnerNode<Self::K, Self::Argument>;
864
865    /// Take the inner node out of the store
866    fn take_inner(&mut self, id: InnerNodeId) -> Box<InnerNode<Self::K, Self::Argument>>;
867
868    /// Put back the inner node
869    fn put_back_inner(&mut self, id: InnerNodeId, node: Box<InnerNode<Self::K, Self::Argument>>);
870
871    /// Create a new empty leaf node and returns its id
872    fn new_empty_leaf(&mut self) -> (LeafNodeId, &mut LeafNode<Self::K, Self::V>);
873
874    /// Reserve a leaf node, it must be assigned later
875    fn reserve_leaf(&mut self) -> LeafNodeId;
876
877    /// Get a refernce to leaf node
878    /// Panics if id is invalid or the node is taken
879    fn get_leaf(&self, id: LeafNodeId) -> &LeafNode<Self::K, Self::V>;
880
881    /// Get a reference to leaf node
882    /// Returns None if id is invalid or the node is taken
883    fn try_get_leaf(&self, id: LeafNodeId) -> Option<&LeafNode<Self::K, Self::V>>;
884
885    /// Get a mut reference to leaf node
886    /// Panics if id is invalid or the node is taken
887    fn get_mut_leaf(&mut self, id: LeafNodeId) -> &mut LeafNode<Self::K, Self::V>;
888
889    /// Take the leaf out of store
890    fn take_leaf(&mut self, id: LeafNodeId) -> Box<LeafNode<Self::K, Self::V>>;
891
892    /// Assign the leaf to the id, the id must exists
893    fn assign_leaf(&mut self, id: LeafNodeId, leaf: Box<LeafNode<Self::K, Self::V>>);
894
895    /// cache leaf
896    fn cache_leaf(&self, leaf_id: LeafNodeId);
897
898    /// try cache for k
899    fn try_cache<Q: ?Sized>(&self, k: &Q) -> Option<LeafNodeId>
900    where
901        Q: Ord,
902        Self::K: Borrow<Q>;
903
904    #[cfg(test)]
905    fn debug(&self)
906    where
907        Self::K: std::fmt::Debug,
908        Self::V: std::fmt::Debug + Clone,
909        Self::Argument: std::fmt::Debug;
910}
911
912/// Key trait
913/// `Clone` is required since Inner Node may dup the key.
914pub trait Key: Clone + Ord {}
915
916impl<T> Key for T where T: Clone + Ord {}
917
918/// ensure NodeStoreVec is send for send v
919fn _ensure_send<V: Send>() {
920    fn _assert_send<T: Send>() {}
921    _assert_send::<BPlusTree<NodeStoreVec<u64, V>>>();
922}
923
924#[cfg(test)]
925mod tests {
926    use std::rc::Rc;
927
928    use rand::seq::SliceRandom;
929
930    use crate::argument::count::Count;
931
932    use super::*;
933
934    #[test]
935    fn test_round_trip_one() {
936        round_trip_one();
937    }
938
939    fn round_trip_one() {
940        let node_store = NodeStoreVec::<i64, i64, Count>::new();
941        let mut tree = BPlusTree::new(node_store);
942
943        let size: i64 = 100000;
944
945        let mut keys = (0..size).collect::<Vec<_>>();
946        keys.shuffle(&mut rand::thread_rng());
947
948        for i in keys {
949            tree.insert(i, i % 13);
950            assert_eq!(tree.root_argument.count(), tree.len());
951        }
952
953        let mut keys = (0..size).collect::<Vec<_>>();
954        keys.shuffle(&mut rand::thread_rng());
955        for i in keys {
956            assert_eq!(*tree.get(&i).unwrap(), i % 13);
957        }
958
959        let mut keys = (0..size).collect::<Vec<_>>();
960        keys.shuffle(&mut rand::thread_rng());
961
962        for i in keys {
963            let k = i;
964            println!("\ndeleting {i}");
965
966            let delete_result = tree.remove(&k);
967            assert!(delete_result.is_some());
968            assert_eq!(tree.root_argument.count(), tree.len());
969        }
970
971        assert!(tree.is_empty());
972
973        // call clear on empty tree
974        tree.clear();
975    }
976
977    #[test]
978    fn test_tree_clear() {
979        let (mut tree, keys) = create_test_tree::<100>();
980        tree.clear();
981        assert!(tree.is_empty());
982
983        // insert after clear
984        for k in keys.clone() {
985            tree.insert(k, k % 13);
986        }
987        assert_eq!(tree.len(), keys.len());
988    }
989
990    #[test]
991    fn test_first_leaf() {
992        let node_store = NodeStoreVec::<i64, i64>::new();
993        let mut tree = BPlusTree::new(node_store);
994        let size: i64 = 500;
995        let keys = (0..size).collect::<Vec<_>>();
996        for i in keys {
997            tree.insert(i, i % 13);
998        }
999
1000        let first_leaf_id = tree.first_leaf().unwrap();
1001        let first_leaf = tree.node_store.get_leaf(first_leaf_id);
1002        assert_eq!(first_leaf.data_at(0).0.clone(), 0);
1003    }
1004
1005    #[test]
1006    fn test_get_mut() {
1007        let (mut tree, _) = create_test_tree::<30>();
1008        let v = tree.get_mut(&1).unwrap();
1009        *v = 100;
1010        assert_eq!(tree.get(&1).unwrap().clone(), 100);
1011    }
1012
1013    #[test]
1014    fn test_cursor() {
1015        let (mut tree, _) = create_test_tree::<30>();
1016
1017        let (cursor, _kv) = tree.get_cursor(&10).unwrap();
1018        assert_eq!(cursor.key().clone(), 10);
1019        assert_eq!(cursor.value(&tree).unwrap().clone(), 10);
1020
1021        {
1022            let prev = cursor.prev(&tree).unwrap();
1023            assert_eq!(prev.key().clone(), 9);
1024
1025            let next = cursor.next(&tree).unwrap();
1026            assert_eq!(next.key().clone(), 11);
1027        }
1028
1029        tree.remove(&10);
1030
1031        {
1032            assert_eq!(cursor.key().clone(), 10);
1033            assert!(cursor.value(&tree).is_none());
1034
1035            let prev = cursor.prev(&tree).unwrap();
1036            assert_eq!(prev.key().clone(), 9);
1037
1038            let next = cursor.next(&tree).unwrap();
1039            assert_eq!(next.key().clone(), 11);
1040        }
1041
1042        let (cursor, kv) = tree.get_cursor(&10).unwrap();
1043        assert_eq!(cursor.key().clone(), 10);
1044        assert!(kv.is_none());
1045    }
1046
1047    pub fn create_test_tree<const N: usize>() -> (BPlusTree<NodeStoreVec<i64, i64>>, Vec<i64>) {
1048        let node_store = NodeStoreVec::<i64, i64>::new();
1049        let mut tree = BPlusTree::new(node_store);
1050
1051        let size: i64 = N as i64;
1052
1053        let mut keys = (0..size).collect::<Vec<_>>();
1054        keys.shuffle(&mut rand::thread_rng());
1055
1056        for i in keys.iter() {
1057            tree.insert(*i, i % 13);
1058        }
1059
1060        assert_eq!(tree.len(), N);
1061
1062        (tree, keys)
1063    }
1064
1065    struct TestKey {
1066        key: i32,
1067        counter: Rc<std::sync::atomic::AtomicU64>,
1068        panic_flag: Rc<std::sync::atomic::AtomicU64>,
1069    }
1070
1071    impl Ord for TestKey {
1072        fn cmp(&self, other: &Self) -> std::cmp::Ordering {
1073            self.key.cmp(&other.key)
1074        }
1075    }
1076
1077    impl PartialOrd for TestKey {
1078        fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
1079            self.key.partial_cmp(&other.key)
1080        }
1081    }
1082
1083    impl PartialEq for TestKey {
1084        fn eq(&self, other: &Self) -> bool {
1085            self.key == other.key
1086        }
1087    }
1088
1089    impl Eq for TestKey {}
1090
1091    impl Clone for TestKey {
1092        fn clone(&self) -> Self {
1093            self.counter
1094                .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
1095            Self {
1096                key: self.key,
1097                panic_flag: self.panic_flag.clone(),
1098                counter: self.counter.clone(),
1099            }
1100        }
1101    }
1102
1103    impl TestKey {
1104        fn new(
1105            key: i32,
1106            counter: Rc<std::sync::atomic::AtomicU64>,
1107            panic_flag: Rc<std::sync::atomic::AtomicU64>,
1108        ) -> Self {
1109            counter.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
1110            Self {
1111                key,
1112                counter,
1113                panic_flag,
1114            }
1115        }
1116    }
1117
1118    impl Drop for TestKey {
1119        fn drop(&mut self) {
1120            self.panic_flag
1121                .fetch_sub(1, std::sync::atomic::Ordering::Relaxed);
1122            self.counter
1123                .fetch_sub(1, std::sync::atomic::Ordering::Relaxed);
1124        }
1125    }
1126
1127    #[derive(Clone)]
1128    struct TestValue {
1129        counter: Rc<std::sync::atomic::AtomicU64>,
1130    }
1131
1132    impl TestValue {
1133        fn new(counter: Rc<std::sync::atomic::AtomicU64>) -> Self {
1134            Self { counter }
1135        }
1136    }
1137
1138    impl Drop for TestValue {
1139        fn drop(&mut self) {
1140            self.counter
1141                .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
1142        }
1143    }
1144
1145    #[test]
1146    fn test_drop() {
1147        let count: u64 = 16000;
1148        // test drop
1149        let node_store = NodeStoreVec::<TestKey, TestValue>::new();
1150        let mut tree = BPlusTree::new(node_store);
1151        let drop_counter = Rc::new(std::sync::atomic::AtomicU64::new(0));
1152        let key_counter = Rc::new(std::sync::atomic::AtomicU64::new(0));
1153        let panic_flag = Rc::new(std::sync::atomic::AtomicU64::new(0));
1154
1155        macro_rules! get {
1156            ($id: ident) => {
1157                $id.load(std::sync::atomic::Ordering::Relaxed)
1158            };
1159        }
1160
1161        for i in 0..count {
1162            tree.insert(
1163                TestKey::new(i as i32, key_counter.clone(), panic_flag.clone()),
1164                TestValue::new(drop_counter.clone()),
1165            );
1166        }
1167
1168        {
1169            let key_count_inserted = key_counter.load(std::sync::atomic::Ordering::Relaxed);
1170            let another_tree = tree.clone();
1171            let key_count_cloned = get!(key_counter);
1172            assert_eq!(key_count_cloned, key_count_inserted * 2);
1173            drop(another_tree);
1174
1175            let key_count_drop_clone = key_counter.load(std::sync::atomic::Ordering::Relaxed);
1176            assert_eq!(key_count_drop_clone, key_count_inserted);
1177        }
1178
1179        {
1180            let keys = tree.iter().map(|(k, _v)| k.clone()).collect::<Vec<_>>();
1181            for k in keys {
1182                panic_flag.store(1, std::sync::atomic::Ordering::Relaxed);
1183
1184                let prev = get!(key_counter);
1185                assert!(tree.remove(&k).is_some());
1186                let delta = prev - get!(key_counter);
1187                if delta > 3 {
1188                    panic!();
1189                }
1190
1191                // one for k
1192                panic_flag.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
1193            }
1194        }
1195
1196        drop(tree);
1197
1198        assert_eq!(
1199            drop_counter.load(std::sync::atomic::Ordering::Relaxed),
1200            count * 2,
1201        );
1202        assert_eq!(key_counter.load(std::sync::atomic::Ordering::Relaxed), 0);
1203    }
1204}
sweep_bptree/tree/mod.rs

sweep_bptree/tree/
mod.rs
