How can I make an optimized version of Take for an Iterator?

I have a data structure whose Iterator could implement .take(n) simply by adjusting a number in the Iterator, which would generally be much more efficient than counting down the number of elements returned in a wrapping Iterator. Same for .skip(n). I thought that slice's Iterator must do an optimization like this, but unless I'm misreading the code (which is certainly possible), it just uses the blanket Iterator version of Take that counts down. Does it really use the slow default implementation? Is there some reason why it doesn't provide a custom implementation?

the signature of Iterator::take() must return Take<Self>, there's not much you can do to override it.

It looks like Take, as you said (unless I'm also reading it wrong).

I suppose the object of your iterator could implement fn take(mut self, n: usize) to change the countdown value directly, but I'm not sure the result after optimization will differ much after optimization. Something to try out? It will only affect a take that directly follows your iter/into_iter, of course, and you'll have to implement that method on all your objects (if you have a mutable iterator, immutable, ...). If there are other adaptors before the take, it'll use Take.

EDIT: same for skip.

It's not clear what sort of take() implementation you are envisioning. Are you proposing that changing the length of the slice (to take a subslice) would be better than Take maintaining its own count of whether it has reached the end of the taken portion? That would, theoretically, avoid duplicating the check in Take with a separate bounds check in the iterator, but I think the performance gains, after optimisation, would be relatively small so I suppose it wasn't thought worth the restructure to allow that alternative implementation. I don't think it would be particularly worthwhile for your iterator either, but perhaps you could give more details as to what your iterator does.

For skip() the obvious optimisation is to just not call next() for the skipped elements, but the implementation of Iterator for Skip actually just uses nth() (which the slice iterator does specialise). So the solution for skip() is that the optimisation should go in nth().

1 Like

In release builds the optimizer will inline Take with the underlying slice iterator and the logic will be optimized together. As a result, an explicit manual specialization may not be any faster. If the argument to take is a constant, the output will even be specialized based on that value. Play around on Compiler Explorer with -C opt-level=2 to see.

3 Likes

Thanks, everyone!