Is there some kind of `ExactSizeIterator::collect_backward()`?

#1

TL; DR: does std or any crate (like itertool) have a collect_backward() for ExactSizeIterator?

I have a piece of code that look more or less like:

let it: impl ExactSizeIterator<Item=T> = ...;
let mut my_vec: Vec<T> = it.collect();
my_vec.reverse();

This could be desugared to:

let mut my_vec: Vec<T> = Vec::with_capacity(it.len());
while let Some(elem) = it.next() {
    my_vec.push(elem); // insert elements from the front
}
my_vec.reverse();

Would it be possible to directly create the vector by inserting the elements from the end. Somethings like an hypothetical.

let my_vec = it.collect_backward();

Since it is ExactSizeIterator, it is possible to create the Vec with the right size, and then use indices to fill directly the right cell.

let mut my_vec = Vec::with_capacity(it.len());
my_vec.resize(it.len());
for i in (0..it.len()).rev {
    my_vec[i] = it.next(); // insert elements from the back
}
assert_eq!(it.next(), None);
#2

If you have a DoubleEndedIterator, you can call rev().collect().

1 Like
#3

It's not a DoubleEndedIterator, only an ExactSize.

I'm iterating on a graph using dijkstra, (so I know how much nodes I crossed), then I walk backward to create the path (a Vec<Node>). I know the size of the iterator, but I can't walk in both directions.

#4

There is no easy way to populate a Vec starting from the back. You could just populate your data in a VecDeque which allows efficient push_front() operation. If you need to have a Vec, you can use into() to convert from VecDeque to Vec. This conversion may copy some memory, but won't do any additional allocations.

If you need to have maximum performance, you can write a function that uses unsafe to populate the Vec from an iterator backwards.

#5

That's unfortunate (but understandable given how niche this is). I think I will just use the implementation I wrote in the OP.

#6

Minor code review nit: if you're doing this, you can just do something like

let mut my_vec = vec![0; it.len()];

(If you do that with 0 specifically, it will use alloc_zeroed internally for a small speedup over different initializers.)

2 Likes
#9

Good call. I also didn't realized when I wrote the code that you can't resize without giving a value in Rust. It make sense, but it's quite anoying in this specific case (I will take a look at MaybeUninit and friends).

#10

Definitely consider whether you can just use a VecDeque instead, first -- that would be fully safe, and in your use case you'd still be able to get a single slice out of it.

If you do need to go unsafe, you might be interested in https://doc.rust-lang.org/nightly/std/vec/struct.Vec.html#method.spare_capacity_mut

5 Likes
#11

It was effectively a good call:

impl<I: Iterator> IteratorUtils for I {}
pub trait IteratorExactUtils: Iterator {
    /// Create a `Vec` by inserting the elements from the end
    fn collect_backward(self) -> Vec<<Self as Iterator>::Item>;
}
impl <I: Iterator + ExactSizeIterator> IteratorExactUtils for I { 
    fn collect_backward(self) -> Vec<<Self as Iterator>::Item> {
        let mut deque = std::collections::VecDeque::with_capacity(self.len());
        for element in self {
            deque.push_front(element)
        }   
        deque.into()
    }
}
#12

Actually the performance of Deque is terrible compared to just reversing a vector:

fn collect_backward(self) -> Vec<<Self as Iterator>::Item> {
    let mut vec = Vec::with_capacity(self.len());
    for element in self {
        vec.push(element)
    }
    vec.reverse();
    vec
}

With vec.reverse(), my code takes 4s per iterations (there is quite a lot of computation in the whole algorithm), but with deque, it takes 40s per iterations, which it 10× slower at the program level.

#13

Building a VecDeque like that leaves an empty space -- if I understand VecDeque correctly -- right at the front of the vector. So when you call into() to convert it into a vector, all the elements have to be moved one space left. (Taking it as a contiguous slice doesn't have this problem because the slice doesn't have to start at the front of the buffer.) So it doesn't necessarily surprise me that reverse on a Vec is faster, although a factor of 10 is pretty big. Maybe something else is going on there.

There might be some way you could fake out VecDeque so the empty space ends up at the end ("back") of the vector and nothing needs to be moved. If you push a dummy value first and then pop it off before pushing the last thing on, maybe?

When I was faced with a similar problem, my solution was to just redefine the Vec so that it was always backwards :grin: I ended up doing a lot less reversals that way and I just had to iterate the Vec in reverse instead (which is super easy and roughly as performant as doing it forward). I don't know if that approach would work for you, but maybe food for thought.

#14

When I was faced with a similar problem, my solution was to just redefine the Vec so that it was always backwards

I kind of did the opposite. Sometime, it's forward, and sometime it's backward, so I had a struct that was wrapping a Vec with a bool to know the direction. But having that boolean propagated in 4 different data structure was anoying, so I decided to "just" construct the vector backward, and get rid of the boolean :wink:

#15

Ah, that is more challenging.

Here's an idea for a type that can be built from a Vec or a reversed Vec but is always iterated in the "forward" direction. Again, maybe not applicable to your situation, but perhaps it could be useful nevertheless.

#16

It doesn't help since I only have an iterator not a DoubleEndedIterator. And I was aware of this technique, it's definitively useful when you want to create an iterator from different one.

#17

Can you show this in a criterion benchmark we can take a look at?

The difference between them should be a reverse vs a memmove, which ought to differ only in constant factors. And if anything, I'd expect the reverse to be slower, since it's awkward to vectorize -- it's particularly slow for 3-byte types, for example.

#18

I am not sure of what information you want to extract from a micro-benchmark. I'm not sure I can correctly extract the same workload I'm doing in a micro-benchmark. My iterator is a bit complicated to implement, and probably slow (each call to next() does an access inside a hash table) so I can't easily emulate it. The element returned by the iterator however is a single usize.

#19

I think this experiment confirms my intuition that the compiler is able to optimize away the my_vec.reverse().

I used the nightly-2020-08-14-x86_64-unknown-linux-gnu version of the compiler, and got exactly the same execution time when using vec.spare_capacity_mut() version vec.reverse()

fn collect_backward(self) -> Vec<<Self as Iterator>::Item> {
    let len = self.len();

    // create the Vec, with the right size
    let mut vec = Vec::with_capacity(self.len());
    let uninit = vec.spare_capacity_mut();

    // Add the elements, starting from the end
    for (i, element) in self.enumerate() {
        uninit[len - i - 1].write(element);
    }

    // Mark the elements of the vector as being initialized.
    unsafe {
        vec.set_len(len);
    }

    vec
}
1 Like
#20

It's possible to have collect_backward for Vec, albeit it's non-trivial, and using VecDeque is much easier. Standard library doesn't provide such a method because it's kinda unusual to need that.

fn collect_backward<I>(iter: I) -> Vec<I::Item>
where
    I: IntoIterator,
    I::IntoIter: ExactSizeIterator,
{
    struct VecState<T> {
        position: usize,
        vec: Vec<T>,
    }
    impl<T> Drop for VecState<T> {
        fn drop(&mut self) {
            // Vec::with_capacity allocation is exact, so we can depend on self.vec.capacity()
            for i in self.position..self.vec.capacity() {
                // This is always initialized due to loop invariant
                unsafe { self.vec.as_mut_ptr().add(i).drop_in_place() }
            }
        }
    }
    let mut iter = iter.into_iter();
    // Query len once to avoid issues with iterators randomizing their length.
    let len = iter.len();
    let mut vec_state = VecState {
        position: len,
        vec: Vec::<I::Item>::with_capacity(len),
    };
    // Loop invariant: vec from vec_state.position to its end is initialized
    for i in (0..len).rev() {
        let elem = iter.next();
        match elem {
            None => panic!("ExactSizeIterator implementation is incorrect"),
            // Safe as i is always less than len (the capacity the vector was
            // allocated with).
            Some(elem) => unsafe { vec_state.vec.as_mut_ptr().add(i).write(elem) },
        }
        // Safe as this position was just written to (if we did follow None branch
        // then we wouldn't reach this code due to panic).
        vec_state.position = i;
    }
    // This function must not panic from this point on.
    // Safe as mem::take for Vec<T> cannot panic.
    let mut vec = mem::take(&mut vec_state.vec);
    // Safe due to loop invariant. At this point vec_state.position must be 0 which
    // means that the range from 0 to len was initialized which is what is required
    // by set_len. set_len cannot panic.
    unsafe { vec.set_len(len) }
    // mem::forget cannot panic.
    mem::forget(vec_state);
    vec
}
2 Likes