Repository: rust-lang/rust
File: library/alloc/src/collections/btree/map.rs (Rust)

Code before:
            .next_back_kv()
            // This should be unwrap(), but that doesn't work because NodeRef
            // doesn't implement Debug. The condition is checked above.
            .ok()?
            .remove_kv_tracking(|| emptied_internal_root = true, self.alloc.clone());
        self.current = Some(pos);
        *self.length -= 1;
        if emptied_internal_root {
            // SAFETY: This is safe since current does not point within the now
            // empty root node.
            let root = unsafe { self.root.reborrow().as_mut().unwrap() };
            root.pop_internal_level(self.alloc.clone());
        }
        Some(kv)
    }
}

impl<'a, K: Ord, V, A: Allocator + Clone> CursorMut<'a, K, V, A> {
    /// Calls a function with a reference to the next element's key and
    /// ownership of its value. The function is expected to return a value
    /// to write back to the next element's value. The cursor is not
    /// advanced forward.
    ///
    /// If the cursor is at the end of the map then the function is not called
    /// and this essentially does not do anything.
    pub(super) fn with_next(&mut self, f: impl FnOnce(&K, V) -> V) {
        // FIXME: This can be optimized to not do all the removing/reinserting
        // logic by using ptr::read, calling `f`, and then using ptr::write.
        // if `f` unwinds, then we need to remove the entry while being careful to
        // not cause UB by moving or dropping the already-dropped `V`
        // for the entry. Some implementation ideas:
        // https://github.com/rust-lang/rust/pull/152418#discussion_r2800232576
        if let Some((k, v)) = self.remove_next() {
            // SAFETY: we remove the K, V out of the next entry,
            // apply 'f' to get a new V, and insert (K, V) back
            // into the next entry that the cursor is pointing at
            let v = f(&k, v);
            unsafe { self.insert_after_unchecked(k, v) };
        }
    }

    /// Inserts a new key-value pair into the map in the gap that the
    /// cursor is currently pointing to.
    ///
    /// After the insertion the cursor will be pointing at the gap after the
    /// newly inserted element.
    ///
    /// # Safety
    ///
    /// You must ensure that the `BTreeMap` invariants are maintained.
    /// Specifically:

Review comment (refactor):
I feel like `with_next` having mutating action on `self` is not particularly intuitive. Maybe we can call this something like `replace_next`? But also it seems like at least the current impl can just be inlined into the call site and we can remove this entirely?

Code after:
            // empty root node.
            let root = unsafe { self.root.reborrow().as_mut().unwrap() };
            root.pop_internal_level(self.alloc.clone());
        }
        Some(kv)
    }
}

impl<'a, K: Ord, V, A: Allocator + Clone> CursorMut<'a, K, V, A> {
    /// Inserts a new key-value pair into the map in the gap that the
    /// cursor is currently pointing to.
    ///
    /// After the insertion the cursor will be pointing at the gap after the
    /// newly inserted element.
    ///
    /// # Safety
    ///
    /// You must ensure that the `BTreeMap` invariants are maintained.
    /// Specifically:
    ///
    /// * The key of the newly inserted element must be unique in the tree.
    /// * All keys in the tree must remain in sorted order.
    #[unstable(feature = "btree_cursors", issue = "107540")]
    pub unsafe fn insert_after_unchecked(&mut self, key: K, value: V) {
        unsafe { self.inner.insert_after_unchecked(key, value) }
    }

    /// Inserts a new key-value pair into the map in the gap that the
    /// cursor is currently pointing to.
    ///
    /// After the insertion the cursor will be pointing at the gap after the
    /// newly inserted element.
    ///
    /// # Safety
    ///
    /// You must ensure that the `BTreeMap` invariants are maintained.
    /// Specifically:
    ///
    /// * The key of the newly inserted element must be unique in the tree.
    /// * All keys in the tree must remain in sorted order.
    #[unstable(feature = "btree_cursors", issue = "107540")]
    pub unsafe fn insert_before_unchecked(&mut self, key: K, value: V) {
        unsafe { self.inner.insert_before_unchecked(key, value) }
    }

    /// Inserts a new key-value pair into the map in the gap that the
    /// cursor is currently pointing to.
    ///
    /// After the insertion the cursor will be pointing at the gap before the
    /// newly inserted element.
    ///