Rust iterators and looking forward (peek/multipeek)

Question

I am trying to use a pattern with iterators in Rust and falling down somewhere, apparently simple.

I would like to iterate through a container and find an element with a predicate [A] (simple), but then look forward using another predicate and get that value [B] and use [B] to mutate [A] in some way. In this case [A] is mutable and [B] can be immutable; this makes no difference to me, only to the borrow checker (rightly).

It would help to understand this with a simple scenario, so I have added a small snippet to let folk see the issue/attempted goal. I have played with itertools and breaking into for/while loops, although I want to remain as idiomatic as possible.

Silly Example scenario

Lookup an even number, find next number that is divisible by 3 and add to the initial number.

#[allow(unused)]
fn is_div_3(num: &u8) -> bool {
    num % 3 == 0
}

fn main() {
    let mut data: Vec<u8> = (0..100).collect();

    let count = data.iter_mut()
        .map(|x| {
            if *x % 2 == 0 {
                // loop through numbers forward to next is_div_3,
                // then x = x + that number
            }
            true
        })
        .count();

    println!("data {:?}, count was {} ", data, count);
}

playground

Answer 1

Sadly I'm a bit late, but here goes.

It's not totally pretty, but it's not as bad as the other suggestion:

let mut data: Vec<u8> = (1..100).collect();

{
    let mut mut_items = data.iter_mut();
    while let Some(x) = mut_items.next() {
        if *x % 2 == 0 {
            let slice = mut_items.into_slice();
            *x += *slice.iter().find(|&x| x % 3 == 0).unwrap();
            mut_items = slice.iter_mut();
        }
    }
}

println!("{:?}", data);

gives

[1, 5, 3, 10, 5, 15, 7, 17, 9, 22, ...]

as with Matthieu M.'s solution.

The key is to use mut_items.into_slice() to "reborrow" the iterator, effectively producing a local (and thus safe) clone of the iterator.

Answer 2

Warning: The iterator presented right below is unsafe because it allows one to obtain multiple aliases to a single mutable element; skip to the second part for the corrected version. (It would be alright if the return type contained immutable references).

If you are willing to write your own window iterator, then it becomes quite easy.

First, the iterator in all its gory details:

use std::marker::PhantomData;

struct WindowIterMut<'a, T>
    where T: 'a
{
    begin: *mut T,
    len: usize,
    index: usize,
    _marker: PhantomData<&'a mut [T]>,
}

impl<'a, T> WindowIterMut<'a, T>
    where T: 'a
{
    pub fn new(slice: &'a mut [T]) -> WindowIterMut<'a, T> {
        WindowIterMut {
            begin: slice.as_mut_ptr(),
            len: slice.len(),
            index: 0,
            _marker: PhantomData,
        }
    }
}

impl<'a, T> Iterator for WindowIterMut<'a, T>
    where T: 'a
{
    type Item = (&'a mut [T], &'a mut [T]);

    fn next(&mut self) -> Option<Self::Item> {
        if self.index > self.len { return None; }

        let slice: &'a mut [T] = unsafe {
            std::slice::from_raw_parts_mut(self.begin, self.len)
        };
        let result = slice.split_at_mut(self.index);

        self.index += 1;

        Some(result)
    }
}

Invoked on [1, 2, 3] it will return (&[], &[1, 2, 3]) then (&[1], &[2, 3]), ... until (&[1, 2, 3], &[]). In short, it iterates over all the potential partitions of the slice (without shuffling).

Which is safe to use as:

fn main() {
    let mut data: Vec<u8> = (1..100).collect();

    for (head, tail) in WindowIterMut::new(&mut data) {
        if let Some(element) = head.last_mut() {
            if *element % 2 == 0 {
                if let Some(n3) = tail.iter().filter(|i| *i % 3 == 0).next() {
                    *element += *n3;
                }
            }
        }
    }

    println!("{:?}", data);
}

Unfortunately it can also be used as:

fn main() {
    let mut data: Vec<u8> = (1..100).collect();

    let mut it = WindowIterMut::new(&mut data);
    let first_0 = { it.next(); &mut it.next().unwrap().0[0] };
    let second_0 = &mut it.next().unwrap().0[0];

    println!("{:?} {:?}", first_0 as *const _, second_0 as *const _);
}

which when run print: 0x7f73a8435000 0x7f73a8435000, show-casing that both mutable references alias the same element.

---

Since we cannot get rid of aliasing, we need to get rid of mutability; or at least defer to interior mutability (Cell here since u8 is Copy).

Fortunately, Cell has no runtime cost, but it does cost a bit in ergonomics (all those .get() and .set()).

I take the opportunity to make the iterator slightly more generic too, and rename it since Window is already a used name for a different concept.

struct FingerIter<'a, T>
    where T: 'a
{
    begin: *const T,
    len: usize,
    index: usize,
    _marker: PhantomData<&'a [T]>,
}

impl<'a, T> FingerIter<'a, T>
    where T: 'a
{
    pub fn new(slice: &'a [T]) -> FingerIter<'a, T> {
        FingerIter {
            begin: slice.as_ptr(),
            len: slice.len(),
            index: 0,
            _marker: PhantomData,
        }
    }
}

impl<'a, T> Iterator for FingerIter<'a, T>
    where T: 'a
{
    type Item = (&'a [T], &'a T, &'a [T]);

    fn next(&mut self) -> Option<Self::Item> {
        if self.index >= self.len { return None; }

        let slice: &'a [T] = unsafe {
            std::slice::from_raw_parts(self.begin, self.len)
        };

        self.index += 1;
        let result = slice.split_at(self.index);

        Some((&result.0[0..self.index-1], result.0.last().unwrap(), result.1))
    }
}

We use it as a building brick:

fn main() {
    let data: Vec<Cell<u8>> = (1..100).map(|i| Cell::new(i)).collect();

    for (_, element, tail) in FingerIter::new(&data) {
        if element.get() % 2 == 0 {
            if let Some(n3) = tail.iter().filter(|i| i.get() % 3 == 0).next() {
                element.set(element.get() + n3.get());
            }
        }
    }

    let data: Vec<u8> = data.iter().map(|cell| cell.get()).collect();

    println!("{:?}", data);
}

On the playpen this prints: [1, 5, 3, 10, 5, 15, 7, 17, 9, 22, ...], which seems correct.