BTreeMap::merge optimized

[asder8215]

This is an optimized version of #151981. See ACP and tracking issue for more information on BTreeMap::merge does.

CC @programmerjake. Let me know what you think of how I'm using CursorMut and IntoIter here and whether the unsafe code here looks good. I decided to use ptr::read() and ptr::write() to grab the value from CursorMut as V than &mut V, use it within the conflict function, and overwrite the content of conflicting key afterward.

I know this needs some polishing, especially with refactoring some redundant looking code in a nicer way, some of which could probably just be public API methods for CursorMut. It does pass all the tests that I currently have for BTreeMap::merge (inspired from BTreeMap::append) though, so that's good.

[programmerjake]

this should really be a method of Cursor*, you need a handler for when conflict panics that removes the entry without dropping whatever you ptr::read from.
something like:

impl<K, V> CursorMut<'a, K, V> {
    /// call `f` with the next entry's key and value, replacing the next entry's value with the returned value. if `f` unwinds, the next entry is removed.
    /// equivalent to a more efficient version of:
    /// ```rust
    /// if let Some((k, v)) = self.remove_next() {
    ///     let v = f(&k, v);
    ///     // Safety: key is unmodified
    ///     unsafe { self.insert_after_unchecked(k, v) };
    /// }
    /// ```
    pub(super) fn with_next(&mut self, f: impl FnOnce(&K, V) -> V) {
        struct RemoveNextOnDrop<'a, 'b, K, V> {
            cursor: &'a mut CursorMut<'b, K, V>,
            forget_next: bool,
        }
        impl<K, V> Drop for RemoveNextOnDrop<'_, '_, K, V> {
            fn drop(&mut self) {
                if self.forget_next {
                    // call an equivalent to CursorMut::remove_next()
                    // except that instead of returning `V`, it never moves or drops it.
                    self.0.forget_next_value();
                }
            }
        }
        let mut remove_next_on_drop = RemoveNextOnDrop {
            cursor: self,
            forget_next: false, // we don't know that we have a next value yet
        };
        if let Some((k, v_mut)) = remove_next_on_drop.cursor.peek_next() {
            remove_next_on_drop.forget_next = true;
            // Safety: we move the V out of the next entry,
            // we marked the entry's value to be forgotten
            // when remove_next_on_drop is dropped that
            // way we avoid returning to the caller leaving
            // a moved-out invalid value if `f` unwinds.
            let v = unsafe { std::ptr::read(v_mut) };
            let v = f(k, v);
            // Safety: move the V back into the next entry
            unsafe { std::ptr::write(v_mut, v) };
            remove_next_on_drop.forget_next = false;
        }
    }
}

The equivalent CursorMutKey method should instead have f be impl FnOnce(K, V) -> (K, V) and needs to forget both the key and value since they were both ptr::read, and not just the value.

[asder8215]

Got you 👍

[asder8215]

Appreciate your patience and feedback in all of this @programmerjake!

[asder8215]

@programmerjake Took your suggestion on with_next() on Cursor* and made and forget_next*() functions. Lmk what you think about this!

Also, I've been wondering if with_next() has uses to be a public accessible function?

[programmerjake]

remove_kv_tracking may panic and drop the key and value, tbh this is very complicated and makes me think it could be better for with_next() instead of calling forget_next[_key_and]_value, to transmute the cursor reference type from &mut CursorMut<K, V> to &mut CursorMut<K, ManuallyDrop<V>> and then just call remove_next. that said, idk if that would be sound, asking on Zulip: #t-libs > could we use transmute tricks in the impl of BTreeMap? @ 💬

[asder8215]

True, I think transmute should be fine since &mut CursorMut<K, V> and &mut CursorMut<K, ManuallyDrop<V>> are the same size and it forgets the original, so we wouldn't be running the destructor on the original. (but of course I can't hold myself to that since I haven't played around with mem::transmute much myself)

[programmerjake]

it might differ with -Zrandomize-layout or other future struct layout changes. I haven't reviewed the whole BTreeMap implementation to see what layout guarantees it has.

[rustbot]

Reminder, once the PR becomes ready for a review, use @rustbot ready.

[asder8215]

@programmerjake Made the change, lmk what you think!

[programmerjake]

looks good enough, though I didn't really review the tests.

View changes since this review

[asder8215]

The tests I took verbatim from append, which seemed thorough enough in examining whether the edges are fixed after insertion, drop behavior, and checking if the key is not overwritten on conflict.

One thing I'll point out is that interestingly this append test:

#[test]
#[cfg_attr(not(panic = "unwind"), ignore = "test requires unwinding support")]
fn test_append_drop_leak() {
    let a = CrashTestDummy::new(0);
    let b = CrashTestDummy::new(1);
    let c = CrashTestDummy::new(2);
    let mut left = BTreeMap::new();
    let mut right = BTreeMap::new();
    left.insert(a.spawn(Panic::Never), ());
    left.insert(b.spawn(Panic::Never), ());
    left.insert(c.spawn(Panic::Never), ());
    right.insert(b.spawn(Panic::InDrop), ()); // first duplicate key, dropped during append
    right.insert(c.spawn(Panic::Never), ());

    catch_unwind(move || left.append(&mut right)).unwrap_err();
    assert_eq!(a.dropped(), 1);
    assert_eq!(b.dropped(), 1); // should be 2 were it not for Rust issue #47949
    assert_eq!(c.dropped(), 2);
}

When I ported a similar version of this for merge:

#[test]
#[cfg_attr(not(panic = "unwind"), ignore = "test requires unwinding support")]
fn test_merge_drop_leak() {
    let a = CrashTestDummy::new(0);
    let b = CrashTestDummy::new(1);
    let c = CrashTestDummy::new(2);
    let mut left = BTreeMap::new();
    let mut right = BTreeMap::new();
    left.insert(a.spawn(Panic::Never), ());
    left.insert(b.spawn(Panic::Never), ());
    left.insert(c.spawn(Panic::Never), ());
    right.insert(b.spawn(Panic::InDrop), ()); // first duplicate key, dropped during append
    right.insert(c.spawn(Panic::Never), ());

    catch_unwind(move || left.merge(right, |_, _, _| ())).unwrap_err();
    assert_eq!(a.dropped(), 1);
    assert_eq!(b.dropped(), 2);
    assert_eq!(c.dropped(), 2);
}

The assert_eq!(b.dropped(), #) gave me 2 instead of 1 (which is expected behavior it seems).

[asder8215]

I added a panic test within conflict, but I think due to #47949, it isn't incrementing the counter for b_val* and c_val* on drop. The left BTreeMap does have a length of 1 though.

[rustbot]

r? @jhpratt

rustbot has assigned @jhpratt.
They will have a look at your PR within the next two weeks and either review your PR or reassign to another reviewer.

Use r? to explicitly pick a reviewer

Why was this reviewer chosen?

The reviewer was selected based on:

• Owners of files modified in this PR: libs
• libs expanded to 7 candidates
• Random selection from Mark-Simulacrum, jhpratt, joboet

[asder8215]

r? @Mark-Simulacrum

Mark was the reviewer of my previous BTreeMap merge PR (which uses double iterator) and I recall that jhpratt mentioned he's stepping off rotation for a bit last week on the previous BTreeMap PR (not sure if it's true right now), so I'm just rolling this review to Mark.

[Mark-Simulacrum]

If you can squash the commits a little (not sure how much useful is in the history, but seems like there's some that ought to get removed), r=me

[asder8215]

@Mark-Simulacrum squashed the commits into 3 separate commits. Lmk if there's anything else I need to do!

[Mark-Simulacrum]

@bors r+

View changes since this review

[rust-bors]

📌 Commit cbdcfca has been approved by Mark-Simulacrum

It is now in the queue for this repository.

[Kobzol]

@rust-timer build 73143b0

For #153074.

[rust-timer]

Missing artifact for sha 73143b06c7c1a162a7e701b4e0863f3b2179f725 (https://ci-artifacts.rust-lang.org/rustc-builds/73143b06c7c1a162a7e701b4e0863f3b2179f725/rustc-nightly-x86_64-unknown-linux-gnu.tar.xz); not built yet, try again later.

[JonathanBrouwer]

@rust-timer build 73143b0

[rust-timer]

Finished benchmarking commit (73143b0): comparison URL.

Overall result: no relevant changes - no action needed

Benchmarking this pull request means it may be perf-sensitive – we'll automatically label it not fit for rolling up. You can override this, but we strongly advise not to, due to possible changes in compiler perf.

@bors rollup=never
@rustbot label: -S-waiting-on-perf -perf-regression

Instruction count

This benchmark run did not return any relevant results for this metric.

Max RSS (memory usage)

Results (primary 5.7%, secondary -4.0%)

A less reliable metric. May be of interest, but not used to determine the overall result above.

	mean	range	count
Regressions ❌ (primary)	5.7%	[5.7%, 5.7%]	1
Regressions ❌ (secondary)	-	-	0
Improvements ✅ (primary)	-	-	0
Improvements ✅ (secondary)	-4.0%	[-4.0%, -4.0%]	1
All ❌✅ (primary)	5.7%	[5.7%, 5.7%]	1

Cycles

Results (secondary -2.5%)

A less reliable metric. May be of interest, but not used to determine the overall result above.

	mean	range	count
Regressions ❌ (primary)	-	-	0
Regressions ❌ (secondary)	-	-	0
Improvements ✅ (primary)	-	-	0
Improvements ✅ (secondary)	-2.5%	[-2.7%, -2.3%]	2
All ❌✅ (primary)	-	-	0

Binary size

This benchmark run did not return any relevant results for this metric.

Bootstrap: 482.035s -> 480.599s (-0.30%)
Artifact size: 397.81 MiB -> 397.80 MiB (-0.00%)