Add Extend::{extend_one,extend_reserve}

[cuviper]

This adds new optional methods on Extend: extend_one add a single
element to the collection, and extend_reserve pre-allocates space for
the predicted number of incoming elements. These are used in Iterator
for partition and unzip as they shuffle elements one-at-a-time into
their respective collections.

[rust-highfive]

r? @LukasKalbertodt

(rust_highfive has picked a reviewer for you, use r? to override)

[cuviper]

This was developed from the discussion on #72085, overlapping but independent of #72159. I showed that it achieved similar performance to an imperative unzip in #72085 (comment).

[Dylan-DPC-zz]

r? @dtolnay

[Mark-Simulacrum]

Bikeshed: push instead of extend_one?

Are we using size hints already to pre-reserve? I would've expected an extend(Some(foo)) to basically compile down to a push - if that doesn't happen today, perhaps worth investigating?

I think the only reason we have the somewhat awkward method names is also because Extend is im the prelude, right? So we don't want to use e.g. push and reserve?

[cuviper]

Are we using size hints already to pre-reserve?

There's no way to do that in partition and unzip as-is, since they only have Default + Extend to work with. So that new pre-allocation does play a big part in the performance gain.

I would've expected an extend(Some(foo)) to basically compile down to a push - if that doesn't happen today, perhaps worth investigating?

Here's a playground of the benchmarks from #72085, and my local results:

test bench_functional        ... bench:       1,292 ns/iter (+/- 31)
test bench_imperative_extend ... bench:         803 ns/iter (+/- 8)
test bench_imperative_push   ... bench:         410 ns/iter (+/- 6)

Note that functional is the status quo unzip, whereas this PR gets it down to push performance.

Looking at the assembly, the biggest difference I see is that extend(Some(_)) has an unconditional call to reserve, I believe from here, whereas push only calls that when it actually needs more capacity. Maybe we could branch like that in extend too, but I think that would pessimize the general case, compared to partition and unzip that are hammering extend in a loop.

That's particular to Vec -- other containers may have their own trade-offs. Note that I did make the default extend_one just call extend(Some(_)) anyway.

I think the only reason we have the somewhat awkward method names is also because Extend is im the prelude, right? So we don't want to use e.g. push and reserve?

Yeah, I think being in the prelude warrants caution in method ambiguity.

[Mark-Simulacrum]

I meant reserving in the Extend impls - those have an iterator to work with, so I would've expected it to be feasible there. But if there's an unconditional reserve already then maybe that is covered.

I want to take a closer look at the diff, but I think in principle this is reasonable, so I'll probably approve it (we can bikeshed names on a tracking issue).

[cuviper]

Oh sure, in general, collections will reserve space in extend. But with unzip, they don't see the whole iterator, just one item at a time.

[Mark-Simulacrum]

Okay, I haven't found the time to do the investigation myself I wanted to, but this seems like a good change regardless. Given the amount of code using extend (and zip etc) which is affected by this, I would like to @bors try @rust-timer queue though just to make sure we're not unduly generating more LLVM IR or something like that causing compile times to increase.

Also, can someone on @rust-lang/libs sign off on adding these APIs (unstable for now)? Presuming we get a go-ahead can you file a tracking issue, @cuviper and update the attributes?

trait Extend {
    /// Extends a collection with exactly one element.
    #[unstable(feature = "extend_one", issue = "none")]
    fn extend_one(&mut self, item: A) {
        self.extend(Some(item));
    }

    /// Reserves capacity in a collection for the given number of additional elements.
    ///
    /// The default implementation does nothing.
    #[unstable(feature = "extend_one", issue = "none")]
    fn extend_reserve(&mut self, _additional: usize) {}
}

[rust-timer]

Awaiting bors try build completion

[bors]

⌛ Trying commit 36b7b8ee5ce9321f60e21c60f2f85d76cd0b1a0a with merge d448c21d56eddbda45277b26ed4f17faf4c5db3c...

[dtolnay]

Looks good to me to land unstable but I haven't gotten to review the whole PR yet, only collect.rs.

[Mark-Simulacrum]

r=me with tracking issue assigned

Implementation looks good to me, at least for an initial draft.

[SimonSapin]

These methods look very out of place in the Extend trait, and the extend_ prefix in their name is apparently unrelated to the meaning of the name of the extend method.

From an API design perspective, this seems to be asking for new Push and Reserve traits. This might deserve to be part of a larger overall design for traits for generic collections. The standard library has mostly avoided doing this so far, if fact Extend is the odd one out and arguably should not have been a trait.

[cuviper]

@SimonSapin for the goal of improving Iterator::unzip, Extend is really the only opportunity.

[Mark-Simulacrum]

I agree that they feel a bit out of place, although they do relatively naturally coincide with extending IMO.

I don't think separate traits make a lot of sense, unfortunately, because we lack specialization (and aren't getting it anytime soon I imagine in flexible enough form). You really want to be able to call e.g. push and have that fallback to just extending if there's no separate push implementation defined.

I would, though, be fine saying that these are internal methods not currently on a (clear) path to stabilization (and perhaps marking them as doc(hidden)).

[cuviper]

I've made the requested changes and filed tracking issue #72631. I guess we should at least wait for that perf run before merging though.

I would, though, be fine saying that these are internal methods not currently on a (clear) path to stabilization (and perhaps marking them as doc(hidden)).

It would be a wart of its own to put third-party collections at a disadvantage. There's nothing technically special here to warrant permanent unstability, versus actual specialization. I think this is a relatively mild API choice, just with a bit of historical baggage that unzip only has access to Extend.

[cuviper]

I'll see if I can reproduce the max-rss locally and track that down. Most of those extend_reserve calls should be perfectly safe though, as long as iterators are accurate in their lower bound, especially since RSS reflects memory that you actually touched.

The one hunch I have is with HashMap/Set, since that uses a heuristic that could allocate too much if a lot of keys overlap. Since items are ~randomly dispersed, it would still show a lot in RSS even if sparsely filled.

I guess it could also be that the new methods just mean more code, but I wouldn't think dyn Extend is a common thing, so otherwise it should only pay for what you use. I also expect the paths that do use these new methods should be simpler than before.

[cuviper]

I blame jemalloc for the RSS variability, and I don't think my change really affected it. Here are my results on the worst in that report, keccak-debug full, with three runs each.

Test	1	2	3
before with jemalloc	1344952	1290632	1106140
before with system allocator	900004	900660	889752
after with jemalloc	1304268	1395768	1141788
after with system allocator	889432	889284	889288

(measuring /bin/time -> maxresident k; system is Fedora 32 x86_64, glibc 2.31)

[Mark-Simulacrum]

Hm, I wonder if we can configure jemalloc to be less variable or perhaps disable it on rust-timer entirely. cc @eddyb / @nnethercote

Regardless seems like perf is fine -- @bors r+

I think we can discuss @SimonSapin's concern (#72162 (comment)) more on the tracking issue if the responses to the original comment didn't resolve it, given that this is just adding some unstable functionality and can easily be reverted if needed.

[bors]

📌 Commit 02226e0 has been approved by Mark-Simulacrum

[nnethercote]

To understand the memory impact of changes I strongly recommend using either Massif or DHAT, if you are on Linux. Basic docs are here. If you try DHAT you need to look at the t-gmax numbers, which represent the global heap peak.

[cuviper]

@nnethercote I just compiled keccak-debug full under massif. Both before and after this change, the peak memory consumption of the rustc process was 786.8 MiB. The new methods here aren't seen in any of the heap snapshots.

[nnethercote]

@cuviper: that's a very strong indication that it's just noise, then. Indeed, https://perf.rust-lang.org/?start=&end=&absolute=true&stat=max-rss shows that keccak's max-rss numbers are extremely noisy.

[RalfJung]

@bors r-
Failed to build in #72732 (comment), likely needs rebasing.

[cuviper]

Rebased.

@bors r=Mark-Simulacrum

[bors]

📌 Commit a51b22a has been approved by Mark-Simulacrum