      1
      2
      3
      4
      5
      6
      7
      8
      9
     10
     11
     12
     13
     14
     15
     16
     17
     18
     19
     20
     21
     22
     23
     24
     25
     26
     27
     28
     29
     30
     31
     32
     33
     34
     35
     36
     37
     38
     39
     40
     41
     42
     43
     44
     45
     46
     47
     48
     49
     50
     51
     52
     53
     54
     55
     56
     57
     58
     59
     60
     61
     62
     63
     64
     65
     66
     67
     68
     69
     70
     71
     72
     73
     74
     75
     76
     77
     78
     79
     80
     81
     82
     83
     84
     85
     86
     87
     88
     89
     90
     91
     92
     93
     94
     95
     96
     97
     98
     99
    100
    101
    102
    103
    104
    105
    106
    107
    108
    109
    110
    111
    112
    113
    114
    115
    116
    117
    118
    119
    120
    121
    122
    123
    124
    125
    126
    127
    128
    129
    130
    131
    132
    133

    use super::map::MIN_LEN;
    use super::node::{marker, ForceResult, Handle, NodeRef};
    use super::unwrap_unchecked;
    use core::mem;
    use core::ptr;

    impl<'a, K: 'a, V: 'a> Handle<NodeRef<marker::Mut<'a>, K, V, marker::LeafOrInternal>, marker::KV> {
        /// Removes a key/value-pair from the map, and returns that pair, as well as
        /// the leaf edge corresponding to that former pair.
        pub fn remove_kv_tracking<F: FnOnce()>(
            self,
            handle_emptied_internal_root: F,
        ) -> ((K, V), Handle<NodeRef<marker::Mut<'a>, K, V, marker::Leaf>, marker::Edge>) {
            let (old_kv, mut pos, was_internal) = match self.force() {
                ForceResult::Leaf(leaf) => {
                    let (old_kv, pos) = leaf.remove();
                    (old_kv, pos, false)
                }
                ForceResult::Internal(mut internal) => {
                    // Replace the location freed in the internal node with an
                    // adjacent KV, and remove that adjacent KV from its leaf.
                    // Always choose the adjacent KV on the left side because
                    // it is typically faster to pop an element from the end
                    // of the KV arrays without needing to shift other elements.

                    let key_loc = internal.kv_mut().0 as *mut K;
                    let val_loc = internal.kv_mut().1 as *mut V;

                    let to_remove = internal.left_edge().descend().last_leaf_edge().left_kv().ok();
                    let to_remove = unsafe { unwrap_unchecked(to_remove) };

                    let (kv, pos) = to_remove.remove();

                    let old_key = unsafe { mem::replace(&mut *key_loc, kv.0) };
                    let old_val = unsafe { mem::replace(&mut *val_loc, kv.1) };

                    ((old_key, old_val), pos, true)
                }
            };

            // Handle underflow
            let mut cur_node = unsafe { ptr::read(&pos).into_node().forget_type() };
            let mut at_leaf = true;
            while cur_node.len() < MIN_LEN {
                match handle_underfull_node(cur_node) {
                    UnderflowResult::AtRoot => break,
                    UnderflowResult::Merged(edge, merged_with_left, offset) => {
                        // If we merged with our right sibling then our tracked
                        // position has not changed. However if we merged with our
                        // left sibling then our tracked position is now dangling.
                        if at_leaf && merged_with_left {
                            let idx = pos.idx() + offset;
                            let node = match unsafe { ptr::read(&edge).descend().force() } {
                                ForceResult::Leaf(leaf) => leaf,
                                ForceResult::Internal(_) => unreachable!(),
                            };
                            pos = unsafe { Handle::new_edge(node, idx) };
                        }

                        let parent = edge.into_node();
                        if parent.len() == 0 {
                            // The parent that was just emptied must be the root,
                            // because nodes on a lower level would not have been
                            // left with a single child.
                            handle_emptied_internal_root();
                            break;
                        } else {
                            cur_node = parent.forget_type();
                            at_leaf = false;
                        }
                    }
                    UnderflowResult::Stole(stole_from_left) => {
                        // Adjust the tracked position if we stole from a left sibling
                        if stole_from_left && at_leaf {
                            // SAFETY: This is safe since we just added an element to our node.
                            unsafe {
                                pos.move_next_unchecked();
                            }
                        }
                        break;
                    }
                }
            }

            // If we deleted from an internal node then we need to compensate for
            // the earlier swap and adjust the tracked position to point to the
            // next element.
            if was_internal {
                pos = unsafe { unwrap_unchecked(pos.next_kv().ok()).next_leaf_edge() };
            }

            (old_kv, pos)
        }
    }

    enum UnderflowResult<'a, K, V> {
        AtRoot,
        Merged(Handle<NodeRef<marker::Mut<'a>, K, V, marker::Internal>, marker::Edge>, bool, usize),
        Stole(bool),
    }

    fn handle_underfull_node<'a, K: 'a, V: 'a>(
        node: NodeRef<marker::Mut<'a>, K, V, marker::LeafOrInternal>,
    ) -> UnderflowResult<'_, K, V> {
        let parent = match node.ascend() {
            Ok(parent) => parent,
            Err(_) => return UnderflowResult::AtRoot,
        };

        // Prefer the left KV if it exists. Merging with the left side is faster,
        // since merging happens towards the left and `node` has fewer elements.
        // Stealing from the left side is faster, since we can pop from the end of
        // the KV arrays.
        let (is_left, mut handle) = match parent.left_kv() {
            Ok(left) => (true, left),
            Err(parent) => {
                let right = unsafe { unwrap_unchecked(parent.right_kv().ok()) };
                (false, right)
            }
        };

        if handle.can_merge() {
            let offset = if is_left { handle.reborrow().left_edge().descend().len() + 1 } else { 0 };
            UnderflowResult::Merged(handle.merge(), is_left, offset)
        } else {
            if is_left {
                handle.steal_left();
            } else {
                handle.steal_right();
            }
            UnderflowResult::Stole(is_left)
        }
    }