For some context, this is inspired by my attempt to solve this SO question.

I have a mutably borrowed slice representing a 2D array, and I want to split the borrow such that I can access all the rows or all the columns at the same time. Splitting into rows is easy, I just have to call split_at_mut repeatedly to obtain a mutable subslice for each row.

But columns are harder because the elements in a column are not contiguous in the slice. My idea is to create a pseudo-slice (Column) that implements Index such that it accesses the elements in a column inside the original slice. The Column struct has a period that is the width of a row, and an offset that is the index of the column. Columns that have the same period but different offsets are guaranteed to be disjoint, and sharing the same mutable borrow should be safe.

The way to create them is to construct a ColumnIterMut iterator, which will yield all the columns for a specified period.

To construct multiple mutable columns sharing the same slice simultaneously, I had to use raw pointers. I've been writing Safe Rust for a fair amount of time now, but I'm still just learning Unsafe Rust. So, most importantly, I would like you to check if my implementation is actually sound. I'm mainly concerned by aliasing rules and conforming to the Stacked Borrows model. The code works and even Miri doesn't complain, but I'm still not 100% sure that it will work in all circumstances.

A less important aspect is, if my code is correct, what additional traits/methods should I implement on Column to make it more user-friendly?

Here is the code:

use core::ops::{Index, IndexMut};
use std::marker::PhantomData;

pub struct ColumnIterMut<'a, T>{
    data: &'a mut [T],
    period: usize,
    offset: usize
}

impl<'a, T> ColumnIterMut<'a, T> {
    pub fn new(data: &'a mut [T], period: usize) -> Self {
        assert!(period > 0);

        Self { data, period, offset: 0 }
    }
}

impl<'a, T> Iterator for ColumnIterMut<'a, T> {
    type Item = Column<'a, T>;

    fn next(&mut self) -> Option<Self::Item> {
        if self.offset < self.period {
            let col = Column{
                ptr: self.data as *mut [T], 
                _lifetime: PhantomData,
                offset: self.offset,
                period: self.period
            };

            self.offset += 1;

            Some(col)
        } else {
            None
        }
    }
}

//INVARIANT: period > 0
//INVARIANT: offset < period
//INVARIANT: ptr is always non-null, well-aligned and points to a valid instance of [T]
//INVARIANT: all Column structs sharing the same slice of data simultaneously
//           must have equal `period`s and distinct `offset`s
pub struct Column<'a, T>{
    ptr: *mut [T],
    _lifetime: PhantomData<&'a mut [T]>,
    
    period: usize,
    offset: usize,
}

impl<'a, T> Column<'a, T> {
    pub fn len(&self) -> usize {
        unsafe{
            ((*self.ptr).len() + self.period - self.offset - 1) / self.period
        }
    }

    pub fn is_empty(&self) -> bool {
        self.len() == 0
    }

    fn map_index(&self, index: usize) -> usize {
        index * self.period + self.offset
    }
}

impl<'a, T> Index<usize> for Column<'a, T> {
    type Output = T;
    
    fn index(&self, index: usize) -> &Self::Output {
        unsafe{
            //SAFETY: if the invariants are maintained, the indices returned by 
            //        `Self::map_index()` will be exclusive to this instance of the struct
            &(*self.ptr)[self.map_index(index)]
        }
    }
}

impl<'a, T> IndexMut<usize> for Column<'a, T> {    
    fn index_mut(&mut self, index: usize) -> &mut Self::Output {
        unsafe{
            //SAFETY: if the invariants are maintained, the indices returned by 
            //        `Self::map_index()` will be exclusive to this instance of the struct
            &mut (*self.ptr)[self.map_index(index)]
        }
    }
}

unsafe impl<'a, T> Send for Column<'a, T> where [T]: Send {}
unsafe impl<'a, T> Sync for Column<'a, T> where [T]: Sync {}

GitHub