Safe buffer access for MPI Receive RequestCollections

test.rs

struct RequestCollection<'a, T: 'a + Copy> {
    handles: Vec<bool>,
    buffers: Vec<&'a mut [T]>,
}

impl<'a, T: 'a + Copy> RequestCollection<'a, T> {
    fn new() -> Self {
        RequestCollection {
            handles: Vec::new(),
            buffers: Vec::new(),
        }
    }

    fn push(&mut self, buffer: &'a mut [T]) {
        self.handles.push(true);
        self.buffers.push(buffer);
    }

    fn mark_completed(&mut self, idx: usize, data: &[T]) {
        self.handles[idx] = false;
        for (x, data) in self.buffers[idx].iter_mut().zip(data) {
            *x = *data;
        }
    }

    fn get_buffer<'b>(&'a self, idx: usize) -> &'b [T]
    where 'a: 'b {
        assert!(!self.handles[idx]);

        unsafe { std::mem::transmute::<&'a [T], &'b [T]>(self.buffers[idx]) }
    }
}

#[test]
fn basic() {
    let mut first_data = [0; 4];

    let mut collection = RequestCollection::new();

    collection.push(&mut first_data);

    collection.mark_completed(0, &[1; 4]);

    let buffer = collection.get_buffer(0);
    assert_eq!([1; 4], buffer);
}

#[test]
#[should_panic]
fn should_panic() {
    let mut first_data = [0; 4];
    let mut second_data = [0; 5];

    let mut collection = RequestCollection::new();

    collection.push(&mut first_data);
    collection.push(&mut second_data);

    collection.mark_completed(0, &[1; 4]);

    let _ = collection.get_buffer(1);
}