Skip to main content

embed_collections/btree/
mod.rs

1//! B+Tree Map - A in memory cache-optimized B+Tree for single-threaded use.
2//!
3//! ## Feature outlines
4//! - Designed for CPU cache efficiency and memory efficiency
5//! - Optimised for numeric key type
6//!   - Typical scenario: [range-tree-rs](https://docs.rs/range-tree-rs)
7//! - A B+tree. Data stores only at leaf level, with links at leaf level.
8//!   - Provides efficient iteration of data
9//!   - Linear search within nodes, respecting cacheline boundaries
10//! - Nodes are filled up in 4 cache lines (256 bytes on x86_64)
11//!   - keys stored in first 128B (with header)
12//!   - Values/pointers stored in last 128B
13//!   - the capacity is calculated according to the size of K, V, with limitations:
14//!     - K & V should <= CACHE_LINE_SIZE - 16  (If K & V is larger should put into `Box`)
15//! - Specially Entry API which allow to modify after moving the cursor to adjacent data.
16//! - The detail design notes are with the source in mod.rs and node.rs
17//!
18//! ## Special APIs
19//!
20//! batch removal:
21//! - [BTreeMap::remove_range()]
22//! - [BTreeMap::remove_range_with()]
23//!
24//! adjacent entry:
25//! - [Entry::peak_forward()]
26//! - [Entry::peak_backward()]
27//! - [Entry::move_forward()]
28//! - [Entry::move_backward()]
29//! - [VacantEntry::peak_forward()]
30//! - [VacantEntry::peak_backward()]
31//! - [VacantEntry::move_forward()]
32//! - [VacantEntry::move_backward()]
33//! - [OccupiedEntry::peak_forward()]
34//! - [OccupiedEntry::peak_backward()]
35//! - [OccupiedEntry::move_forward()]
36//! - [OccupiedEntry::move_backward()]
37
38/*
39
40# Designer notes
41
42 Target scenario:  To maintain slab tree for disk, lets say 1 billion nodes, this design is aim for high fan-out.
43
44 Since There're no parent pointer, no fence keys. So we maintain a cache to accelerate the look
45 up for parent nodes.
46 Since we support combining cursor moving in Entry API (try to merge with previous or next adjacent
47 node). But user can call `remove()` on moved cursor.
48
49## Acceleration Search for finding the parent using PathCache
50
51- Assume PathCache is for Node A, Node B is the right brother of node A. A's ptr is at ParentA[idx]
52  - To find parent for Node B
53    - If idx < last (Cap - 1), then A and B has common parent,
54    - otherwise A and B have no common parent. Should continue iteration to upper PathCache. There will be common ancestor until idx < last
55
56- Assume PathCache is for Node B, Node A is the left brother of node B, B's ptr is at ParentB[idx]
57  - To find parent for Node A
58    - If idx > 0, then A and B has common parent.
59    - otherwise A and B have no common parent. Should continue iteration to upper PathCache. There will be common ancestor until idx > 0
60
61## Rebalance
62
63- When insert on a full LeafNode, we try to move items to left/right brother, to delay split operation.
64
65- One entry removed, current node usage < 50% (the threshold can be adjust according to tree size)
66
67  - we don't need to borrow data from brothers (to avoid trashing), and we have already done borrowing on insert.
68
69  - try to merge data, with left + current < cap, or current + right < cap, or left + current + right == 2 * cap
70
71  - No need to mere 3 -> 1, because before reaching 30% average usage, we already checked 3 -> 2.
72
73- The threshold to check merge for InterNode can be lower (like 30%), to avoid trashing
74
75
76## Future ideas
77
78- dynamically incress InterNode size, or variable size
79- borrow from leaf / right on InterNode split, increase fill ratio to 80%
80- 2 - 3 split (increase fill ratio to 90%)
81*/
82
83use core::borrow::Borrow;
84use core::cell::UnsafeCell;
85use core::fmt::Debug;
86use core::ops::{Bound, RangeBounds};
87mod entry;
88pub use entry::*;
89mod node;
90use node::*;
91mod inter;
92use inter::*;
93mod leaf;
94use leaf::*;
95mod iter;
96#[allow(unused_imports)]
97use crate::{print_log, trace_log};
98use iter::RangeBase;
99pub use iter::{IntoIter, Iter, IterMut, Keys, Range, RangeMut, Values, ValuesMut};
100
101#[cfg(test)]
102mod tests;
103
104/// B+Tree Map for single-threaded usage, optimized for numeric type.
105pub struct BTreeMap<K: Ord + Clone + Sized, V: Sized> {
106    /// Root node (may be null for empty tree)
107    root: Option<Node<K, V>>,
108    /// Number of elements in the tree
109    len: usize,
110    // use unsafe to avoid borrow problems
111    cache: UnsafeCell<PathCache<K, V>>,
112    // count of leaf nodes
113    leaf_count: usize,
114    #[cfg(test)]
115    triggers: u32,
116}
117
118#[cfg(test)]
119#[repr(u32)]
120enum TestFlag {
121    LeafSplit = 1,
122    InterSplit = 1 << 1,
123    LeafMoveLeft = 1 << 2,
124    LeafMoveRight = 1 << 3,
125    LeafMergeLeft = 1 << 4,
126    LeafMergeRight = 1 << 5,
127    InterMoveLeft = 1 << 6,
128    InterMoveLeftFirst = 1 << 7,
129    InterMoveRight = 1 << 8,
130    InterMoveRightLast = 1 << 9,
131    InterMergeLeft = 1 << 10,
132    InterMergeRight = 1 << 11,
133    UpdateSepKey = 1 << 12,
134    RemoveOnlyChild = 1 << 13,
135    RemoveChildFirst = 1 << 14,
136    RemoveChildMid = 1 << 15,
137    RemoveChildLast = 1 << 16,
138}
139
140unsafe impl<K: Ord + Clone + Sized + Send, V: Sized + Send> Send for BTreeMap<K, V> {}
141unsafe impl<K: Ord + Clone + Sized + Send, V: Sized + Send> Sync for BTreeMap<K, V> {}
142
143#[cfg(feature = "std")]
144impl<K: Ord + Clone + Sized, V: Sized> std::panic::RefUnwindSafe for BTreeMap<K, V> {}
145
146impl<K: Ord + Sized + Clone, V: Sized> BTreeMap<K, V> {
147    /// Create a new empty BTreeMap
148    pub fn new() -> Self {
149        Self {
150            root: None,
151            len: 0,
152            cache: UnsafeCell::new(PathCache::<K, V>::new()),
153            leaf_count: 0,
154            #[cfg(test)]
155            triggers: 0,
156        }
157    }
158
159    /// Returns the number of elements in the map
160    #[inline(always)]
161    pub fn len(&self) -> usize {
162        self.len
163    }
164
165    /// Returns true if the map is empty
166    #[inline(always)]
167    pub fn is_empty(&self) -> bool {
168        self.len == 0
169    }
170
171    /// return (cap_of_inter_node, cap_of_leaf_node)
172    #[inline]
173    pub const fn cap() -> (u32, u32) {
174        let inter_cap = InterNode::<K, V>::cap();
175        let leaf_cap = LeafNode::<K, V>::cap();
176        (inter_cap, leaf_cap)
177    }
178
179    /// Return the number of leaf nodes
180    #[inline]
181    pub fn leaf_count(&self) -> usize {
182        self.leaf_count
183    }
184
185    #[cfg(all(test, feature = "std"))]
186    pub fn print_trigger_flags(&self) {
187        let mut s = String::from("");
188        if self.triggers & TestFlag::InterSplit as u32 > 0 {
189            s += "InterSplit,";
190        }
191        if self.triggers & TestFlag::LeafSplit as u32 > 0 {
192            s += "LeafSplit,";
193        }
194        if self.triggers & TestFlag::LeafMoveLeft as u32 > 0 {
195            s += "LeafMoveLeft,";
196        }
197        if self.triggers & TestFlag::LeafMoveRight as u32 > 0 {
198            s += "LeafMoveRight,";
199        }
200        if self.triggers & TestFlag::InterMoveLeft as u32 > 0 {
201            s += "InterMoveLeft,";
202        }
203        if self.triggers & TestFlag::InterMoveRight as u32 > 0 {
204            s += "InterMoveRight,";
205        }
206        if s.len() > 0 {
207            print_log!("{s}");
208        }
209        let mut s = String::from("");
210        if self.triggers & TestFlag::InterMergeLeft as u32 > 0 {
211            s += "InterMergeLeft,";
212        }
213        if self.triggers & TestFlag::InterMergeRight as u32 > 0 {
214            s += "InterMergeRight,";
215        }
216        if self.triggers & TestFlag::RemoveOnlyChild as u32 > 0 {
217            s += "RemoveOnlyChild,";
218        }
219        if self.triggers
220            & (TestFlag::RemoveChildFirst as u32
221                | TestFlag::RemoveChildMid as u32
222                | TestFlag::RemoveChildLast as u32)
223            > 0
224        {
225            s += "RemoveChild,";
226        }
227        if s.len() > 0 {
228            print_log!("{s}");
229        }
230    }
231
232    /// Return the average fill ratio of leaf nodes
233    ///
234    /// The range is [0.0, 100]
235    #[inline]
236    pub fn get_fill_ratio(&self) -> f32 {
237        if self.len == 0 {
238            0.0
239        } else {
240            let cap = LeafNode::<K, V>::cap() as usize * self.leaf_count;
241            self.len as f32 / cap as f32 * 100.0
242        }
243    }
244
245    /// When root is leaf, returns 1, otherwise return the number of layers of inter node
246    #[inline(always)]
247    pub fn height(&self) -> u32 {
248        if let Some(Node::Inter(node)) = &self.root {
249            return node.height() + 1;
250        }
251        1
252    }
253
254    #[inline(always)]
255    fn get_cache(&self) -> &mut PathCache<K, V> {
256        unsafe { (&mut *self.cache.get()) as _ }
257    }
258
259    /// Returns an entry to the key in the map
260    #[inline]
261    pub fn entry(&mut self, key: K) -> Entry<'_, K, V> {
262        let cache = self.get_cache();
263        cache.clear();
264        let leaf = match &self.root {
265            None => return Entry::Vacant(VacantEntry { tree: self, key, idx: 0, leaf: None }),
266            Some(root) => root.find_leaf_with_cache(cache, &key),
267        };
268        let (idx, is_equal) = leaf.search(&key);
269        if is_equal {
270            Entry::Occupied(OccupiedEntry { tree: self, idx, leaf })
271        } else {
272            Entry::Vacant(VacantEntry { tree: self, key, idx, leaf: Some(leaf) })
273        }
274    }
275
276    #[inline(always)]
277    fn find<Q>(&self, key: &Q) -> Option<(LeafNode<K, V>, u32)>
278    where
279        K: Borrow<Q>,
280        Q: Ord + ?Sized,
281    {
282        let leaf = match &self.root {
283            None => return None,
284            Some(root) => root.find_leaf(key),
285        };
286        let (idx, is_equal) = leaf.search(key);
287        trace_log!("find leaf {leaf:?} {idx} exist {is_equal}");
288        if is_equal { Some((leaf, idx)) } else { None }
289    }
290
291    /// Returns true if the map contains the key
292    #[inline(always)]
293    pub fn contains_key<Q>(&self, key: &Q) -> bool
294    where
295        K: Borrow<Q>,
296        Q: Ord + ?Sized,
297    {
298        self.find::<Q>(key).is_some()
299    }
300
301    /// Returns a reference to the value corresponding to the key
302    pub fn get<Q>(&self, key: &Q) -> Option<&V>
303    where
304        K: Borrow<Q>,
305        Q: Ord + ?Sized,
306    {
307        if let Some((leaf, idx)) = self.find::<Q>(key) {
308            let value = unsafe { leaf.value_ptr(idx) };
309            debug_assert!(!value.is_null());
310            Some(unsafe { (*value).assume_init_ref() })
311        } else {
312            None
313        }
314    }
315
316    /// Returns a mutable reference to the value corresponding to the key
317    pub fn get_mut<Q>(&mut self, key: &Q) -> Option<&mut V>
318    where
319        K: Borrow<Q>,
320        Q: Ord + ?Sized,
321    {
322        if let Some((leaf, idx)) = self.find::<Q>(key) {
323            let value = unsafe { leaf.value_ptr(idx) };
324            debug_assert!(!value.is_null());
325            Some(unsafe { (*value).assume_init_mut() })
326        } else {
327            None
328        }
329    }
330
331    /// Insert a key-value pair into the map
332    /// Returns the old value if the key already existed
333    pub fn insert(&mut self, key: K, value: V) -> Option<V> {
334        // use entry api to insert
335        match self.entry(key) {
336            Entry::Occupied(mut entry) => Some(entry.insert(value)),
337            Entry::Vacant(entry) => {
338                entry.insert(value);
339                None
340            }
341        }
342    }
343
344    /// Remove a key from the map, returning the value if it existed
345    pub fn remove<Q>(&mut self, key: &Q) -> Option<V>
346    where
347        K: Borrow<Q>,
348        Q: Ord + ?Sized,
349    {
350        if let Some(root) = &self.root {
351            let cache = self.get_cache();
352            cache.clear();
353            let mut leaf = root.find_leaf_with_cache::<Q>(cache, key);
354            let (idx, is_equal) = leaf.search(key);
355            if is_equal {
356                trace_log!("{leaf:?} remove {idx}");
357                let val = leaf.remove_value_no_borrow(idx);
358                self.len -= 1;
359                // Check for underflow and handle merge
360                let new_count = leaf.key_count();
361                let min_count = LeafNode::<K, V>::cap() >> 1;
362                if new_count < min_count && !root.is_leaf() {
363                    // The cache should already contain the path from the entry lookup
364                    self.handle_leaf_underflow(leaf, true);
365                }
366                Some(val)
367            } else {
368                None
369            }
370        } else {
371            None
372        }
373    }
374
375    /// Perform removal in batch mode and return the last item removed
376    ///
377    /// NOTE: this function speed up by skipping the underflow of LeafNode until the last operation.
378    pub fn remove_range<R>(&mut self, range: R) -> Option<(K, V)>
379    where
380        R: RangeBounds<K>,
381    {
382        self.remove_range_with::<R, _>(range, |_, _| {})
383    }
384
385    /// Perform removal in batch mode and return the last item removed
386    ///
387    /// On each removal, callback function `cb` is invoke with (ref of key, ref of value)
388    ///
389    /// NOTE: this function speed up by skipping the underflow of LeafNode until the last operation.
390    #[inline]
391    pub fn remove_range_with<R, F>(&mut self, range: R, mut cb: F) -> Option<(K, V)>
392    where
393        R: RangeBounds<K>,
394        F: FnMut(&K, &V),
395    {
396        macro_rules! end_contains {
397            ($key: expr) => {{
398                match range.end_bound() {
399                    Bound::Excluded(k) => $key < k,
400                    Bound::Included(k) => $key <= k,
401                    Bound::Unbounded => true,
402                }
403            }};
404        }
405        // Note: do not use find_leaf_with_bound, it's a different behavior
406        let mut ent = match range.start_bound() {
407            Bound::Excluded(k) => match self.entry(k.clone()).move_forward() {
408                Ok(ent) => {
409                    if end_contains!(ent.key()) {
410                        ent
411                    } else {
412                        return None;
413                    }
414                }
415                Err(_) => return None,
416            },
417            Bound::Included(k) => match self.entry(k.clone()) {
418                Entry::Occupied(ent) => ent,
419                Entry::Vacant(ent) => match ent.move_forward() {
420                    Ok(ent) => {
421                        if end_contains!(ent.key()) {
422                            ent
423                        } else {
424                            return None;
425                        }
426                    }
427                    Err(_) => return None,
428                },
429            },
430            Bound::Unbounded => {
431                if let Some(ent) = self.first_entry() {
432                    ent
433                } else {
434                    return None;
435                }
436            }
437        };
438        loop {
439            if let Some((_next_k, _next_v)) = ent.peak_forward() {
440                if end_contains!(_next_k) {
441                    let next_key = _next_k.clone();
442                    let (_k, _v) = ent._remove_entry(false);
443                    cb(&_k, &_v);
444                    if let Entry::Occupied(_ent) = self.entry(next_key) {
445                        ent = _ent;
446                        continue;
447                    } else {
448                        unreachable!();
449                    }
450                }
451            }
452            let (_k, _v) = ent._remove_entry(true);
453            cb(&_k, &_v);
454            return Some((_k, _v));
455        }
456    }
457
458    #[inline(always)]
459    fn get_root_unwrap(&self) -> &Node<K, V> {
460        self.root.as_ref().unwrap()
461    }
462
463    /// update the separate_key in parent after borrowing space from left/right node
464    #[inline(always)]
465    fn update_ancestor_sep_key<const MOVE: bool>(&mut self, sep_key: K) {
466        // if idx == 0, this is the leftmost ptr in the InterNode, we go up until finding a
467        // split key
468        let ret = if MOVE {
469            self.get_cache()
470                .move_to_ancenstor(|_node, idx| -> bool { idx > 0 }, dummy_post_callback)
471        } else {
472            self.get_cache().peak_ancenstor(|_node, idx| -> bool { idx > 0 })
473        };
474        if let Some((parent, parent_idx)) = ret {
475            trace_log!("update_ancestor_sep_key move={MOVE} at {parent:?}:{}", parent_idx - 1);
476            #[cfg(test)]
477            {
478                self.triggers |= TestFlag::UpdateSepKey as u32;
479            }
480            parent.change_key(parent_idx - 1, sep_key);
481        }
482    }
483
484    /// Insert with split handling - called when leaf is full
485    fn insert_with_split(
486        &mut self, key: K, value: V, mut leaf: LeafNode<K, V>, idx: u32,
487    ) -> *mut V {
488        debug_assert!(leaf.is_full());
489        let cap = LeafNode::<K, V>::cap();
490        if idx < cap {
491            // random insert, try borrow space from left and right
492            if let Some(mut left_node) = leaf.get_left_node()
493                && !left_node.is_full()
494            {
495                trace_log!("insert {leaf:?}:{idx} borrow left {left_node:?}");
496                let val_p = if idx == 0 {
497                    // leaf is not change, but since the insert pos is leftmost of this node, mean parent
498                    // separate_key <= key, need to update its separate_key
499                    left_node.insert_no_split_with_idx(left_node.key_count(), key, value)
500                } else {
501                    leaf.insert_borrow_left(&mut left_node, idx, key, value)
502                };
503                #[cfg(test)]
504                {
505                    self.triggers |= TestFlag::LeafMoveLeft as u32;
506                }
507                self.update_ancestor_sep_key::<true>(leaf.clone_first_key());
508                return val_p;
509            }
510        } else {
511            // insert into empty new node, left is probably full, right is probably none
512        }
513        if let Some(mut right_node) = leaf.get_right_node()
514            && !right_node.is_full()
515        {
516            trace_log!("insert {leaf:?}:{idx} borrow right {right_node:?}");
517            let val_p = if idx == cap {
518                // leaf is not change, in this condition, right_node is the leftmost child
519                // of its parent, key < right_node.get_keys()[0]
520                right_node.insert_no_split_with_idx(0, key, value)
521            } else {
522                leaf.borrow_right(&mut right_node);
523                leaf.insert_no_split_with_idx(idx, key, value)
524            };
525            #[cfg(test)]
526            {
527                self.triggers |= TestFlag::LeafMoveRight as u32;
528            }
529            self.get_cache().move_right();
530            self.update_ancestor_sep_key::<true>(right_node.clone_first_key());
531            return val_p;
532        }
533        #[cfg(test)]
534        {
535            self.triggers |= TestFlag::LeafSplit as u32;
536        }
537        let (new_leaf, ptr_v) = leaf.insert_with_split(idx, key, value);
538        self.leaf_count += 1;
539        let split_key = unsafe { (*new_leaf.key_ptr(0)).assume_init_ref().clone() };
540        self.propagate_split(split_key, leaf.get_ptr(), new_leaf.get_ptr());
541        ptr_v
542    }
543
544    /// Propagate node split up the tree using iteration (non-recursive)
545    /// First tries to borrow space from left/right sibling before splitting
546    ///
547    /// left_ptr: existing child
548    /// right_ptr: new_child split from left_ptr
549    /// promote_key: sep_key to split left_ptr & right_ptr
550    #[inline(always)]
551    fn propagate_split(
552        &mut self, mut promote_key: K, mut left_ptr: *mut NodeHeader,
553        mut right_ptr: *mut NodeHeader,
554    ) {
555        let mut height = 0;
556        // If we have parent nodes in cache, process them iteratively
557        while let Some((mut parent, idx)) = self.get_cache().pop() {
558            if !parent.is_full() {
559                trace_log!("propagate_split {parent:?}:{idx} insert {right_ptr:p}");
560                // should insert next to left_ptr
561                parent.insert_no_split_with_idx(idx, promote_key, right_ptr);
562                return;
563            } else {
564                // Parent is full, try to borrow space from sibling through grand_parent
565                if let Some((mut grand, grand_idx)) = self.get_cache().peak_next() {
566                    // Try to borrow from left sibling of parent
567                    if grand_idx > 0 {
568                        let mut left_parent = grand.get_child_as_inter(grand_idx - 1);
569                        if !left_parent.is_full() {
570                            #[cfg(test)]
571                            {
572                                self.triggers |= TestFlag::InterMoveLeft as u32;
573                            }
574                            if idx == 0 {
575                                trace_log!(
576                                    "propagate_split rotate_left {grand:?}:{} first ->{left_parent:?} left {left_ptr:p} insert {idx} {right_ptr:p}",
577                                    grand_idx - 1
578                                );
579                                // special case: split from first child of parent
580                                let demote_key = grand.change_key(grand_idx - 1, promote_key);
581                                debug_assert_eq!(parent.get_child_ptr(0), left_ptr);
582                                unsafe { (*parent.child_ptr(0)) = right_ptr };
583                                left_parent.append(demote_key, left_ptr);
584                                #[cfg(test)]
585                                {
586                                    self.triggers |= TestFlag::InterMoveLeftFirst as u32;
587                                }
588                            } else {
589                                trace_log!(
590                                    "propagate_split insert_rotate_left {grand:?}:{grand_idx} -> {left_parent:?} insert {idx} {right_ptr:p}"
591                                );
592                                parent.insert_rotate_left(
593                                    &mut grand,
594                                    grand_idx,
595                                    &mut left_parent,
596                                    idx,
597                                    promote_key,
598                                    right_ptr,
599                                );
600                            }
601                            return;
602                        }
603                    }
604                    // Try to borrow from right sibling of parent
605                    if grand_idx < grand.key_count() {
606                        let mut right_parent = grand.get_child_as_inter(grand_idx + 1);
607                        if !right_parent.is_full() {
608                            #[cfg(test)]
609                            {
610                                self.triggers |= TestFlag::InterMoveRight as u32;
611                            }
612                            if idx == parent.key_count() {
613                                trace_log!(
614                                    "propagate_split rotate_right last {grand:?}:{grand_idx} -> {right_parent:?}:0 insert right {right_parent:?}:0 {right_ptr:p}"
615                                );
616                                // split from last child of parent
617                                let demote_key = grand.change_key(grand_idx, promote_key);
618                                right_parent.insert_at_front(right_ptr, demote_key);
619                                #[cfg(test)]
620                                {
621                                    self.triggers |= TestFlag::InterMoveRightLast as u32;
622                                }
623                            } else {
624                                trace_log!(
625                                    "propagate_split rotate_right {grand:?}:{grand_idx} -> {right_parent:?}:0 insert {parent:?}:{idx} {right_ptr:p}"
626                                );
627                                parent.rotate_right(&mut grand, grand_idx, &mut right_parent);
628                                parent.insert_no_split_with_idx(idx, promote_key, right_ptr);
629                            }
630                            return;
631                        }
632                    }
633                }
634                height += 1;
635
636                // Cannot borrow from siblings, need to split internal node
637                let (right, _promote_key) = parent.insert_split(promote_key, right_ptr);
638                promote_key = _promote_key;
639                right_ptr = right.get_ptr();
640                left_ptr = parent.get_ptr();
641                #[cfg(test)]
642                {
643                    self.triggers |= TestFlag::InterSplit as u32;
644                }
645                // Continue to next parent in cache (loop will pop next parent)
646            }
647        }
648        // No more parents in cache, create new root
649        let new_root = InterNode::<K, V>::new_root(height + 1, promote_key, left_ptr, right_ptr);
650        let _old_root = self.root.replace(Node::Inter(new_root)).expect("should have old root");
651        #[cfg(debug_assertions)]
652        {
653            match _old_root {
654                Node::Inter(node) => {
655                    debug_assert_eq!(height, node.height(), "old inter root at {height}");
656                    debug_assert_eq!(node.get_ptr(), left_ptr);
657                }
658                Node::Leaf(node) => {
659                    debug_assert_eq!(height, 0);
660                    debug_assert_eq!(node.get_ptr(), left_ptr);
661                }
662            }
663        }
664    }
665
666    /// Handle leaf node underflow by merging with sibling
667    /// Uses PathCache to accelerate parent lookup
668    /// Following the merge strategy from Designer Notes:
669    /// - Try merge with left sibling (if left + current <= cap)
670    /// - Try merge with right sibling (if current + right <= cap)
671    /// - Try 3-node merge (if left + current + right <= 2 * cap)
672    fn handle_leaf_underflow(&mut self, mut leaf: LeafNode<K, V>, merge: bool) {
673        debug_assert!(!self.get_root_unwrap().is_leaf());
674        let cur_count = leaf.key_count();
675        let cap = LeafNode::<K, V>::cap();
676        debug_assert!(cur_count <= cap >> 1);
677        let mut can_unlink: bool = false;
678        let (mut left_avail, mut right_avail) = (0, 0);
679        let mut merge_right = false;
680        if cur_count == 0 {
681            trace_log!("handle_leaf_underflow {leaf:?} unlink");
682            // if the right and left are full, or they not exist, can come to this
683            can_unlink = true;
684        }
685        if merge {
686            if !can_unlink && let Some(mut left_node) = leaf.get_left_node() {
687                let left_count = left_node.key_count();
688                if left_count + cur_count <= cap {
689                    trace_log!(
690                        "handle_leaf_underflow {leaf:?} merge left {left_node:?} {cur_count}"
691                    );
692                    leaf.copy_left(&mut left_node, cur_count);
693                    can_unlink = true;
694                    #[cfg(test)]
695                    {
696                        self.triggers |= TestFlag::LeafMergeLeft as u32;
697                    }
698                } else {
699                    left_avail = cap - left_count;
700                }
701            }
702            if !can_unlink && let Some(mut right_node) = leaf.get_right_node() {
703                let right_count = right_node.key_count();
704                if right_count + cur_count <= cap {
705                    trace_log!(
706                        "handle_leaf_underflow {leaf:?} merge right {right_node:?} {cur_count}"
707                    );
708                    leaf.copy_right::<false>(&mut right_node, 0, cur_count);
709                    can_unlink = true;
710                    merge_right = true;
711                    #[cfg(test)]
712                    {
713                        self.triggers |= TestFlag::LeafMergeRight as u32;
714                    }
715                } else {
716                    right_avail = cap - right_count;
717                }
718            }
719            // if we require left_avail + right_avail > cur_count, not possible to construct a 3-2
720            // merge, only either triggering merge left or merge right.
721            if !can_unlink
722                && left_avail > 0
723                && right_avail > 0
724                && left_avail + right_avail == cur_count
725            {
726                let mut left_node = leaf.get_left_node().unwrap();
727                let mut right_node = leaf.get_right_node().unwrap();
728                debug_assert!(left_avail < cur_count);
729                trace_log!("handle_leaf_underflow {leaf:?} merge left {left_node:?} {left_avail}");
730                leaf.copy_left(&mut left_node, left_avail);
731                trace_log!(
732                    "handle_leaf_underflow {leaf:?} merge right {right_node:?} {}",
733                    cur_count - left_avail
734                );
735                leaf.copy_right::<false>(&mut right_node, left_avail, cur_count - left_avail);
736                merge_right = true;
737                can_unlink = true;
738                #[cfg(test)]
739                {
740                    self.triggers |=
741                        TestFlag::LeafMergeLeft as u32 | TestFlag::LeafMergeRight as u32;
742                }
743            }
744        }
745        if !can_unlink {
746            return;
747        }
748        self.leaf_count -= 1;
749        let right_sep = if merge_right {
750            let right_node = leaf.get_right_node().unwrap();
751            Some(right_node.clone_first_key())
752        } else {
753            None
754        };
755        let no_right = leaf.unlink().is_null();
756        leaf.dealloc::<false>();
757        let (mut parent, mut idx) = self.get_cache().pop().unwrap();
758        trace_log!("handle_leaf_underflow pop parent {parent:?}:{idx}");
759        if parent.key_count() == 0 {
760            if let Some((grand, grand_idx)) = self.remove_only_child(parent) {
761                trace_log!("handle_leaf_underflow remove_only_child until {grand:?}:{grand_idx}");
762                parent = grand;
763                idx = grand_idx;
764            } else {
765                trace_log!("handle_leaf_underflow remove_only_child all");
766                return;
767            }
768        }
769        self.remove_child_from_inter(&mut parent, idx, right_sep, no_right);
770        if parent.key_count() <= 1 {
771            self.handle_inter_underflow(parent);
772        }
773    }
774
775    /// To simplify the logic, we perform delete first.
776    /// return the Some(node) when need to rebalance
777    #[inline]
778    fn remove_child_from_inter(
779        &mut self, node: &mut InterNode<K, V>, delete_idx: u32, right_sep: Option<K>,
780        _no_right: bool,
781    ) {
782        self.get_cache().assert_center();
783        debug_assert!(node.key_count() > 0, "{:?} {}", node, node.key_count());
784        if delete_idx == node.key_count() {
785            trace_log!("remove_child_from_inter {node:?}:{delete_idx} last");
786            #[cfg(test)]
787            {
788                self.triggers |= TestFlag::RemoveChildLast as u32;
789            }
790            // delete the last child of this node
791            node.remove_last_child();
792            if let Some(key) = right_sep
793                && let Some((grand_parent, grand_idx)) =
794                    self.get_cache().peak_ancenstor(|_node: &InterNode<K, V>, idx: u32| -> bool {
795                        _node.key_count() > idx
796                    })
797            {
798                #[cfg(test)]
799                {
800                    self.triggers |= TestFlag::UpdateSepKey as u32;
801                }
802                trace_log!("remove_child_from_inter change_key {grand_parent:?}:{grand_idx}");
803                // key idx = child idx - 1 , and + 1 for right node
804                grand_parent.change_key(grand_idx, key);
805            }
806        } else if delete_idx > 0 {
807            trace_log!("remove_child_from_inter {node:?}:{delete_idx} mid");
808            node.remove_mid_child(delete_idx);
809            #[cfg(test)]
810            {
811                self.triggers |= TestFlag::RemoveChildMid as u32;
812            }
813            // sep key of right node shift left
814            if let Some(key) = right_sep {
815                trace_log!("remove_child_from_inter change_key {node:?}:{}", delete_idx - 1);
816                node.change_key(delete_idx - 1, key);
817                #[cfg(test)]
818                {
819                    self.triggers |= TestFlag::UpdateSepKey as u32;
820                }
821            }
822        } else {
823            trace_log!("remove_child_from_inter {node:?}:{delete_idx} first");
824            // delete_idx is the first but not the last
825            let mut sep_key = node.remove_first_child();
826            #[cfg(test)]
827            {
828                self.triggers |= TestFlag::RemoveChildFirst as u32;
829            }
830            if let Some(key) = right_sep {
831                sep_key = key;
832            }
833            self.update_ancestor_sep_key::<false>(sep_key);
834        }
835    }
836
837    #[inline]
838    fn handle_inter_underflow(&mut self, mut node: InterNode<K, V>) {
839        let cap = InterNode::<K, V>::cap();
840        while node.key_count() <= InterNode::<K, V>::UNDERFLOW_CAP {
841            if node.key_count() == 0 {
842                let root_height = self.get_root_unwrap().height();
843                if node.height() == root_height
844                    || self
845                        .get_cache()
846                        .peak_ancenstor(|_node: &InterNode<K, V>, _idx: u32| -> bool {
847                            _node.key_count() > 0
848                        })
849                        .is_none()
850                {
851                    let root = unsafe { Node::from_header(*node.child_ptr(0)) };
852                    trace_log!("handle_inter_underflow downgrade root {root:?}");
853                    let _old_root = self.root.replace(root);
854                    debug_assert!(_old_root.is_some());
855
856                    while let Some((parent, _)) = self.get_cache().pop() {
857                        parent.dealloc::<false>();
858                    }
859                    node.dealloc::<false>(); // they all have no key
860                }
861                return;
862            } else {
863                if let Some((mut grand, grand_idx)) = self.get_cache().pop() {
864                    if grand_idx > 0 {
865                        let mut left = grand.get_child_as_inter(grand_idx - 1);
866                        // the sep key should pull down,  key+1 + key + 1 > cap + 1
867                        if left.key_count() + node.key_count() < cap {
868                            #[cfg(test)]
869                            {
870                                self.triggers |= TestFlag::InterMergeLeft as u32;
871                            }
872                            trace_log!(
873                                "handle_inter_underflow {node:?} merge left {left:?} parent {grand:?}:{grand_idx}"
874                            );
875                            left.merge(node, &mut grand, grand_idx);
876                            node = grand;
877                            continue;
878                        }
879                    }
880                    if grand_idx < grand.key_count() {
881                        let right = grand.get_child_as_inter(grand_idx + 1);
882                        // the sep key should pull down,  key+1 + key + 1 > cap + 1
883                        if right.key_count() + node.key_count() < cap {
884                            #[cfg(test)]
885                            {
886                                self.triggers |= TestFlag::InterMergeRight as u32;
887                            }
888                            trace_log!(
889                                "handle_inter_underflow {node:?} cap {cap} merge right {right:?} parent {grand:?}:{}",
890                                grand_idx + 1
891                            );
892                            node.merge(right, &mut grand, grand_idx + 1);
893                            node = grand;
894                            continue;
895                        }
896                    }
897                }
898                return;
899            }
900        }
901    }
902
903    #[inline]
904    fn remove_only_child(&mut self, node: InterNode<K, V>) -> Option<(InterNode<K, V>, u32)> {
905        debug_assert_eq!(node.key_count(), 0);
906        #[cfg(test)]
907        {
908            self.triggers |= TestFlag::RemoveOnlyChild as u32;
909        }
910        if let Some((parent, idx)) = self.get_cache().move_to_ancenstor(
911            |node: &InterNode<K, V>, _idx: u32| -> bool { node.key_count() != 0 },
912            InterNode::<K, V>::dealloc::<false>,
913        ) {
914            node.dealloc::<true>();
915            Some((parent, idx))
916        } else {
917            node.dealloc::<true>();
918            // we are empty, my ancestor are all empty and delete by move_to_ancenstor
919            self.root = None;
920            None
921        }
922    }
923
924    /// Dump the entire tree structure for debugging
925    #[cfg(all(test, feature = "std"))]
926    pub fn dump(&self)
927    where
928        K: Debug,
929        V: Debug,
930    {
931        print_log!("=== BTreeMap Dump ===");
932        print_log!("Length: {}", self.len());
933        if let Some(root) = &self.root {
934            self.dump_node(root, 0);
935        } else {
936            print_log!("(empty)");
937        }
938        print_log!("=====================");
939    }
940
941    #[cfg(all(test, feature = "std"))]
942    fn dump_node(&self, node: &Node<K, V>, depth: usize)
943    where
944        K: Debug,
945        V: Debug,
946    {
947        match node {
948            Node::Leaf(leaf) => {
949                print!("{:indent$}", "", indent = depth * 2);
950                print_log!("{}", leaf);
951            }
952            Node::Inter(inter) => {
953                print!("{:indent$}", "", indent = depth * 2);
954                print_log!("{}", inter);
955                // Dump children
956                let count = inter.key_count() as u32;
957                for i in 0..=count {
958                    unsafe {
959                        let child_ptr = *inter.child_ptr(i);
960                        if !child_ptr.is_null() {
961                            let child_node = if (*child_ptr).is_leaf() {
962                                Node::Leaf(LeafNode::<K, V>::from_header(child_ptr))
963                            } else {
964                                Node::Inter(InterNode::<K, V>::from_header(child_ptr))
965                            };
966                            self.dump_node(&child_node, depth + 1);
967                        }
968                    }
969                }
970            }
971        }
972    }
973
974    /// Validate the entire tree structure
975    /// Uses the same traversal logic as Drop to avoid recursion
976    pub fn validate(&self)
977    where
978        K: Debug,
979        V: Debug,
980    {
981        if self.root.is_none() {
982            assert_eq!(self.len, 0, "Empty tree should have len 0");
983            return;
984        }
985
986        let mut total_keys = 0usize;
987        let mut prev_leaf_max: Option<K> = None;
988
989        match &self.root {
990            Some(Node::Leaf(leaf)) => {
991                total_keys += leaf.validate(None, None);
992            }
993            Some(Node::Inter(inter)) => {
994                // Do not use btree internal PathCache (might distrupt test scenario)
995                let mut cache = PathCache::new();
996                let mut cur = inter.clone();
997                loop {
998                    cache.push(cur.clone(), 0);
999                    cur.validate();
1000                    match cur.get_child(0) {
1001                        Node::Leaf(leaf) => {
1002                            // Validate first leaf with no min/max bounds from parent
1003                            let min_key: Option<K> = None;
1004                            let max_key = if inter.key_count() > 0 {
1005                                unsafe { Some((*inter.key_ptr(0)).assume_init_ref().clone()) }
1006                            } else {
1007                                None
1008                            };
1009                            total_keys += leaf.validate(min_key.as_ref(), max_key.as_ref());
1010                            if let Some(ref prev_max) = prev_leaf_max {
1011                                let first_key = unsafe { (*leaf.key_ptr(0)).assume_init_ref() };
1012                                assert!(
1013                                    prev_max < first_key,
1014                                    "{:?} Leaf keys not in order: prev max {:?} >= current min {:?}",
1015                                    leaf,
1016                                    prev_max,
1017                                    first_key
1018                                );
1019                            }
1020                            prev_leaf_max = unsafe {
1021                                Some(
1022                                    (*leaf.key_ptr(leaf.key_count() - 1)).assume_init_ref().clone(),
1023                                )
1024                            };
1025                            break;
1026                        }
1027                        Node::Inter(child_inter) => {
1028                            cur = child_inter;
1029                        }
1030                    }
1031                }
1032
1033                // Continue traversal like Drop does
1034                while let Some((parent, idx)) =
1035                    cache.move_right_and_pop_l1(dummy_post_callback::<K, V>)
1036                {
1037                    cache.push(parent.clone(), idx);
1038                    if let Node::Leaf(leaf) = parent.get_child(idx) {
1039                        // Calculate bounds for this leaf
1040                        let min_key = if idx > 0 {
1041                            unsafe { Some((*parent.key_ptr(idx - 1)).assume_init_ref().clone()) }
1042                        } else {
1043                            None
1044                        };
1045                        let max_key = if idx < parent.key_count() {
1046                            unsafe { Some((*parent.key_ptr(idx)).assume_init_ref().clone()) }
1047                        } else {
1048                            None
1049                        };
1050                        total_keys += leaf.validate(min_key.as_ref(), max_key.as_ref());
1051
1052                        // Check ordering with previous leaf
1053                        if let Some(ref prev_max) = prev_leaf_max {
1054                            let first_key = unsafe { (*leaf.key_ptr(0)).assume_init_ref() };
1055                            assert!(
1056                                prev_max < first_key,
1057                                "{:?} Leaf keys not in order: prev max {:?} >= current min {:?}",
1058                                leaf,
1059                                prev_max,
1060                                first_key
1061                            );
1062                        }
1063                        prev_leaf_max = unsafe {
1064                            Some((*leaf.key_ptr(leaf.key_count() - 1)).assume_init_ref().clone())
1065                        };
1066                    } else {
1067                        panic!("{parent:?} child {:?} is not leaf", parent.get_child(idx));
1068                    }
1069                }
1070            }
1071            None => unreachable!(),
1072        }
1073        assert_eq!(
1074            total_keys, self.len,
1075            "Total keys in tree ({}) doesn't match len ({})",
1076            total_keys, self.len
1077        );
1078    }
1079
1080    /// Returns the first key-value pair in the map
1081    /// Returns `None` if the map is empty
1082    #[inline]
1083    pub fn first_key_value(&self) -> Option<(&K, &V)> {
1084        let leaf = match &self.root {
1085            None => return None,
1086            Some(root) => root.find_first_leaf(None),
1087        };
1088        if leaf.key_count() == 0 {
1089            return None;
1090        }
1091        unsafe {
1092            let key = (*leaf.key_ptr(0)).assume_init_ref();
1093            let value = (*leaf.value_ptr(0)).assume_init_ref();
1094            Some((key, value))
1095        }
1096    }
1097
1098    /// Returns the last key-value pair in the map
1099    /// Returns `None` if the map is empty
1100    #[inline]
1101    pub fn last_key_value(&self) -> Option<(&K, &V)> {
1102        let leaf = match &self.root {
1103            None => return None,
1104            Some(root) => root.find_last_leaf(None),
1105        };
1106        let count = leaf.key_count();
1107        if count == 0 {
1108            return None;
1109        }
1110        unsafe {
1111            let last_idx = count - 1;
1112            let key = (*leaf.key_ptr(last_idx)).assume_init_ref();
1113            let value = (*leaf.value_ptr(last_idx)).assume_init_ref();
1114            Some((key, value))
1115        }
1116    }
1117
1118    /// Returns an entry to the first key in the map
1119    /// Returns `None` if the map is empty
1120    #[inline]
1121    pub fn first_entry(&mut self) -> Option<OccupiedEntry<'_, K, V>> {
1122        let cache = self.get_cache();
1123        cache.clear();
1124        let leaf = match &self.root {
1125            None => return None,
1126            Some(root) => {
1127                // Populate cache with the path to first leaf
1128                root.find_first_leaf(Some(cache))
1129            }
1130        };
1131        if leaf.key_count() == 0 {
1132            return None;
1133        }
1134        Some(OccupiedEntry { tree: self, idx: 0, leaf })
1135    }
1136
1137    /// Returns an entry to the last key in the map
1138    /// Returns `None` if the map is empty
1139    #[inline]
1140    pub fn last_entry(&mut self) -> Option<OccupiedEntry<'_, K, V>> {
1141        let cache = self.get_cache();
1142        cache.clear();
1143        let leaf = match &self.root {
1144            None => return None,
1145            Some(root) => {
1146                // Populate cache with the path to last leaf
1147                root.find_last_leaf(Some(cache))
1148            }
1149        };
1150        let count = leaf.key_count();
1151        if count == 0 {
1152            return None;
1153        }
1154        Some(OccupiedEntry { tree: self, idx: count - 1, leaf })
1155    }
1156
1157    /// Removes and returns the first key-value pair in the map
1158    /// Returns `None` if the map is empty
1159    #[inline]
1160    pub fn pop_first(&mut self) -> Option<(K, V)> {
1161        match self.first_entry() {
1162            Some(entry) => Some(entry.remove_entry()),
1163            _ => None,
1164        }
1165    }
1166
1167    /// Removes and returns the last key-value pair in the map
1168    /// Returns `None` if the map is empty
1169    #[inline]
1170    pub fn pop_last(&mut self) -> Option<(K, V)> {
1171        match self.last_entry() {
1172            Some(entry) => Some(entry.remove_entry()),
1173            _ => None,
1174        }
1175    }
1176
1177    /// Returns an iterator over the map's entries
1178    #[inline]
1179    pub fn iter(&self) -> Iter<'_, K, V> {
1180        Iter::new(self.find_first_and_last_leaf(), self.len)
1181    }
1182
1183    /// Returns a mutable iterator over the map's entries
1184    #[inline]
1185    pub fn iter_mut(&mut self) -> IterMut<'_, K, V> {
1186        IterMut::new(self.find_first_and_last_leaf(), self.len)
1187    }
1188
1189    /// Returns an iterator over the map's keys
1190    #[inline]
1191    pub fn keys(&self) -> Keys<'_, K, V> {
1192        Keys::new(self.iter())
1193    }
1194
1195    /// Returns an iterator over the map's values
1196    #[inline]
1197    pub fn values(&self) -> Values<'_, K, V> {
1198        Values::new(self.iter())
1199    }
1200
1201    /// Returns a mutable iterator over the map's values
1202    #[inline]
1203    pub fn values_mut(&mut self) -> ValuesMut<'_, K, V> {
1204        ValuesMut::new(self.iter_mut())
1205    }
1206
1207    #[inline]
1208    fn find_first_and_last_leaf(&self) -> Option<(LeafNode<K, V>, LeafNode<K, V>)> {
1209        let root = self.root.as_ref()?;
1210        Some((root.find_first_leaf(None), root.find_last_leaf(None)))
1211    }
1212
1213    /// Internal helper to find range bounds
1214    /// Returns (front_leaf, front_idx, back_leaf, back_idx) where both leaves are Some or both are None
1215    #[inline]
1216    fn find_range_bounds<R>(&self, range: R) -> Option<RangeBase<'_, K, V>>
1217    where
1218        R: RangeBounds<K>,
1219    {
1220        let root = self.root.as_ref()?;
1221        let (front_leaf, front_idx) = root.find_leaf_with_bound(range.start_bound(), true);
1222        let (back_leaf, back_idx) = root.find_leaf_with_bound(range.end_bound(), false);
1223        Some(RangeBase::new(front_leaf, front_idx, back_leaf, back_idx))
1224    }
1225
1226    /// Returns an iterator over a sub-range of entries in the map
1227    #[inline]
1228    pub fn range<R>(&self, range: R) -> Range<'_, K, V>
1229    where
1230        R: RangeBounds<K>,
1231    {
1232        Range::new(self.find_range_bounds(range))
1233    }
1234
1235    /// Returns a mutable iterator over a sub-range of entries in the map
1236    #[inline]
1237    pub fn range_mut<R>(&mut self, range: R) -> RangeMut<'_, K, V>
1238    where
1239        R: RangeBounds<K>,
1240    {
1241        RangeMut::new(self.find_range_bounds(range))
1242    }
1243}
1244
1245impl<K: Ord + Clone + Sized, V: Sized> Default for BTreeMap<K, V> {
1246    fn default() -> Self {
1247        Self::new()
1248    }
1249}
1250
1251impl<K: Ord + Clone + Sized, V: Sized> Drop for BTreeMap<K, V> {
1252    fn drop(&mut self) {
1253        let mut cache = self.get_cache().take();
1254        let mut cur = match self.root.take() {
1255            None => return,
1256            Some(root) => root.find_first_leaf(Some(&mut cache)),
1257        };
1258        cur.dealloc::<true>();
1259        // To navigate to next leaf,
1260        // return None when reach the end
1261        while let Some((parent, idx)) = cache.move_right_and_pop_l1(InterNode::dealloc::<true>) {
1262            cache.push(parent.clone(), idx);
1263            cur = parent.get_child_as_leaf(idx);
1264            cur.dealloc::<true>();
1265        }
1266    }
1267}
1268
1269impl<K: Ord + Clone + Sized, V: Sized> IntoIterator for BTreeMap<K, V> {
1270    type Item = (K, V);
1271    type IntoIter = IntoIter<K, V>;
1272
1273    #[inline]
1274    fn into_iter(self) -> Self::IntoIter {
1275        IntoIter::new(self, true)
1276    }
1277}
1278impl<'a, K: Ord + Clone + Sized, V: Sized> IntoIterator for &'a BTreeMap<K, V> {
1279    type Item = (&'a K, &'a V);
1280    type IntoIter = Iter<'a, K, V>;
1281
1282    #[inline]
1283    fn into_iter(self) -> Self::IntoIter {
1284        self.iter()
1285    }
1286}
1287
1288impl<'a, K: Ord + Clone + Sized, V: Sized> IntoIterator for &'a mut BTreeMap<K, V> {
1289    type Item = (&'a K, &'a mut V);
1290    type IntoIter = IterMut<'a, K, V>;
1291
1292    #[inline]
1293    fn into_iter(self) -> Self::IntoIter {
1294        self.iter_mut()
1295    }
1296}