ACP: Add `BTreeMap::append_with`

[asder8215]

Proposal

Problem statement

When we need to merge two BTreeMaps into one BTreeMap (preserving Ord behavior) in an efficient manner, we do this through BTreeMap::append (an O(n + m) operation or just linear time). This method is more efficient than grabbing an owned (K, V) pair from other BTreeMap and using BTreeMap::insert to insert it into self BTreeMap, which is an O(mlg(n)) operation (m is the number of (K, V) pairs in other BTreeMap and n is the number of (K, V) pairs in self BTreeMap) .

In using BTreeMap::append though, we opt in for the behavior that:

If a key from other is already present in self, the respective value from self will be overwritten with the respective value from other.

In some cases, however, we want explicit control on the value that the conflicting key should hold.

Motivating examples or use cases

• As mentioned by Mozhewen, Kpreid, and Vorpal here, they have cases where they would desire control on the value of the conflicting key:

I'm currently working on a project that makes heavy use of the BTreeMap of the standard library. One of my needs is to merge two maps into one where values of the same key are combined using a user-specified function f:

// pseudo-code
m1 = {"key1": "foo", "key2": "baz", ...}
m2 = {"key1": "bar", "key3": "dog", ...}

mergeWith(f, m1, m2) = {
  "key1": f("foo", "bar"),
  "key2": "baz",
  "key3": "dog",
  ...
}

- Mozhewen

I needed to know specifically if things existed in the left, right or both map (and then also compare the values to generate diffs of those values describing what needs to be done to go from left to right).

- Vorpal

My use case is much the same, except that my maps are already patches — or rather, “transactions” — and the goal is to produce a single transaction which has the simultaneous effects of both inputs.

- Kpreid (replying to what Vorpal said above)

• Another example of when we may need this is when our values for a conflicting key in two different BTreeMaps represent frequencies. For example, I may have two BTreeMaps contain (String, usize) key-value pairs that I want to combine the two BTreeMap to collect the frequencies of the Strings while preserving alphabetical or reverse-alphabetical order of the keys for various reasons.

Solution sketch

We can introduce a function like BTreeMap::append_with() that takes in an impl FnMut to decide on what value should be held in a conflicting key between two BTreeMaps.

Changes that could be made in rust/library/alloc/src/collections/btree/map.rs (implementation same as BTreeMap::append() except for root.append_from_sorted_iters_with()).

impl <K, V, A: Allocator + Clone> BTreeMap<K, V, A> {
    pub fn append_with(&mut self, other: &mut BTreeMap<K, V, A>, conflict: impl FnMut(&K, V, V) -> V) { 
            // Do we have to append anything at all?
            if other.is_empty() {
                return;
            }

            // We can just swap `self` and `other` if `self` is empty.
            if self.is_empty() {
                mem::swap(self, other);
                return;
           }

           let self_iter = mem::replace(self, Self::new_in((*self.alloc).clone())).into_iter();
           let other_iter = mem::replace(other, Self::new_in((*self.alloc).clone())).into_iter();
           let root = self.root.get_or_insert_with(|| Root::new((*self.alloc).clone()));
           root.append_from_sorted_iters_with(self_iter, other_iter, &mut self.length, (*self.alloc).clone(), conflict)
    };
}

Changes that could be made in rust/library/alloc/src/collections/btree/append.rs (as a consequence of introducing FnMut conflict)

impl<K, V> Root<K, V> {
    pub(super) fn append_from_sorted_iters_with<I, A: Allocator + Clone>(
        &mut self,
        left: I,
        right: I,
        length: &mut usize,
        alloc: A,
        f: impl FnMut(&K, V, V) -> V,
    ) where
        K: Ord,
        I: Iterator<Item = (K, V)> + FusedIterator,
    {
        // We prepare to merge `left` and `right` into a sorted sequence in linear time.
        let iter = MergeIterWith { inner: MergeIterInner::new(left, right), f };

        // Meanwhile, we build a tree from the sorted sequence in linear time.
        self.bulk_push(iter, length, alloc)
    }
}

struct MergeIterWith<F, K, V, I: Iterator<Item = (K, V)>> {
    inner: MergeIterInner<I>,
    f: F,
}

impl<F, K: Ord, V, I> Iterator for MergeIterWith<F, K, V, I>
where
    F: FnMut(&K, V, V) -> V,
    I: Iterator<Item = (K, V)> + FusedIterator,
{
    type Item = (K, V);

    /// If two keys are equal, returns the key from the left and uses `f` to return
    /// a value given knowledge of conflicting key and values from left and right
    fn next(&mut self) -> Option<(K, V)> {
        let (a_next, b_next) = self.inner.nexts(|a: &(K, V), b: &(K, V)| K::cmp(&a.0, &b.0));
        match (a_next, b_next) {
            (Some((a_k, a_v)), Some((_, b_v))) => Some({
                let next_val = (self.f)(&a_k, a_v, b_v);
                (a_k, next_val)
            }),
            (Some(a), None) => Some(a),
            (None, Some(b)) => Some(b),
            (None, None) => None,
        }
    }
}

You can see a solution implementation of what BTreeMap::append_with() could look like in this PR I made.

Alternatives

Alternatives to doing something akin to BTreeMap::append_with() with existing methods could be to use a combination of an owned iterator of other BTreeMap, and using BTreeMap::get and BTreeMap::insert on self BTreeMap to decide on what value the conflicting key should hold:

use std::collections::BTreeMap;
    
let mut a = BTreeMap::new();
a.insert(1, String::from("a"));
a.insert(2, String::from("b"));
a.insert(3, String::from("c")); // Note: Key (3) also present in b.
    
let mut b = BTreeMap::new();
b.insert(3, String::from("d")); // Note: Key (3) also present in a.
b.insert(4, String::from("e"));
b.insert(5, String::from("f"));

// concatenate a's value and b's value
for (b_key, b_val) in b { // takes O(m)
    let a_val = a.get(&b_key); // takes O(lg(n))
    // below takes O(lg(n))
    if let Some(a_val) = a_val {
        // concatenation occurs here
        a.insert(b_key, format!("{a_val}{b_val}"));
    } else {
        a.insert(b_key, b_val);
    }
}

Doing this however leads to a time complexity of O(m(2lg(n))), where m is the number of (K, V) pairs in other BTreeMap and n is the number of (K, V) pairs in self BTreeMap. This would be a suboptimal approach to an append_with() equivalent.

We can also consider alternate impls of FnMut for the solution sketch I made like:

• impl FnMut(K, &mut V, V) as suggested by Simonas in here
• impl FnMut(V, V) -> V
• impl FnMut(&mut V, V) -> V

The last two impl FnMut provides no context to the user on what the conflicting key is and only gives them access to self's and other's V. In some cases, I think having context about the conflicting key may be desirable and would warrant different behaviors on what V should be for that key, so I would not opt for this callback. As for impl FnMut(K, &mut V, V) vs. impl FnMut(&K, V, V) -> V, I mentioned this in my solution implementation PR:

I think having the conflict function own self's value and returning an owned value might give a lot more flexibility to the user. It also feels a bit more intuitive to me as the user that a value of type V is being put into self when there's conflicting keys between self and other.

For example, if my values are String types and I want to combine my values when keys are conflicting, then they can either mutate selfs value with .push_str() passing in other's value and return self's value, or they can use format!() to return a concatenation of self's and other's value.

Now, a consequence for this approach is that it's very well possible that the user decides to return a value V that has nothing to do with self or other when it comes to a conflicting key. However, should we decide to do conflict: impl FnMut(&K, &mut V, V), I could have mutated self's value easily to produce a whole different value that has nothing to do with the original value of self or other.

Links and related work

• BTreeMap::append_with that allows custom value merges rust#147700
• https://internals.rust-lang.org/t/btreemap-methods-insert-append-to-combine-values-instead-of-overwriting/21353/4

What happens now?

This issue contains an API change proposal (or ACP) and is part of the libs-api team feature lifecycle. Once this issue is filed, the libs-api team will review open proposals as capability becomes available. Current response times do not have a clear estimate, but may be up to several months.

Possible responses

The libs team may respond in various different ways. First, the team will consider the problem (this doesn't require any concrete solution or alternatives to have been proposed):

• We think this problem seems worth solving, and the standard library might be the right place to solve it.
• We think that this probably doesn't belong in the standard library.

Second, if there's a concrete solution:

• We think this specific solution looks roughly right, approved, you or someone else should implement this. (Further review will still happen on the subsequent implementation PR.)
• We're not sure this is the right solution, and the alternatives or other materials don't give us enough information to be sure about that. Here are some questions we have that aren't answered, or rough ideas about alternatives we'd want to see discussed.

[programmerjake]

imo the signature for the merge conflict callback should be FnMut(K, V, V) -> V since you have 2 keys so you can pass ownership of one of them, or maybe even FnMut(&K, V, K, V) -> V since keys can compare equal but have different internal state.

[pitaj]

FnMut(&K, V, K, V) -> V since keys can compare equal but have different internal state.

Or even FnMut((K, V), (K, V)) -> (K, V) so you can pick both the key and value you want.

[programmerjake]

Or even FnMut((K, V), (K, V)) -> (K, V) so you can pick both the key and value you want.

but then append_with has to handle what happens when you return a unrelated key that conflicts with some other key, removing the advantage of knowing that the keys are in sorted order and have no duplicates.

[scottmcm]

so you can pick both the key and value you want.

The worry I'd have with anything returning the key is that now it's possible for me to violate the "it's still sorted" criterion, which presumably would require pessimizing the implementation since even for trusted types it'd need to check that I didn't just return rand().

Taking impl FnMut(&K, &mut V, (K, V)), maybe, having the closure update the existing value?

[asder8215]

imo the signature for the merge conflict callback should be FnMut(K, V, V) -> V since you have 2 keys so you can pass ownership of one of them, or maybe even FnMut(&K, V, K, V) -> V since keys can compare equal but have different internal state.

I'm okay with using FnMut(K, V, V) -> V or FnMut(&K, V, V) -> V; if I did the former, I think it would make sense to pass ownership of other's key because it is other BTreeMap being appended onto self BTreeMap (the K portion of the returned (K, V) in MergeIterWith's next() should come from self BTreeMap).

I can see your point in using FnMut(&K, V, K, V) -> V since K::cmp(a, b) can return Ordering::Equal based on what's being compared even if both a and b have other differing fields. I'm really curious on when someone would want to implement comparing their keys in this manner; I would imagine that people would put anything that differs in the value of the conflicting key. It also brings up the question on if we go for this approach, are we actually having keys that 'conflict' (if they are not duplicates)?

Or even FnMut((K, V), (K, V)) -> (K, V) so you can pick both the key and value you want.

I disagree with using FnMut((K, V), (K, V)) -> (K, V). This would give the user an opportunity to return a K that could be completely unrelated to the conflicting key, which like programmerjake and scottmcm pointed out removes the benefit that self BTreeMap's keys are in sorted order.

Taking impl FnMut(&K, &mut V, (K, V)), maybe, having the closure update the existing value?

I think using a tuple to contain other BTreeMap conflicting key-val pairs is clean, which going back to programmerjake's suggestion on FnMut(&K, V, K, V) -> V, it might be neater to use FnMut(&K, V, (K, V)) -> V? I still feel like it would be more intuitive for the callback to pass in V rather than &mut V from self and return a V as per my reasoning for impl FnMut(K, &mut V, V) vs impl FnMut(&K, V, V) -> V.

I think an important thing to note is that if we do go for something like FnMut(K, V, V) -> V or FnMut(&K, V, V) -> V (or variants of using one K for FnMut), it would need to be documented that the conflicting keys from self BTreeMap and other BTreeMap does not necessarily mean the keys are duplicate and that self BTreeMap preserves its own key in this case (this could definitely be worded better).

[kpreid]

Regarding the signature of the conflict function, it is important that both Vs are passed in by value (not &mut V). This is because it may be the case (and is, for my use case) that constructing the combined V value requires, or greatly benefits from, ownership of both original Vs.

---

What is the actual time complexity of this operation going to be? The first paragraph says that it is O(n + m), but that is worse than O(n lg(m)) which you can achieve through individual insertions. In particular, is the complexity the same when n and m are swapped? If it is not, it may be desirable to automatically perform swaps of the two inputs.

[programmerjake]

if the inputs are automatically swapped, I think it's important that they are un-swapped again when passed to the conflict function, so code can rely on knowing which value comes from which input.

[the8472]

Deconstructing both maps and then constructing a new one is far from ideal when n >> m or m >> n. Iterating over the smaller map and doing targeted merges/insertions would be more efficient in that scenario.

Additionally, if the fraction of conflicting values is ~1 then you're paying node deallocation/allocation overhead instead of doing the merge in-place.

So an optimal implementation would be more complicated than proposed, but that doesn't affect the necessary signatures.

[asder8215]

What is the actual time complexity of this operation going to be? The first paragraph says that it is O(n + m), but that is worse than O(n lg(m)) which you can achieve through individual insertions. In particular, is the complexity the same when n and m are swapped? If it is not, it may be desirable to automatically perform swaps of the two inputs.

The worst case time complexity for BTreeMap::append_with will be O(n+m) if our self BTreeMap and other BTreeMap have completely different keys. Like BTreeMap::append, BTreeMap::append_with takes the union of two sorted iterators using a two pointer approach. You can see this from BTreeMap::append:

// An iterator for merging two sorted sequences into one
struct MergeIter<K, V, I: Iterator<Item = (K, V)>>(MergeIterInner<I>);

impl<K: Ord, V, I> Iterator for MergeIter<K, V, I>
where
    I: Iterator<Item = (K, V)> + FusedIterator,
{
    type Item = (K, V);

    /// If two keys are equal, returns the key from the left and the value from the right.
    fn next(&mut self) -> Option<(K, V)> {
        let (a_next, b_next) = self.0.nexts(|a: &(K, V), b: &(K, V)| K::cmp(&a.0, &b.0));
        match (a_next, b_next) {
            (Some((a_k, _)), Some((_, b_v))) => Some((a_k, b_v)),
            (Some(a), None) => Some(a),
            (None, Some(b)) => Some(b),
            (None, None) => None,
        }
    }
}

impl<I: Iterator> MergeIterInner<I> {
/// Creates a new core for an iterator merging a pair of sources.
    pub(super) fn new(a: I, b: I) -> Self {
        MergeIterInner { a, b, peeked: None }
    }

    /// Returns the next pair of items stemming from the pair of sources
    /// being merged. If both returned options contain a value, that value
    /// is equal and occurs in both sources. If one of the returned options
    /// contains a value, that value doesn't occur in the other source (or
    /// the sources are not strictly ascending). If neither returned option
    /// contains a value, iteration has finished and subsequent calls will
    /// return the same empty pair.
    pub(super) fn nexts<Cmp: Fn(&I::Item, &I::Item) -> Ordering>(
        &mut self,
        cmp: Cmp,
    ) -> (Option<I::Item>, Option<I::Item>)
    where
        I: FusedIterator,
    {
        let mut a_next;
        let mut b_next;
        match self.peeked.take() {
            Some(Peeked::A(next)) => {
                a_next = Some(next);
                b_next = self.b.next();
            }
            Some(Peeked::B(next)) => {
                b_next = Some(next);
                a_next = self.a.next();
            }
            None => {
                a_next = self.a.next();
                b_next = self.b.next();
            }
        }
        if let (Some(a1), Some(b1)) = (&a_next, &b_next) {
            match cmp(a1, b1) {
                Ordering::Less => self.peeked = b_next.take().map(Peeked::B),
                Ordering::Greater => self.peeked = a_next.take().map(Peeked::A),
                Ordering::Equal => (),
            }
        }
        (a_next, b_next)
    }

    /// Returns a pair of upper bounds for the `size_hint` of the final iterator.
    pub(super) fn lens(&self) -> (usize, usize)
    where
        I: ExactSizeIterator,
    {
        match self.peeked {
            Some(Peeked::A(_)) => (1 + self.a.len(), self.b.len()),
            Some(Peeked::B(_)) => (self.a.len(), 1 + self.b.len()),
            _ => (self.a.len(), self.b.len()),
        }
    }
}

The only difference for BTreeMap::append_with is that it uses the conflict function to handle the case that (a_next, b_next) = (Some(K_a, V_a), Some(K_b, V_b). Swapping whether we are appending self BTreeMap onto other BTreeMap or vice versa won't change this worst case time complexity (or the time complexity in general since we end up iterating through each keys from self BTreeMap and other BTreeMap regardless).

If you do the BTreeMap::get and BTreeMap::insert() approach, your time complexity could be O(n * lg(m)) or O(m * lg(n)) depending on the whether you're inserting all elements from the first BTreeMap onto the second BTreeMap or the other way around.

And you're right that O(n * lg(m)) or O(m * lg(n)) could be more efficient than O(n + m) in some cases, but that occurs when m and n are small values or m is significantly smaller than n or n is significantly smaller than m. If m = n, O(n + m) would be better almost always.

Here's a 3D graph of O(n * lg(m)) (in purple) vs O(m * lg(n)) (in magenta) vs O(n + m) (in red)

In terms of gunning for maximum efficiency, we could fallback to using BTreeMap::insert and BTreeMap::get approach if O(n * lg(m)) / O(m * lg(n)) < O(n + m). However, this could be out of scope with what BTreeMap::append_with is trying to accomplish in being something similar to extend() but with a callback function (and we would probably need to change BTreeMap::append to do something similar in a separate PR). It could just be the user themselves could go out their way to implement a maximally efficient approach to combining BTreeMaps provided these methods.

[Amanieu]

We discussed this in the @rust-lang/libs-api meeting and are happy to accept this proposal. Some notes from the meeting:

• We prefer the name merge instead of append_with since this effectively provides a conflict resolution function during the merge.
• We would like other to be taken by value instead of by &mut. We generally consider the use of &mut in append for the second argument to be a mistake.
• It's fine to keep this as taking &mut self.
• It's fine to keep the signature of the conflict function as FnMut(&K, V, V).

Please open a tracking issue and open a PR to rust-lang/rust to add it as an unstable feature. You can close this ACP once the tracking issue has been created.

[asder8215]

It's fine to keep the signature of the conflict function as FnMut(&K, V, V).

Just to clarify, did you mean FnMut(&K, V, V) -> V or FnMut(&K, V, V)?

[Amanieu]

It's fine to keep the signature of the conflict function as FnMut(&K, V, V).

Just to clarify, did you mean FnMut(&K, V, V) -> V or FnMut(&K, V, V)?

FnMut(&K, V, V) -> V