use rayon::prelude::*; // 1.1.0 // 0.9.0

struct MultiMutIter<'a, T, I> {
    last_idx: Option<usize>,
    idx_iter: I,
    slice: &'a mut [T],
}

impl<'a, I, T> Iterator for MultiMutIter<'a, T, I>
where
    I: Iterator<Item = usize>,
{
    type Item = &'a mut T;

    fn next(&mut self) -> Option<Self::Item> {
        let idx = self.idx_iter.next()?;
        let offset = if let Some(last_idx) = self.last_idx {
            if idx <= last_idx {
                panic!("indexes were not unique or not sorted");
            }
            idx - last_idx - 1
        } else {
            idx
        };
        self.last_idx = Some(idx);
        // Make sure new_slice's lifetime isn't limited to &self
        let new_slice = std::mem::replace(&mut self.slice, &mut []);
        let new_slice = new_slice.split_at_mut(offset).1;
        let (elem, new_slice) = new_slice.split_first_mut().unwrap();
        self.slice = new_slice;
        Some(elem)
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        self.idx_iter.size_hint()
    }
}

fn get_multi_mut<T, I>(slice: &mut [T], idx_iter: I) -> MultiMutIter<T, I>
where
    I: Iterator<Item = usize>,
{
    MultiMutIter {
        last_idx: None,
        idx_iter,
        slice,
    }
}

fn main() {
    let mut vecs = vec![
        vec![1, 2, 3],
        vec![1, 2, 3],
        vec![1, 2, 3],
        vec![1, 2, 3],
        vec![1, 2, 3],
        vec![1, 2, 3],
    ];
    let mut indices = vec![4, 3, 1];
    // Ensure all indices are unique to ensure safety:
    indices.sort_unstable();
    let len = indices.len();
    indices.dedup();
    assert!(indices.len() == len);
    let mut_vecs: Vec<_> = get_multi_mut(&mut vecs, indices.into_iter()).collect();

    mut_vecs.into_par_iter().for_each(|vec| {
        vec.push(4);
    });
    println!("vecs: {:?}", vecs);
}