Mutable column iterator

[yosefm]

I have a struct representing a vec-backed 2D image, inspired by the image crate:

pub struct Array2D<ElementType> 
    where ElementType: Primitive
{
    data: Vec<ElementType>,
    width: u32,
    height: u32,
}

I want to iterate it mutably and immutably. Iterating by rows is simple:

impl<ElementType: Primitive> Array2D<ElementType> {
    pub fn rows(&self) -> Chunks<ElementType> {
        self.data.chunks(self.width as usize)
    }
    
    pub fn rows_mut(&mut self) -> ChunksMut<ElementType> {
        self.data.chunks_mut(self.width as usize)
    }
}

But trying to do the same thing over columns defeated me. I tried to write my own iterator:

pub struct ColumnIterMut<'a, ElemT> 
{
    current_col: usize,
    stride: usize,
    underlying: &'a mut Vec<ElemT>,
}
impl<'a, ElemT: Primitive> Iterator for ColumnIterMut<'a, ElemT> {
    type Item = StepBy<IterMut<'a, ElemT>>;
    
    fn next(&mut self) -> Option<Self::Item> {
        if self.current_col >= self.stride {
            None
        } else {
            let slice = &mut self.underlying[self.current_col as usize ..];
            let ret = Some(
                slice.iter_mut().step_by(self.stride as usize));
            self.current_col += 1;
            ret
        }
    }
}

This does not compile, apparently because the mutable reference inside the function can't outlive the function, although it's a reference to a &'a vector. If I remove the mutability everywhere, this iterator class compiles, even though the same lifetimes are involved.

I'm looking either for suggetions on how to mutably iterate over columns, or how to get my iterator attempt to succeed. Thanks in advance.

[najamelan]

Hi,

I would suggest you put your question in the help category. I sure look for posts in that category with no answers to avoid people being left behind.

Another thing I wonder. If I was facing this problem, the first thing goes through my head is: "Sure, I can make this datastructure, but how long is it going to take me to make it decent and polished?" Is that worth it? Like 2h with tests, docs, debugging, etc? 4h maybe? 6 if I'm being perfectionist?

And then I would think bugger, someone surely spent 6+ hours on something like this to make it really smart, and to have better performance than a naive impl, say if you want to add an extra element to a row, since everything is in the same vec, you going to re-allocate and shift everything?

So doing a quick search on lib.rs shows that there is a number of crates out there already: md_grid, grid_2d, array2d, gridd, gridly, ...

Maybe these don't suit your needs, why? In any case some of them have mutable iterators on both rows and columns, so I would think if I still wanted to implement this I would at least have some ideas of how it can be done...

[alice]

The challenge is that a mutable reference cannot overlap with any other reference unless one is a borrow of the other, but due to the 'a lifetime, the StepBy iterators returned by the column iterator are a borrow of the vector, not the column iterator itself, and thus the mutable reference inside the StepBy iterator overlaps with your underlying reference without either being a borrow of the other, which is not allowed.

[yosefm]

I actually couldn't find any package that gives mutable column iteration. So it might have been worth it to do my own even if my goal was not primarily educational, which it is.

The best example of something similar to what I want is std::Vec::ChunksMut. However, it relies on the premise that chunks do not overlap, which does not hold in the strided case.

I wonder what to make of the fact that nobody seems to have done mutable column iteration. Is it just too hard to do in Rust? Because I can't convince myself that it's never needed. Think of calculating a distance map in a memory-constrained application. The best-in-class (as far as I know) algorithm works in-place and serves this use-case. But it requires mutable column iteration. (For this case there's the distance-transform crate, but it doesn't work in place).

[yosefm]

Here's my best attempt so far, using a fair bit of unsafe code to circumvent the borrow-checker. It makes it possible to hold two mutable references to a column (or element) at the same time, if you don't use the iterator through a for-loop. I can live with that, though I still wonder if it's possible to get a tighter result.

pub struct StridedIterMut<'a, T> {
    underlying: &'a mut Vec<T>,
    stride: usize,
    len: usize,
    current: *mut T,
}

impl<'a, T> StridedIterMut<'a, T> {
    fn new(underlying: &'a mut Vec<T>, start: usize, stride: usize, len: usize) -> Self {
        let Icurrent = unsafe {
            underlying.as_mut_ptr().add(start)
        };
        StridedIterMut {
            underlying, stride, len, current: Icurrent
        }
    }
    
    pub fn reset_pos(&mut self, start: usize, len: usize) {
        self.current = unsafe {
            self.underlying.as_mut_ptr().add(start)
        };
        self.len = len;
    }
}

impl<'a, T> Iterator for StridedIterMut<'a, T> {
    type Item = &'a mut T;
    
    fn next(&mut self) -> Option<Self::Item>{
        if self.len == 0 {
            None
        }
        else {
            unsafe {
                let ret = Some(&mut *self.current);
                self.current = self.current.add(self.stride);
                self.len -= 1;
                ret
            }
        }
    }
}

pub struct ColumnsIterMut<'a, T> {
    col_iter: StridedIterMut<'a, T>,
    len: usize,
    column_len: usize,
    current: usize,
}

impl<'a, T: Primitive> ColumnsIterMut<'a, T> {
    fn new(
        underlying: &'a mut Vec<T>, 
        start: usize, stride: usize, 
        len: usize, column_len: usize) -> Self 
    {
        ColumnsIterMut {
            col_iter: StridedIterMut::new(underlying, start, stride, column_len),
            len, column_len, current: start
        }
    }
}

impl<'a, T: Primitive> Iterator for ColumnsIterMut<'a, T> {
    type Item = &'a mut StridedIterMut<'a, T>;
    
    fn next(&mut self) -> Option<Self::Item>{
        if self.len == self.current {
            None
        }
        else {
            self.col_iter.reset_pos(self.current, self.column_len);
            self.current += 1;
            unsafe {
                let tmp = &mut self.col_iter as *mut StridedIterMut<'a, T>;
                Some(&mut *tmp)
            }
        }
    }
}

[najamelan]

sorry, when I looked at the api docs of grid_2d, I saw several mut iterators and thought I saw columns, but indeed they don't have that.

I had a try in the playground. The lifetime problem boils down to this: The contract of the Iterator trait basically says: "If you have a &mut to Self, you can call next, for any lifetime of the &mut Self". Where as you would like to say:

fn next(self: &'a mut Self) -> Option<Self::Item>

Basically saying I only implement next when the reference outlives 'a, which doesn't uphold the contract of the trait. So I thought you need to specify that you only implement iterator in that case, rather than saying I only implement next. So I tried with:

impl<'a, ElemT: Primitive> Iterator for &'a mut ColumnIterMut<'a, ElemT>

But unfortunately I don't manage to dereference the self to get a &'a mut ColumnIterMut<'a, ElemT> out of it.

As from this earlier thread though it seems it's not possible without GAT's, which would allow you to tie the lifetime in the associated type to the lifetime of self. The use of a the streaming-iterator crate is suggested.

As an alternative, you could make ColumnIterMut static by using Rc<RefCell<Vec<ElemT>>> or Arc<Mutex<Vec<ElemT>>>. It obviously implies locking and some overhead, but it would guarantee that if you hold a ColumnIterMut, all required memory is still alive.