Speed up creating and extending packed arrays from iterators up to 63×

[ttencate]

This uses the iterator size hint to pre-allocate, which leads to 63× speedup in the best case. If the hint is pessimistic, it reads into a buffer to avoid repeated push() calls, which is still 44x as fast as the previous implementation.

[GodotRust]

API docs are being generated and will be shortly available at: https://godot-rust.github.io/docs/gdext/pr-1023

[Bromeon]

Thanks a lot, this sounds like a great improvement! 🚀

Could you elaborate the role of the intermediate stack buffer? Since it's possible to resize the packed array based on size_hint(), why not do that and write directly from the iterator to self.as_mut_slice()?

Also, ParamType::owned_to_arg() is no longer occurring in the resulting code, is that not necessary for genericity?

[ttencate]

Could you elaborate the role of the intermediate stack buffer? Since it's possible to resize the packed array based on size_hint(), why not do that and write directly from the iterator to self.as_mut_slice()?

That's what the "Fast part" does. The buffer is only needed if there are more items after that.

I guess there might be iterators whose size_hint() gets updated intelligently during iteration, but mostly I'd expect it to decrement down to 0 every time next() is called. So once we have processed size_hint() items, we have no idea how many more are coming, and that's when we proceed to the "Slower part" and start using the buffer.

The alternative (which I implemented initially) is to grow the array in increments of 32 elements, write to as_mut_slice(), and then shrink it again at the end. The problem with that is that you might end up over-allocating, because Godot extends the memory allocation in powers of two. In some cases, this would make the returned array use twice as much memory as in the current implementation – memory that is potentially retained for a very long time. The buffer avoids that nicely, at the expense of some extra copying.

Also, ParamType::owned_to_arg() is no longer occurring in the resulting code, is that not necessary for genericity?

Apparently not. We only implement Extend<$Element>, not something like Extend<E> where E: Into<$Element>.

[Bromeon]

So once we have processed size_hint() items, we have no idea how many more are coming, and that's when we proceed to the "Slower part" and start using the buffer.

If that's the slow part that only happens on "bad" implementations of size_hint(), is the complexity really necessary or can we fall back to the naive approach for remaining elements (less maintenance, lower risks of bugs)?

Do you know how often this occurs in practice?

[ttencate]

There are at least two categories of iterators that are common in the wild, for which we'd want good performance:

1. Exact size is known, e.g. just iterating over a Vec or BTreeSet. This uses the fast part and does nothing on the slower part, since the iterator is finished.
2. Exact size is not known and the lower bound is 0, e.g. Filter, FlatMap, FromFn. This does nothing on the fast part and then finishes the iterator through the slower part. Note that the slower path is only 1.5× slower than the fast path, and still 44× as fast as naive repeated push() calls.

This PR is sufficient to handle them both efficiently. We could eliminate the fast part (case 1) and not lose a lot of performance (maybe incur some memory fragmentation), but that's actually the straightforward and obvious part, so the maintainability gain is small.

This PR also happens to deal efficiently with anything in between, i.e. iterators that report a nonzero lower bound but may return more elements. One example of those would be a Chain of the above two cases.

[Bromeon]

Sounds good, thanks for elaborating! The 2kB buffer (512 ints) is probably also not a big issue, even on mobile/Wasm?

[ttencate]

A cursory search shows stack sizes of at least 1 MB on all platforms. If it becomes a problem after all, it's easy enough to adjust.

[Dheatly23]

If buffer is full, then iterator is advanced but the item is discarded.

Reference: https://doc.rust-lang.org/src/core/iter/adapters/zip.rs.html#165-170

Suggested change

	for (src, dst) in iter::zip(&mut iter, buf.iter_mut().skip(1)) {
	for (dst, src) in iter::zip(buf.iter_mut().skip(1), &mut iter) {

[ttencate]

😱 Yikes, great catch! Maybe this is why I intuitively wrote it in a more explicit way to begin with. iter::zip looks so symmetrical (which is why I prefer it over Iterator::zip) but in this case, that's misleading.

I've updated the test to catch this, and rewrote the loop to be more explicit. The new test also caught another bug that all three of us missed: len += buf_len; was missing at the end of the loop. But I'm confident that it is correct now.

[Dheatly23]

Looks good to me.

Small issue i noticed: if iterator panics, data in buffer will not be dropped. It's not a safety issue, but it would be nice to drop buffer properly.

struct Buffer<const N: usize, T> {
    buf: [MaybeUninit<T>; N],
    len: usize,
}

impl<const N: usize, T> Default for Buffer<N, T> {
    fn default() -> Self {
        Self {
            buf: [const { MaybeUninit::uninit() }; N],
            len: 0,
        }
    }
}

impl<const N: usize, T> Drop for Buffer<N, T> {
    fn drop(&mut self) {
        assert!(self.len <= N);
        if N > 0 {
            unsafe {
                ptr::drop_in_place(ptr::slice_from_raw_parts_mut(
                    self.buf[0].as_mut_ptr(),
                    self.len,
                ));
            }
        }
    }
}

[ttencate]

Great catch yet again. This looks like a good opportunity to make that Buffer into a safe-ish abstraction, which makes the logic in extend() a lot easier to follow as well! PTAL.

[Bromeon]

This PR adds now a significant amount of complexity, so there's quite a chance that we introduce bugs. Given the performance gains, it's probably OK to iron those out over time, but maybe we should bookmark this in case regressions appear in the future 🙂

But I wonder -- are there no higher-level ways to achieve this with the standard library? It seems to be a pattern that occurs every now and then when implementing Extend... Does anyone know how std or other crates handle it, also with low-level unsafe all the time?

[Bromeon]

Could this be simplified with something like

for item in iter.take(N - buf.len()) { ... }

[ttencate]

No, take consumes the iterator by value. Maybe something like (&mut iter).take(N - buf.len()) would work, but I think the version above expresses the intent more clearly and is less prone to bugs.

[ttencate]

This PR adds now a significant amount of complexity, so there's quite a chance that we introduce bugs. Given the performance gains, it's probably OK to iron those out over time, but maybe we should bookmark this in case regressions appear in the future 🙂

It'll appear in the changelog, right?

But I wonder -- are there no higher-level ways to achieve this with the standard library? It seems to be a pattern that occurs every now and then when implementing Extend... Does anyone know how std or other crates handle it, also with low-level unsafe all the time?

Here's the std::Vec implementation: https://doc.rust-lang.org/src/alloc/vec/mod.rs.html#3512-3534 It polls the size_hint() continuously, which makes sense in that implementation.

The problem is that such an implementation would make at least one Godot API call per element, so it wouldn't be any faster than we currently have.

The high-level way would be to collect the iterator into a Vec first, and then memcpying that into the Packed*Array. I haven't benchmarked that strategy, but I think it would be around the same speed as my implementation. However, it requires an allocation the same size as the output, and possibly intermediate allocations as the Vec grows. If you'd rather have that simplicity, I'd be happy to bench it.