Iterator over elements around specific index in Vec<Vec<Object>> - rust

Question

I have a grid: Vec<Vec<Object>> and a pair of x/y indices. I want to find all the elements surrounding the one indexed.
Unfortunately, I can't simply loop over the elements because that ends up borrowing the Vec twice and the borrow checker screams at me:
let mut cells = Vec::with_capacity(8);
for cx in xstart..xend {
for cy in ystart..yend {
if cx != x || cy != y {
cells.push(&mut squares[cy as usize][cx as usize]);
}
}
}
cells.into_iter()
My best attempt at changing this into an iterator chain has also failed spectacularly:
let xstart = if x == 0 { x } else { x - 1 };
let xlen = if x + 2 > squares[0].len() { x + 1 } else { 3 };
let ystart = if y == 0 { y } else { y - 1 };
let ylen = if y + 2 > squares.len() { y + 1 } else { 3 };
let xrel = x - xstart;
let yrel = y - ystart;
squares.iter().enumerate()
.skip(ystart).take(ylen).flat_map(|(i, ref row)|
row.iter().enumerate()
.skip(xstart).take(xlen).filter(|&(j, &c)| i != yrel || j != xrel))
Does anyone know how I can do this?

Answer 1

Personally, I am not sure I would be comfortable working with an iterator when the relative positions of the elements can be important. Instead, I would seek to create a "view" of those elements.
The gist can be found here, but the idea is simple so here are the core structures.
#[derive(Debug)]
struct NeighbourhoodRow<'a, T>
where T: 'a
{
pub left : Option<&'a mut T>,
pub center : Option<&'a mut T>,
pub right : Option<&'a mut T>,
}
#[derive(Debug)]
struct Neighbourhood<'a, T>
where T: 'a
{
pub top : NeighbourhoodRow<'a, T>,
pub center : NeighbourhoodRow<'a, T>,
pub bottom : NeighbourhoodRow<'a, T>,
}
To build them, I use a healthy dose of split_at_mut:
fn take_centered_trio<'a, T>(row: &'a mut [T], x: usize) ->
(Option<&'a mut T>, Option<&'a mut T>, Option<&'a mut T>)
{
fn extract<'a, T>(row: &'a mut [T], x: usize) -> (Option<&'a mut T>, &'a mut [T]) {
if x+1 > row.len() {
(None, row)
} else {
let (h, t) = row.split_at_mut(x+1);
(Some(&mut h[x]), t)
}
}
let (prev, row) = if x > 0 { extract(row, x-1) } else { (None, row) };
let (elem, row) = extract(row, 0);
let (next, _ ) = extract(row, 0);
(prev, elem, next)
}
and the rest is just some uninteresting constructors.
Of course, you can then build some kind of iterator over those.

Answer 2

In the end I made a custom iterator with the help of the guys in #rust
I've typed my struct out to give you the actual code. As pointed out by the guys in #rust you cannot return &mut safely from an iterator without using a different iterator that uses unsafe anyway, and given that the math here is simple enough to ensure it doesn't go wrong an unsafe was the way to go.
type FieldSquare = u8;
use std::iter::Iterator;
pub struct SurroundingSquaresIter<'a> {
squares: &'a mut Vec<Vec<FieldSquare>>,
center_x: usize,
center_y: usize,
current_x: usize,
current_y: usize,
}
pub trait HasSurroundedSquares<'a> {
fn surrounding_squares(&'a mut self, x: usize, y:usize) -> SurroundingSquaresIter<'a>;
}
impl<'a> HasSurroundedSquares<'a> for Vec<Vec<FieldSquare>> {
fn surrounding_squares(&'a mut self, x: usize, y:usize) -> SurroundingSquaresIter<'a> {
SurroundingSquaresIter {
squares: self,
center_x: x,
center_y: y,
current_x: if x == 0 { x } else { x - 1 },
current_y: if y == 0 { y } else { y - 1 },
}
}
}
impl<'a> Iterator for SurroundingSquaresIter<'a> {
type Item = &'a mut FieldSquare;
fn next(&mut self) -> Option<&'a mut FieldSquare> {
if self.current_y + 1 > self.squares.len() || self.current_y > self.center_y + 1 {
return None;
}
let ret_x = self.current_x;
let ret_y = self.current_y;
if self.current_x < self.center_x + 1 && self.current_x + 1 < self.squares[self.current_y].len() {
self.current_x += 1;
}
else {
self.current_x = if self.center_x == 0 { self.center_x } else { self.center_x - 1 };
self.current_y += 1;
}
if ret_x == self.center_x && ret_y == self.center_y {
return self.next();
}
Some(unsafe { &mut *(&mut self.squares[ret_y][ret_x] as *mut _) })
}
}

Answer 3

You want to get mutable references to all surrounding elements, right? I don't think this is possible to do it directly. The problem is, Rust cannot statically prove that you want mutable references to different cells. If it ignored this, then, for example, you could make a slight mistake in indexing and get two mutable references to the same data, which is something Rust guarantees to prevent. Hence it disallows this.
On the language level this is caused by IndexMut trait. You can see how its only method's self parameter lifetime is tied to the result lifetime:
fn index_mut(&'a mut self, index: Idx) -> &'a mut Self::Output;
This means that if this method is called (implicitly through an indexing operation) then the whole object will be borrowed mutably until the resulting reference goes out of scope. This prevents calling &mut a[i] multiple times.
The most simple and safest way to fix this would be to refactor your code in a "double buffering" manner - you have two instances of the field and copy data between each other on the each step. Alternatively, you can create a temporary field on each step and replace the main one with it after all computations but it is probably less efficient than swapping two fields.
Another way to solve this would be, naturally, using raw *mut pointers. This is unsafe and should only be used directly as the last resort. You can use unsafety, however, to implement a safe abstraction, something like
fn index_multiple_mut<'a, T>(input: &'a mut [Vec<T>], indices: &[(usize, usize)]) -> Vec<&'a mut T>
where you first check that all indices are different and then use unsafe with some pointer casts (with transmute, probably) to create the resulting vector.
A third possible way would be to use split_at_mut() method in some clever way, but I'm not that sure that it is possible, and if it is, it is likely not very convenient.