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        // Note: do not use find_leaf_with_bound, it's a different behavior
397        let mut ent = match range.start_bound() {
398            Bound::Excluded(k) => match self.entry(k.clone()).move_forward() {
399                Ok(ent) => {
400                    if !range.contains(ent.key()) {
401                        return None;
402                    }
403                    ent
404                }
405                Err(_) => return None,
406            },
407            Bound::Included(k) => match self.entry(k.clone()) {
408                Entry::Occupied(ent) => ent,
409                Entry::Vacant(ent) => match ent.move_forward() {
410                    Ok(ent) => {
411                        if !range.contains(ent.key()) {
412                            return None;
413                        }
414                        ent
415                    }
416                    Err(_) => return None,
417                },
418            },
419            Bound::Unbounded => {
420                if let Some(ent) = self.first_entry() {
421                    ent
422                } else {
423                    return None;
424                }
425            }
426        };
427        loop {
428            if let Some((_next_k, _next_v)) = ent.peak_forward() {
429                if range.contains(_next_k) {
430                    let next_key = _next_k.clone();
431                    let (_k, _v) = ent._remove_entry(false);
432                    cb(&_k, &_v);
433                    if let Entry::Occupied(_ent) = self.entry(next_key) {
434                        ent = _ent;
435                        continue;
436                    } else {
437                        unreachable!();
438                    }
439                }
440            }
441            let (_k, _v) = ent._remove_entry(true);
442            cb(&_k, &_v);
443            return Some((_k, _v));
444        }
445    }
446
447    #[inline(always)]
448    fn get_root_unwrap(&self) -> &Node<K, V> {
449        self.root.as_ref().unwrap()
450    }
451
452    /// update the separate_key in parent after borrowing space from left/right node
453    #[inline(always)]
454    fn update_ancestor_sep_key<const MOVE: bool>(&mut self, sep_key: K) {
455        // if idx == 0, this is the leftmost ptr in the InterNode, we go up until finding a
456        // split key
457        let ret = if MOVE {
458            self.get_cache()
459                .move_to_ancenstor(|_node, idx| -> bool { idx > 0 }, dummy_post_callback)
460        } else {
461            self.get_cache().peak_ancenstor(|_node, idx| -> bool { idx > 0 })
462        };
463        if let Some((parent, parent_idx)) = ret {
464            trace_log!("update_ancestor_sep_key move={MOVE} at {parent:?}:{}", parent_idx - 1);
465            #[cfg(test)]
466            {
467                self.triggers |= TestFlag::UpdateSepKey as u32;
468            }
469            parent.change_key(parent_idx - 1, sep_key);
470        }
471    }
472
473    /// Insert with split handling - called when leaf is full
474    fn insert_with_split(
475        &mut self, key: K, value: V, mut leaf: LeafNode<K, V>, idx: u32,
476    ) -> *mut V {
477        debug_assert!(leaf.is_full());
478        let cap = LeafNode::<K, V>::cap();
479        if idx < cap {
480            // random insert, try borrow space from left and right
481            if let Some(mut left_node) = leaf.get_left_node()
482                && !left_node.is_full()
483            {
484                trace_log!("insert {leaf:?}:{idx} borrow left {left_node:?}");
485                let val_p = if idx == 0 {
486                    // leaf is not change, but since the insert pos is leftmost of this node, mean parent
487                    // separate_key <= key, need to update its separate_key
488                    left_node.insert_no_split_with_idx(left_node.key_count(), key, value)
489                } else {
490                    leaf.insert_borrow_left(&mut left_node, idx, key, value)
491                };
492                #[cfg(test)]
493                {
494                    self.triggers |= TestFlag::LeafMoveLeft as u32;
495                }
496                self.update_ancestor_sep_key::<true>(leaf.clone_first_key());
497                return val_p;
498            }
499        } else {
500            // insert into empty new node, left is probably full, right is probably none
501        }
502        if let Some(mut right_node) = leaf.get_right_node()
503            && !right_node.is_full()
504        {
505            trace_log!("insert {leaf:?}:{idx} borrow right {right_node:?}");
506            let val_p = if idx == cap {
507                // leaf is not change, in this condition, right_node is the leftmost child
508                // of its parent, key < right_node.get_keys()[0]
509                right_node.insert_no_split_with_idx(0, key, value)
510            } else {
511                leaf.borrow_right(&mut right_node);
512                leaf.insert_no_split_with_idx(idx, key, value)
513            };
514            #[cfg(test)]
515            {
516                self.triggers |= TestFlag::LeafMoveRight as u32;
517            }
518            self.get_cache().move_right();
519            self.update_ancestor_sep_key::<true>(right_node.clone_first_key());
520            return val_p;
521        }
522        #[cfg(test)]
523        {
524            self.triggers |= TestFlag::LeafSplit as u32;
525        }
526        let (new_leaf, ptr_v) = leaf.insert_with_split(idx, key, value);
527        self.leaf_count += 1;
528        let split_key = unsafe { (*new_leaf.key_ptr(0)).assume_init_ref().clone() };
529        self.propagate_split(split_key, leaf.get_ptr(), new_leaf.get_ptr());
530        ptr_v
531    }
532
533    /// Propagate node split up the tree using iteration (non-recursive)
534    /// First tries to borrow space from left/right sibling before splitting
535    ///
536    /// left_ptr: existing child
537    /// right_ptr: new_child split from left_ptr
538    /// promote_key: sep_key to split left_ptr & right_ptr
539    #[inline(always)]
540    fn propagate_split(
541        &mut self, mut promote_key: K, mut left_ptr: *mut NodeHeader,
542        mut right_ptr: *mut NodeHeader,
543    ) {
544        let mut height = 0;
545        // If we have parent nodes in cache, process them iteratively
546        while let Some((mut parent, idx)) = self.get_cache().pop() {
547            if !parent.is_full() {
548                trace_log!("propagate_split {parent:?}:{idx} insert {right_ptr:p}");
549                // should insert next to left_ptr
550                parent.insert_no_split_with_idx(idx, promote_key, right_ptr);
551                return;
552            } else {
553                // Parent is full, try to borrow space from sibling through grand_parent
554                if let Some((mut grand, grand_idx)) = self.get_cache().peak_next() {
555                    // Try to borrow from left sibling of parent
556                    if grand_idx > 0 {
557                        let mut left_parent = grand.get_child_as_inter(grand_idx - 1);
558                        if !left_parent.is_full() {
559                            #[cfg(test)]
560                            {
561                                self.triggers |= TestFlag::InterMoveLeft as u32;
562                            }
563                            if idx == 0 {
564                                trace_log!(
565                                    "propagate_split rotate_left {grand:?}:{} first ->{left_parent:?} left {left_ptr:p} insert {idx} {right_ptr:p}",
566                                    grand_idx - 1
567                                );
568                                // special case: split from first child of parent
569                                let demote_key = grand.change_key(grand_idx - 1, promote_key);
570                                debug_assert_eq!(parent.get_child_ptr(0), left_ptr);
571                                unsafe { (*parent.child_ptr(0)) = right_ptr };
572                                left_parent.append(demote_key, left_ptr);
573                                #[cfg(test)]
574                                {
575                                    self.triggers |= TestFlag::InterMoveLeftFirst as u32;
576                                }
577                            } else {
578                                trace_log!(
579                                    "propagate_split insert_rotate_left {grand:?}:{grand_idx} -> {left_parent:?} insert {idx} {right_ptr:p}"
580                                );
581                                parent.insert_rotate_left(
582                                    &mut grand,
583                                    grand_idx,
584                                    &mut left_parent,
585                                    idx,
586                                    promote_key,
587                                    right_ptr,
588                                );
589                            }
590                            return;
591                        }
592                    }
593                    // Try to borrow from right sibling of parent
594                    if grand_idx < grand.key_count() {
595                        let mut right_parent = grand.get_child_as_inter(grand_idx + 1);
596                        if !right_parent.is_full() {
597                            #[cfg(test)]
598                            {
599                                self.triggers |= TestFlag::InterMoveRight as u32;
600                            }
601                            if idx == parent.key_count() {
602                                trace_log!(
603                                    "propagate_split rotate_right last {grand:?}:{grand_idx} -> {right_parent:?}:0 insert right {right_parent:?}:0 {right_ptr:p}"
604                                );
605                                // split from last child of parent
606                                let demote_key = grand.change_key(grand_idx, promote_key);
607                                right_parent.insert_at_front(right_ptr, demote_key);
608                                #[cfg(test)]
609                                {
610                                    self.triggers |= TestFlag::InterMoveRightLast as u32;
611                                }
612                            } else {
613                                trace_log!(
614                                    "propagate_split rotate_right {grand:?}:{grand_idx} -> {right_parent:?}:0 insert {parent:?}:{idx} {right_ptr:p}"
615                                );
616                                parent.rotate_right(&mut grand, grand_idx, &mut right_parent);
617                                parent.insert_no_split_with_idx(idx, promote_key, right_ptr);
618                            }
619                            return;
620                        }
621                    }
622                }
623                height += 1;
624
625                // Cannot borrow from siblings, need to split internal node
626                let (right, _promote_key) = parent.insert_split(promote_key, right_ptr);
627                promote_key = _promote_key;
628                right_ptr = right.get_ptr();
629                left_ptr = parent.get_ptr();
630                #[cfg(test)]
631                {
632                    self.triggers |= TestFlag::InterSplit as u32;
633                }
634                // Continue to next parent in cache (loop will pop next parent)
635            }
636        }
637        // No more parents in cache, create new root
638        let new_root = InterNode::<K, V>::new_root(height + 1, promote_key, left_ptr, right_ptr);
639        let _old_root = self.root.replace(Node::Inter(new_root)).expect("should have old root");
640        #[cfg(debug_assertions)]
641        {
642            match _old_root {
643                Node::Inter(node) => {
644                    debug_assert_eq!(height, node.height(), "old inter root at {height}");
645                    debug_assert_eq!(node.get_ptr(), left_ptr);
646                }
647                Node::Leaf(node) => {
648                    debug_assert_eq!(height, 0);
649                    debug_assert_eq!(node.get_ptr(), left_ptr);
650                }
651            }
652        }
653    }
654
655    /// Handle leaf node underflow by merging with sibling
656    /// Uses PathCache to accelerate parent lookup
657    /// Following the merge strategy from Designer Notes:
658    /// - Try merge with left sibling (if left + current <= cap)
659    /// - Try merge with right sibling (if current + right <= cap)
660    /// - Try 3-node merge (if left + current + right <= 2 * cap)
661    fn handle_leaf_underflow(&mut self, mut leaf: LeafNode<K, V>, merge: bool) {
662        debug_assert!(!self.get_root_unwrap().is_leaf());
663        let cur_count = leaf.key_count();
664        let cap = LeafNode::<K, V>::cap();
665        debug_assert!(cur_count <= cap >> 1);
666        let mut can_unlink: bool = false;
667        let (mut left_avail, mut right_avail) = (0, 0);
668        let mut merge_right = false;
669        if cur_count == 0 {
670            trace_log!("handle_leaf_underflow {leaf:?} unlink");
671            // if the right and left are full, or they not exist, can come to this
672            can_unlink = true;
673        }
674        if merge {
675            if !can_unlink && let Some(mut left_node) = leaf.get_left_node() {
676                let left_count = left_node.key_count();
677                if left_count + cur_count <= cap {
678                    trace_log!(
679                        "handle_leaf_underflow {leaf:?} merge left {left_node:?} {cur_count}"
680                    );
681                    leaf.copy_left(&mut left_node, cur_count);
682                    can_unlink = true;
683                    #[cfg(test)]
684                    {
685                        self.triggers |= TestFlag::LeafMergeLeft as u32;
686                    }
687                } else {
688                    left_avail = cap - left_count;
689                }
690            }
691            if !can_unlink && let Some(mut right_node) = leaf.get_right_node() {
692                let right_count = right_node.key_count();
693                if right_count + cur_count <= cap {
694                    trace_log!(
695                        "handle_leaf_underflow {leaf:?} merge right {right_node:?} {cur_count}"
696                    );
697                    leaf.copy_right::<false>(&mut right_node, 0, cur_count);
698                    can_unlink = true;
699                    merge_right = true;
700                    #[cfg(test)]
701                    {
702                        self.triggers |= TestFlag::LeafMergeRight as u32;
703                    }
704                } else {
705                    right_avail = cap - right_count;
706                }
707            }
708            // if we require left_avail + right_avail > cur_count, not possible to construct a 3-2
709            // merge, only either triggering merge left or merge right.
710            if !can_unlink
711                && left_avail > 0
712                && right_avail > 0
713                && left_avail + right_avail == cur_count
714            {
715                let mut left_node = leaf.get_left_node().unwrap();
716                let mut right_node = leaf.get_right_node().unwrap();
717                debug_assert!(left_avail < cur_count);
718                trace_log!("handle_leaf_underflow {leaf:?} merge left {left_node:?} {left_avail}");
719                leaf.copy_left(&mut left_node, left_avail);
720                trace_log!(
721                    "handle_leaf_underflow {leaf:?} merge right {right_node:?} {}",
722                    cur_count - left_avail
723                );
724                leaf.copy_right::<false>(&mut right_node, left_avail, cur_count - left_avail);
725                merge_right = true;
726                can_unlink = true;
727                #[cfg(test)]
728                {
729                    self.triggers |=
730                        TestFlag::LeafMergeLeft as u32 | TestFlag::LeafMergeRight as u32;
731                }
732            }
733        }
734        if !can_unlink {
735            return;
736        }
737        self.leaf_count -= 1;
738        let right_sep = if merge_right {
739            let right_node = leaf.get_right_node().unwrap();
740            Some(right_node.clone_first_key())
741        } else {
742            None
743        };
744        let no_right = leaf.unlink().is_null();
745        leaf.dealloc::<false>();
746        let (mut parent, mut idx) = self.get_cache().pop().unwrap();
747        trace_log!("handle_leaf_underflow pop parent {parent:?}:{idx}");
748        if parent.key_count() == 0 {
749            if let Some((grand, grand_idx)) = self.remove_only_child(parent) {
750                trace_log!("handle_leaf_underflow remove_only_child until {grand:?}:{grand_idx}");
751                parent = grand;
752                idx = grand_idx;
753            } else {
754                trace_log!("handle_leaf_underflow remove_only_child all");
755                return;
756            }
757        }
758        self.remove_child_from_inter(&mut parent, idx, right_sep, no_right);
759        if parent.key_count() <= 1 {
760            self.handle_inter_underflow(parent);
761        }
762    }
763
764    /// To simplify the logic, we perform delete first.
765    /// return the Some(node) when need to rebalance
766    #[inline]
767    fn remove_child_from_inter(
768        &mut self, node: &mut InterNode<K, V>, delete_idx: u32, right_sep: Option<K>,
769        _no_right: bool,
770    ) {
771        self.get_cache().assert_center();
772        debug_assert!(node.key_count() > 0, "{:?} {}", node, node.key_count());
773        if delete_idx == node.key_count() {
774            trace_log!("remove_child_from_inter {node:?}:{delete_idx} last");
775            #[cfg(test)]
776            {
777                self.triggers |= TestFlag::RemoveChildLast as u32;
778            }
779            // delete the last child of this node
780            node.remove_last_child();
781            if let Some(key) = right_sep
782                && let Some((grand_parent, grand_idx)) =
783                    self.get_cache().peak_ancenstor(|_node: &InterNode<K, V>, idx: u32| -> bool {
784                        _node.key_count() > idx
785                    })
786            {
787                #[cfg(test)]
788                {
789                    self.triggers |= TestFlag::UpdateSepKey as u32;
790                }
791                trace_log!("remove_child_from_inter change_key {grand_parent:?}:{grand_idx}");
792                // key idx = child idx - 1 , and + 1 for right node
793                grand_parent.change_key(grand_idx, key);
794            }
795        } else if delete_idx > 0 {
796            trace_log!("remove_child_from_inter {node:?}:{delete_idx} mid");
797            node.remove_mid_child(delete_idx);
798            #[cfg(test)]
799            {
800                self.triggers |= TestFlag::RemoveChildMid as u32;
801            }
802            // sep key of right node shift left
803            if let Some(key) = right_sep {
804                trace_log!("remove_child_from_inter change_key {node:?}:{}", delete_idx - 1);
805                node.change_key(delete_idx - 1, key);
806                #[cfg(test)]
807                {
808                    self.triggers |= TestFlag::UpdateSepKey as u32;
809                }
810            }
811        } else {
812            trace_log!("remove_child_from_inter {node:?}:{delete_idx} first");
813            // delete_idx is the first but not the last
814            let mut sep_key = node.remove_first_child();
815            #[cfg(test)]
816            {
817                self.triggers |= TestFlag::RemoveChildFirst as u32;
818            }
819            if let Some(key) = right_sep {
820                sep_key = key;
821            }
822            self.update_ancestor_sep_key::<false>(sep_key);
823        }
824    }
825
826    #[inline]
827    fn handle_inter_underflow(&mut self, mut node: InterNode<K, V>) {
828        let cap = InterNode::<K, V>::cap();
829        while node.key_count() <= InterNode::<K, V>::UNDERFLOW_CAP {
830            if node.key_count() == 0 {
831                let root_height = self.get_root_unwrap().height();
832                if node.height() == root_height
833                    || self
834                        .get_cache()
835                        .peak_ancenstor(|_node: &InterNode<K, V>, _idx: u32| -> bool {
836                            _node.key_count() > 0
837                        })
838                        .is_none()
839                {
840                    let root = unsafe { Node::from_header(*node.child_ptr(0)) };
841                    trace_log!("handle_inter_underflow downgrade root {root:?}");
842                    let _old_root = self.root.replace(root);
843                    debug_assert!(_old_root.is_some());
844
845                    while let Some((parent, _)) = self.get_cache().pop() {
846                        parent.dealloc::<false>();
847                    }
848                    node.dealloc::<false>(); // they all have no key
849                }
850                return;
851            } else {
852                if let Some((mut grand, grand_idx)) = self.get_cache().pop() {
853                    if grand_idx > 0 {
854                        let mut left = grand.get_child_as_inter(grand_idx - 1);
855                        // the sep key should pull down,  key+1 + key + 1 > cap + 1
856                        if left.key_count() + node.key_count() < cap {
857                            #[cfg(test)]
858                            {
859                                self.triggers |= TestFlag::InterMergeLeft as u32;
860                            }
861                            trace_log!(
862                                "handle_inter_underflow {node:?} merge left {left:?} parent {grand:?}:{grand_idx}"
863                            );
864                            left.merge(node, &mut grand, grand_idx);
865                            node = grand;
866                            continue;
867                        }
868                    }
869                    if grand_idx < grand.key_count() {
870                        let right = grand.get_child_as_inter(grand_idx + 1);
871                        // the sep key should pull down,  key+1 + key + 1 > cap + 1
872                        if right.key_count() + node.key_count() < cap {
873                            #[cfg(test)]
874                            {
875                                self.triggers |= TestFlag::InterMergeRight as u32;
876                            }
877                            trace_log!(
878                                "handle_inter_underflow {node:?} cap {cap} merge right {right:?} parent {grand:?}:{}",
879                                grand_idx + 1
880                            );
881                            node.merge(right, &mut grand, grand_idx + 1);
882                            node = grand;
883                            continue;
884                        }
885                    }
886                }
887                return;
888            }
889        }
890    }
891
892    #[inline]
893    fn remove_only_child(&mut self, node: InterNode<K, V>) -> Option<(InterNode<K, V>, u32)> {
894        debug_assert_eq!(node.key_count(), 0);
895        #[cfg(test)]
896        {
897            self.triggers |= TestFlag::RemoveOnlyChild as u32;
898        }
899        if let Some((parent, idx)) = self.get_cache().move_to_ancenstor(
900            |node: &InterNode<K, V>, _idx: u32| -> bool { node.key_count() != 0 },
901            InterNode::<K, V>::dealloc::<false>,
902        ) {
903            node.dealloc::<true>();
904            Some((parent, idx))
905        } else {
906            node.dealloc::<true>();
907            // we are empty, my ancestor are all empty and delete by move_to_ancenstor
908            self.root = None;
909            None
910        }
911    }
912
913    /// Dump the entire tree structure for debugging
914    #[cfg(all(test, feature = "std"))]
915    pub fn dump(&self)
916    where
917        K: Debug,
918        V: Debug,
919    {
920        print_log!("=== BTreeMap Dump ===");
921        print_log!("Length: {}", self.len());
922        if let Some(root) = &self.root {
923            self.dump_node(root, 0);
924        } else {
925            print_log!("(empty)");
926        }
927        print_log!("=====================");
928    }
929
930    #[cfg(all(test, feature = "std"))]
931    fn dump_node(&self, node: &Node<K, V>, depth: usize)
932    where
933        K: Debug,
934        V: Debug,
935    {
936        match node {
937            Node::Leaf(leaf) => {
938                print!("{:indent$}", "", indent = depth * 2);
939                print_log!("{}", leaf);
940            }
941            Node::Inter(inter) => {
942                print!("{:indent$}", "", indent = depth * 2);
943                print_log!("{}", inter);
944                // Dump children
945                let count = inter.key_count() as u32;
946                for i in 0..=count {
947                    unsafe {
948                        let child_ptr = *inter.child_ptr(i);
949                        if !child_ptr.is_null() {
950                            let child_node = if (*child_ptr).is_leaf() {
951                                Node::Leaf(LeafNode::<K, V>::from_header(child_ptr))
952                            } else {
953                                Node::Inter(InterNode::<K, V>::from_header(child_ptr))
954                            };
955                            self.dump_node(&child_node, depth + 1);
956                        }
957                    }
958                }
959            }
960        }
961    }
962
963    /// Validate the entire tree structure
964    /// Uses the same traversal logic as Drop to avoid recursion
965    pub fn validate(&self)
966    where
967        K: Debug,
968        V: Debug,
969    {
970        if self.root.is_none() {
971            assert_eq!(self.len, 0, "Empty tree should have len 0");
972            return;
973        }
974
975        let mut total_keys = 0usize;
976        let mut prev_leaf_max: Option<K> = None;
977
978        match &self.root {
979            Some(Node::Leaf(leaf)) => {
980                total_keys += leaf.validate(None, None);
981            }
982            Some(Node::Inter(inter)) => {
983                // Do not use btree internal PathCache (might distrupt test scenario)
984                let mut cache = PathCache::new();
985                let mut cur = inter.clone();
986                loop {
987                    cache.push(cur.clone(), 0);
988                    cur.validate();
989                    match cur.get_child(0) {
990                        Node::Leaf(leaf) => {
991                            // Validate first leaf with no min/max bounds from parent
992                            let min_key: Option<K> = None;
993                            let max_key = if inter.key_count() > 0 {
994                                unsafe { Some((*inter.key_ptr(0)).assume_init_ref().clone()) }
995                            } else {
996                                None
997                            };
998                            total_keys += leaf.validate(min_key.as_ref(), max_key.as_ref());
999                            if let Some(ref prev_max) = prev_leaf_max {
1000                                let first_key = unsafe { (*leaf.key_ptr(0)).assume_init_ref() };
1001                                assert!(
1002                                    prev_max < first_key,
1003                                    "{:?} Leaf keys not in order: prev max {:?} >= current min {:?}",
1004                                    leaf,
1005                                    prev_max,
1006                                    first_key
1007                                );
1008                            }
1009                            prev_leaf_max = unsafe {
1010                                Some(
1011                                    (*leaf.key_ptr(leaf.key_count() - 1)).assume_init_ref().clone(),
1012                                )
1013                            };
1014                            break;
1015                        }
1016                        Node::Inter(child_inter) => {
1017                            cur = child_inter;
1018                        }
1019                    }
1020                }
1021
1022                // Continue traversal like Drop does
1023                while let Some((parent, idx)) =
1024                    cache.move_right_and_pop_l1(dummy_post_callback::<K, V>)
1025                {
1026                    cache.push(parent.clone(), idx);
1027                    if let Node::Leaf(leaf) = parent.get_child(idx) {
1028                        // Calculate bounds for this leaf
1029                        let min_key = if idx > 0 {
1030                            unsafe { Some((*parent.key_ptr(idx - 1)).assume_init_ref().clone()) }
1031                        } else {
1032                            None
1033                        };
1034                        let max_key = if idx < parent.key_count() {
1035                            unsafe { Some((*parent.key_ptr(idx)).assume_init_ref().clone()) }
1036                        } else {
1037                            None
1038                        };
1039                        total_keys += leaf.validate(min_key.as_ref(), max_key.as_ref());
1040
1041                        // Check ordering with previous leaf
1042                        if let Some(ref prev_max) = prev_leaf_max {
1043                            let first_key = unsafe { (*leaf.key_ptr(0)).assume_init_ref() };
1044                            assert!(
1045                                prev_max < first_key,
1046                                "{:?} Leaf keys not in order: prev max {:?} >= current min {:?}",
1047                                leaf,
1048                                prev_max,
1049                                first_key
1050                            );
1051                        }
1052                        prev_leaf_max = unsafe {
1053                            Some((*leaf.key_ptr(leaf.key_count() - 1)).assume_init_ref().clone())
1054                        };
1055                    } else {
1056                        panic!("{parent:?} child {:?} is not leaf", parent.get_child(idx));
1057                    }
1058                }
1059            }
1060            None => unreachable!(),
1061        }
1062        assert_eq!(
1063            total_keys, self.len,
1064            "Total keys in tree ({}) doesn't match len ({})",
1065            total_keys, self.len
1066        );
1067    }
1068
1069    /// Returns the first key-value pair in the map
1070    /// Returns `None` if the map is empty
1071    #[inline]
1072    pub fn first_key_value(&self) -> Option<(&K, &V)> {
1073        let leaf = match &self.root {
1074            None => return None,
1075            Some(root) => root.find_first_leaf(None),
1076        };
1077        if leaf.key_count() == 0 {
1078            return None;
1079        }
1080        unsafe {
1081            let key = (*leaf.key_ptr(0)).assume_init_ref();
1082            let value = (*leaf.value_ptr(0)).assume_init_ref();
1083            Some((key, value))
1084        }
1085    }
1086
1087    /// Returns the last key-value pair in the map
1088    /// Returns `None` if the map is empty
1089    #[inline]
1090    pub fn last_key_value(&self) -> Option<(&K, &V)> {
1091        let leaf = match &self.root {
1092            None => return None,
1093            Some(root) => root.find_last_leaf(None),
1094        };
1095        let count = leaf.key_count();
1096        if count == 0 {
1097            return None;
1098        }
1099        unsafe {
1100            let last_idx = count - 1;
1101            let key = (*leaf.key_ptr(last_idx)).assume_init_ref();
1102            let value = (*leaf.value_ptr(last_idx)).assume_init_ref();
1103            Some((key, value))
1104        }
1105    }
1106
1107    /// Returns an entry to the first key in the map
1108    /// Returns `None` if the map is empty
1109    #[inline]
1110    pub fn first_entry(&mut self) -> Option<OccupiedEntry<'_, K, V>> {
1111        let cache = self.get_cache();
1112        cache.clear();
1113        let leaf = match &self.root {
1114            None => return None,
1115            Some(root) => {
1116                // Populate cache with the path to first leaf
1117                root.find_first_leaf(Some(cache))
1118            }
1119        };
1120        if leaf.key_count() == 0 {
1121            return None;
1122        }
1123        Some(OccupiedEntry { tree: self, idx: 0, leaf })
1124    }
1125
1126    /// Returns an entry to the last key in the map
1127    /// Returns `None` if the map is empty
1128    #[inline]
1129    pub fn last_entry(&mut self) -> Option<OccupiedEntry<'_, K, V>> {
1130        let cache = self.get_cache();
1131        cache.clear();
1132        let leaf = match &self.root {
1133            None => return None,
1134            Some(root) => {
1135                // Populate cache with the path to last leaf
1136                root.find_last_leaf(Some(cache))
1137            }
1138        };
1139        let count = leaf.key_count();
1140        if count == 0 {
1141            return None;
1142        }
1143        Some(OccupiedEntry { tree: self, idx: count - 1, leaf })
1144    }
1145
1146    /// Removes and returns the first key-value pair in the map
1147    /// Returns `None` if the map is empty
1148    #[inline]
1149    pub fn pop_first(&mut self) -> Option<(K, V)> {
1150        match self.first_entry() {
1151            Some(entry) => Some(entry.remove_entry()),
1152            _ => None,
1153        }
1154    }
1155
1156    /// Removes and returns the last key-value pair in the map
1157    /// Returns `None` if the map is empty
1158    #[inline]
1159    pub fn pop_last(&mut self) -> Option<(K, V)> {
1160        match self.last_entry() {
1161            Some(entry) => Some(entry.remove_entry()),
1162            _ => None,
1163        }
1164    }
1165
1166    /// Returns an iterator over the map's entries
1167    #[inline]
1168    pub fn iter(&self) -> Iter<'_, K, V> {
1169        Iter::new(self.find_first_and_last_leaf(), self.len)
1170    }
1171
1172    /// Returns a mutable iterator over the map's entries
1173    #[inline]
1174    pub fn iter_mut(&mut self) -> IterMut<'_, K, V> {
1175        IterMut::new(self.find_first_and_last_leaf(), self.len)
1176    }
1177
1178    /// Returns an iterator over the map's keys
1179    #[inline]
1180    pub fn keys(&self) -> Keys<'_, K, V> {
1181        Keys::new(self.iter())
1182    }
1183
1184    /// Returns an iterator over the map's values
1185    #[inline]
1186    pub fn values(&self) -> Values<'_, K, V> {
1187        Values::new(self.iter())
1188    }
1189
1190    /// Returns a mutable iterator over the map's values
1191    #[inline]
1192    pub fn values_mut(&mut self) -> ValuesMut<'_, K, V> {
1193        ValuesMut::new(self.iter_mut())
1194    }
1195
1196    #[inline]
1197    fn find_first_and_last_leaf(&self) -> Option<(LeafNode<K, V>, LeafNode<K, V>)> {
1198        let root = self.root.as_ref()?;
1199        Some((root.find_first_leaf(None), root.find_last_leaf(None)))
1200    }
1201
1202    /// Internal helper to find range bounds
1203    /// Returns (front_leaf, front_idx, back_leaf, back_idx) where both leaves are Some or both are None
1204    #[inline]
1205    fn find_range_bounds<R>(&self, range: R) -> Option<RangeBase<'_, K, V>>
1206    where
1207        R: RangeBounds<K>,
1208    {
1209        let root = self.root.as_ref()?;
1210        let (front_leaf, front_idx) = root.find_leaf_with_bound(range.start_bound(), true);
1211        let (back_leaf, back_idx) = root.find_leaf_with_bound(range.end_bound(), false);
1212        Some(RangeBase::new(front_leaf, front_idx, back_leaf, back_idx))
1213    }
1214
1215    /// Returns an iterator over a sub-range of entries in the map
1216    #[inline]
1217    pub fn range<R>(&self, range: R) -> Range<'_, K, V>
1218    where
1219        R: RangeBounds<K>,
1220    {
1221        Range::new(self.find_range_bounds(range))
1222    }
1223
1224    /// Returns a mutable iterator over a sub-range of entries in the map
1225    #[inline]
1226    pub fn range_mut<R>(&mut self, range: R) -> RangeMut<'_, K, V>
1227    where
1228        R: RangeBounds<K>,
1229    {
1230        RangeMut::new(self.find_range_bounds(range))
1231    }
1232}
1233
1234impl<K: Ord + Clone + Sized, V: Sized> Default for BTreeMap<K, V> {
1235    fn default() -> Self {
1236        Self::new()
1237    }
1238}
1239
1240impl<K: Ord + Clone + Sized, V: Sized> Drop for BTreeMap<K, V> {
1241    fn drop(&mut self) {
1242        let mut cache = self.get_cache().take();
1243        let mut cur = match self.root.take() {
1244            None => return,
1245            Some(root) => root.find_first_leaf(Some(&mut cache)),
1246        };
1247        cur.dealloc::<true>();
1248        // To navigate to next leaf,
1249        // return None when reach the end
1250        while let Some((parent, idx)) = cache.move_right_and_pop_l1(InterNode::dealloc::<true>) {
1251            cache.push(parent.clone(), idx);
1252            cur = parent.get_child_as_leaf(idx);
1253            cur.dealloc::<true>();
1254        }
1255    }
1256}
1257
1258impl<K: Ord + Clone + Sized, V: Sized> IntoIterator for BTreeMap<K, V> {
1259    type Item = (K, V);
1260    type IntoIter = IntoIter<K, V>;
1261
1262    #[inline]
1263    fn into_iter(self) -> Self::IntoIter {
1264        IntoIter::new(self, true)
1265    }
1266}
1267impl<'a, K: Ord + Clone + Sized, V: Sized> IntoIterator for &'a BTreeMap<K, V> {
1268    type Item = (&'a K, &'a V);
1269    type IntoIter = Iter<'a, K, V>;
1270
1271    #[inline]
1272    fn into_iter(self) -> Self::IntoIter {
1273        self.iter()
1274    }
1275}
1276
1277impl<'a, K: Ord + Clone + Sized, V: Sized> IntoIterator for &'a mut BTreeMap<K, V> {
1278    type Item = (&'a K, &'a mut V);
1279    type IntoIter = IterMut<'a, K, V>;
1280
1281    #[inline]
1282    fn into_iter(self) -> Self::IntoIter {
1283        self.iter_mut()
1284    }
1285}