Custom cyclic iterator over vector of vectors and mutability woes

#1

Hello there,

I'm still breaking into my Rust shoes and am in need of a bit of advice, as I've painted myself into a classic mutable ref borrowing corner, but can't see my way out of it this time around.

I need to wrap the iteration over a vector of vectors (Vec<WorkGroup>), where the outer vector should be cycled (i.e. restarted from the beginning once a None is yielded), in a single iterator. The item yielded by the iterator (WaveFront) needs to be mutable (the code using the iterator modifies the contents of the vectors, but not the vectors themselves). The iterator's lifetime will always be shorter than the vector of vectors it works on.

The item mutability requirement seems to "creep up" the hierarchy of things, eventually turning almost every reference involved into a mutable one. I tried core::iter::Cycle first, but that clones the original iterator (which is an IterMut). I tried only holding onto indices into the respective vecs, but that gives me the classic mutable-borrow-in-previous-iteration. There are also self reference lifetime shenanigans that prevent actually implementing the Iterator trait.

How should I solve a conundrum like this?

#[derive(Clone)]
struct WaveFront<'a> {
    /// Representative dummy - there is an immutable slice ref here indeed.
    foo: &'a [u32]
}

struct WorkGroup<'a> {
    wavefronts: Vec<WaveFront<'a>>
}

// This is what I'd like to have, but `Cycle` keeps the original iterator
// around and clones it once the current iterator returns `None`, and you can't
// clone a mutable iterator.
/*struct WaveFrontIterator<'a> {
	wg_it: core::iter::Cycle<std::slice::IterMut<'a, WorkGroup<'a>>>,
	wf_it: std::slice::IterMut<'a, WaveFront<'a>>
}

impl<'a> WaveFrontIterator<'a> {
	fn new(wgs: &'a mut Vec<WorkGroup<'a>>) -> Self {
		let wg_it = wgs.iter_mut().cycle();
		let wf_it = wg_it.next().as_mut().unwrap().wavefronts.iter();
		WaveFrontIterator {
			wg_it: wg_it,
			wf_it: wf_it
		}
	}
}

impl<'a> WaveFrontIterator<'a> {
	pub fn next(&mut self) -> Option<&'a mut WaveFront<'a>> {
		match self.wf_it.next() {
			Some(wf) => Some(wf),
			None => {
				self.wf_it = self.wg_it.next().unwrap();
				self.wf_it.next()
			}
		}
	}
}*/

struct WaveFrontIterator<'a> {
	pub wgs: &'a mut Vec<WorkGroup<'a>>,
	pub wg_idx: usize,
	pub wf_idx: usize
}

impl<'a> WaveFrontIterator<'a> {
	fn new(wgs: &'a mut Vec<WorkGroup<'a>>) -> Self {
		WaveFrontIterator {
			wgs: wgs,
			wg_idx: 0usize,
			wf_idx: 0usize
		}
	}
}

// Can't implement this as `trait Iterator` because we place a lifetime constraint on the `self` reference.
impl<'a> /*Iterator for*/ WaveFrontIterator<'a> {
    //type Item = &'a mut WaveFront<'a>;
	pub fn next(&'a mut self) -> Option<&'a mut WaveFront<'a>> {
		if self.wf_idx >= self.wgs[self.wg_idx].wavefronts.len() {
			self.wf_idx = 0;
			self.wg_idx += 1;
			if self.wg_idx >= self.wgs.len() {
				self.wg_idx = 0;
			}
		}
		let wf = &mut self.wgs[self.wg_idx].wavefronts[self.wf_idx];
		self.wf_idx += 1;
		return Some(wf);
	}
}

fn main() {
    static DWORDS: &[u32] = &[0, 1, 2];
    let mut wgs: Vec<WorkGroup> = vec![];
    for _ in 0..32 {
        wgs.push(WorkGroup{ wavefronts: vec![WaveFront{ foo: DWORDS }; 64] });
    }
    let mut wf_it = WaveFrontIterator::new(&mut wgs);
    let mut index = 0usize;
    loop {
        let wf = wf_it.next();
        if wf.is_none() {
            break;
        }
        index += 1;
        println!("Hello wavefront {index}!");
    }
}

(Playground)

Errors:

   Compiling playground v0.0.1 (/playground)
error[E0499]: cannot borrow `wf_it` as mutable more than once at a time
  --> src/main.rs:84:18
   |
84 |         let wf = wf_it.next();
   |                  ^^^^^^^^^^^^ `wf_it` was mutably borrowed here in the previous iteration of the loop

For more information about this error, try `rustc --explain E0499`.
error: could not compile `playground` due to previous error
#2

You need to use a distinct lifetime for self:

pub fn next<'b>(&'b mut self) -> Option<&'b mut WaveFront<'a>>

Because you want to borrow from self, but you don't need the whole duration of 'a, just a short borrow for this item.

As an aside, using the same lifetime for the Vec and WorkGroup (&'a mut Vec<WorkGroup<'a>> looks wrong. Also, I'd use a slice and not a Vec in the iterator.

3 Likes
#3

&'a mut Thing<'a> translates as "exclusively borrow Thing for the rest of it's valid lifetime." Once done, you can never use Thing again, except through the exclusive reference. [1] It has extremely niche utility and is thus an anti-pattern, as @chrefr indicated.

Playground separating the lifetimes.

  1. Including "using" via destructors -- so if you have one, it won't compile. ↩ī¸Ž

4 Likes
#4

Since the Iterator trait won't work for this usecase, you may find it more comfortable to use an interior iteration pattern instead:

impl<'a> WorkGroup<'a> {
    fn process<T,F>(wgs: &mut [Self], mut f:F)->T
    where F: FnMut(&mut WaveFront<'a>)->ControlFlow<T> {
        loop {
            for wg in wgs.iter_mut() {
                for wf in wg.wavefronts.iter_mut() {
                    if let ControlFlow::Break(x) = f(wf) {
                        return x
                    }
                }
            }
        }
    }
}

fn main() {
    static DWORDS: &[u32] = &[0, 1, 2];
    let mut wgs: Vec<WorkGroup> = vec![];
    for _ in 0..32 {
        wgs.push(WorkGroup{ wavefronts: vec![WaveFront{ foo: DWORDS }; 64] });
    }

    let mut index = 0usize;
    
    WorkGroup::process(&mut wgs, |_wf| {
        index += 1;
        println!("Hello wavefront {index}!");
        if index >= 1000 { ControlFlow::Break(()) }
        else { ControlFlow::Continue(()) }
    })
}
4 Likes