@@ -1780,18 +1780,12 @@ where
1780
1780
where
1781
1781
F : FnMut ( & K , & mut V ) -> bool ,
1782
1782
{
1783
- while let Some ( kv) = unsafe { self . next_kv ( ) } {
1784
- let ( k, v) = unsafe { ptr :: read ( & kv ) } . into_kv_mut ( ) ;
1783
+ while let Some ( mut kv) = unsafe { self . next_kv ( ) } {
1784
+ let ( k, v) = kv . kv_mut ( ) ;
1785
1785
if pred ( k, v) {
1786
1786
* self . length -= 1 ;
1787
1787
let ( k, v, leaf_edge_location) = kv. remove_kv_tracking ( ) ;
1788
- // `remove_kv_tracking` has either preserved or invalidated `self.cur_leaf_edge`
1789
- if let Some ( node) = leaf_edge_location {
1790
- match search:: search_tree ( node, & k) {
1791
- search:: SearchResult :: Found ( _) => unreachable ! ( ) ,
1792
- search:: SearchResult :: GoDown ( leaf) => self . cur_leaf_edge = Some ( leaf) ,
1793
- }
1794
- } ;
1788
+ self . cur_leaf_edge = Some ( leaf_edge_location) ;
1795
1789
return Some ( ( k, v) ) ;
1796
1790
}
1797
1791
self . cur_leaf_edge = Some ( kv. next_leaf_edge ( ) ) ;
@@ -2698,108 +2692,124 @@ impl<'a, K: Ord, V> OccupiedEntry<'a, K, V> {
2698
2692
2699
2693
impl < ' a , K : ' a , V : ' a > Handle < NodeRef < marker:: Mut < ' a > , K , V , marker:: LeafOrInternal > , marker:: KV > {
2700
2694
/// Removes a key/value-pair from the map, and returns that pair, as well as
2701
- /// the whereabouts of the leaf edge corresponding to that former pair:
2702
- /// if None is returned, the leaf edge is still the left leaf edge of the KV handle;
2703
- /// if a node is returned, it heads the subtree where the leaf edge may be found.
2695
+ /// the leaf edge corresponding to that former pair.
2704
2696
fn remove_kv_tracking (
2705
2697
self ,
2706
- ) -> ( K , V , Option < NodeRef < marker:: Mut < ' a > , K , V , marker:: LeafOrInternal > > ) {
2707
- let mut levels_down_handled: isize ;
2708
- let ( small_leaf, old_key, old_val) = match self . force ( ) {
2698
+ ) -> ( K , V , Handle < NodeRef < marker:: Mut < ' a > , K , V , marker:: Leaf > , marker:: Edge > ) {
2699
+ let ( mut pos, old_key, old_val, was_internal) = match self . force ( ) {
2709
2700
Leaf ( leaf) => {
2710
- levels_down_handled = 1 ; // handled at same level, but affects only the right side
2711
2701
let ( hole, old_key, old_val) = leaf. remove ( ) ;
2712
- ( hole. into_node ( ) , old_key, old_val)
2702
+ ( hole, old_key, old_val, false )
2713
2703
}
2714
2704
Internal ( mut internal) => {
2715
2705
// Replace the location freed in the internal node with the next KV,
2716
2706
// and remove that next KV from its leaf.
2717
- levels_down_handled = unsafe { ptr:: read ( & internal) . into_node ( ) . height ( ) } as isize ;
2718
2707
2719
2708
let key_loc = internal. kv_mut ( ) . 0 as * mut K ;
2720
2709
let val_loc = internal. kv_mut ( ) . 1 as * mut V ;
2721
2710
2722
- let to_remove = internal. right_edge ( ) . descend ( ) . first_leaf_edge ( ) . right_kv ( ) . ok ( ) ;
2711
+ // Deleting from the left side is typically faster since we can
2712
+ // just pop an element from the end of the KV array without
2713
+ // needing to shift the other values.
2714
+ let to_remove = internal. left_edge ( ) . descend ( ) . last_leaf_edge ( ) . left_kv ( ) . ok ( ) ;
2723
2715
let to_remove = unsafe { unwrap_unchecked ( to_remove) } ;
2724
2716
2725
2717
let ( hole, key, val) = to_remove. remove ( ) ;
2726
2718
2727
2719
let old_key = unsafe { mem:: replace ( & mut * key_loc, key) } ;
2728
2720
let old_val = unsafe { mem:: replace ( & mut * val_loc, val) } ;
2729
2721
2730
- ( hole. into_node ( ) , old_key, old_val)
2722
+ ( hole, old_key, old_val, true )
2731
2723
}
2732
2724
} ;
2733
2725
2734
2726
// Handle underflow
2735
- let mut cur_node = small_leaf. forget_type ( ) ;
2727
+ let mut cur_node = unsafe { ptr:: read ( & pos) . into_node ( ) . forget_type ( ) } ;
2728
+ let mut at_leaf = true ;
2736
2729
while cur_node. len ( ) < node:: MIN_LEN {
2737
2730
match handle_underfull_node ( cur_node) {
2738
- AtRoot ( root) => {
2739
- cur_node = root;
2740
- break ;
2741
- }
2742
- EmptyParent ( _) => unreachable ! ( ) ,
2743
- Merged ( parent) => {
2744
- levels_down_handled -= 1 ;
2731
+ AtRoot => break ,
2732
+ Merged ( edge, merged_with_left, offset) => {
2733
+ // If we merged with our right sibling then our tracked
2734
+ // position has not changed. However if we merged with our
2735
+ // left sibling then our tracked position is now dangling.
2736
+ if at_leaf && merged_with_left {
2737
+ let idx = pos. idx ( ) + offset;
2738
+ let node = match unsafe { ptr:: read ( & edge) . descend ( ) . force ( ) } {
2739
+ Leaf ( leaf) => leaf,
2740
+ Internal ( _) => unreachable ! ( ) ,
2741
+ } ;
2742
+ pos = unsafe { Handle :: new_edge ( node, idx) } ;
2743
+ }
2744
+
2745
+ let parent = edge. into_node ( ) ;
2745
2746
if parent. len ( ) == 0 {
2746
2747
// We must be at the root
2747
- let root = parent. into_root_mut ( ) ;
2748
- root. pop_level ( ) ;
2749
- cur_node = root. as_mut ( ) ;
2748
+ parent. into_root_mut ( ) . pop_level ( ) ;
2750
2749
break ;
2751
2750
} else {
2752
2751
cur_node = parent. forget_type ( ) ;
2752
+ at_leaf = false ;
2753
2753
}
2754
2754
}
2755
- Stole ( internal_node) => {
2756
- levels_down_handled -= 1 ;
2757
- cur_node = internal_node. forget_type ( ) ;
2755
+ Stole ( stole_from_left) => {
2756
+ // Adjust the tracked position if we stole from a left sibling
2757
+ if stole_from_left && at_leaf {
2758
+ // SAFETY: This is safe since we just added an element to our node.
2759
+ unsafe {
2760
+ pos. next_unchecked ( ) ;
2761
+ }
2762
+ }
2763
+
2758
2764
// This internal node might be underfull, but only if it's the root.
2759
2765
break ;
2760
2766
}
2761
2767
}
2762
2768
}
2763
2769
2764
- let leaf_edge_location = if levels_down_handled > 0 { None } else { Some ( cur_node) } ;
2765
- ( old_key, old_val, leaf_edge_location)
2770
+ // If we deleted from an internal node then we need to compensate for
2771
+ // the earlier swap and adjust the tracked position to point to the
2772
+ // next element.
2773
+ if was_internal {
2774
+ pos = unsafe { unwrap_unchecked ( pos. next_kv ( ) . ok ( ) ) . next_leaf_edge ( ) } ;
2775
+ }
2776
+
2777
+ ( old_key, old_val, pos)
2766
2778
}
2767
2779
}
2768
2780
2769
2781
enum UnderflowResult < ' a , K , V > {
2770
- AtRoot ( NodeRef < marker:: Mut < ' a > , K , V , marker:: LeafOrInternal > ) ,
2771
- EmptyParent ( NodeRef < marker:: Mut < ' a > , K , V , marker:: Internal > ) ,
2772
- Merged ( NodeRef < marker:: Mut < ' a > , K , V , marker:: Internal > ) ,
2773
- Stole ( NodeRef < marker:: Mut < ' a > , K , V , marker:: Internal > ) ,
2782
+ AtRoot ,
2783
+ Merged ( Handle < NodeRef < marker:: Mut < ' a > , K , V , marker:: Internal > , marker:: Edge > , bool , usize ) ,
2784
+ Stole ( bool ) ,
2774
2785
}
2775
2786
2776
2787
fn handle_underfull_node < K , V > (
2777
2788
node : NodeRef < marker:: Mut < ' _ > , K , V , marker:: LeafOrInternal > ,
2778
2789
) -> UnderflowResult < ' _ , K , V > {
2779
2790
let parent = match node. ascend ( ) {
2780
2791
Ok ( parent) => parent,
2781
- Err ( root ) => return AtRoot ( root ) ,
2792
+ Err ( _ ) => return AtRoot ,
2782
2793
} ;
2783
2794
2784
2795
let ( is_left, mut handle) = match parent. left_kv ( ) {
2785
2796
Ok ( left) => ( true , left) ,
2786
- Err ( parent) => match parent. right_kv ( ) {
2787
- Ok ( right) => ( false , right) ,
2788
- Err ( parent) => {
2789
- return EmptyParent ( parent. into_node ( ) ) ;
2790
- }
2791
- } ,
2797
+ Err ( parent) => {
2798
+ let right = unsafe { unwrap_unchecked ( parent. right_kv ( ) . ok ( ) ) } ;
2799
+ ( false , right)
2800
+ }
2792
2801
} ;
2793
2802
2794
2803
if handle. can_merge ( ) {
2795
- Merged ( handle. merge ( ) . into_node ( ) )
2804
+ let offset = if is_left { handle. reborrow ( ) . left_edge ( ) . descend ( ) . len ( ) + 1 } else { 0 } ;
2805
+ Merged ( handle. merge ( ) , is_left, offset)
2796
2806
} else {
2797
2807
if is_left {
2798
2808
handle. steal_left ( ) ;
2799
2809
} else {
2800
2810
handle. steal_right ( ) ;
2801
2811
}
2802
- Stole ( handle . into_node ( ) )
2812
+ Stole ( is_left )
2803
2813
}
2804
2814
}
2805
2815
0 commit comments