Combinations Iterator

combi.rs

// Outputs:
// [1, 2, 3], ([4, 5, 6])
// [1, 2, 4], ([5, 6, 3])
// [1, 2, 5], ([6, 3, 4])
// [1, 2, 6], ([3, 4, 5])
// [1, 3, 4], ([5, 6, 2])
// [1, 3, 5], ([6, 4, 2])
// [1, 3, 6], ([4, 5, 2])
// [1, 4, 5], ([6, 3, 2])
// [1, 4, 6], ([5, 3, 2])
// [1, 5, 6], ([4, 3, 2])
// [2, 3, 4], ([5, 6, 1])
// [2, 3, 5], ([6, 4, 1])
// [2, 3, 6], ([4, 5, 1])
// [2, 4, 5], ([6, 3, 1])
// [2, 4, 6], ([5, 3, 1])
// [2, 5, 6], ([4, 3, 1])
// [3, 4, 5], ([6, 2, 1])
// [3, 4, 6], ([5, 2, 1])
// [3, 5, 6], ([4, 2, 1])
// [4, 5, 6], ([3, 2, 1])

fn main() {
    let mut v = vec![1, 2, 3, 4, 5, 6];
    let r = 3;

    unsafe {
        let ptr = v.as_mut_ptr();
        out(&v, r);
        for _ in CombiIterator3::new_ptr(ptr, v.len()) {
            out(&v, r);
        }
    }
    
    // At this point, v is back in the original order
}

pub fn out(elems: &[usize], n: usize) {
    println!("{:?}, ({:?})", &elems[..n], &elems[n..]);
}

struct CombiIterator1<T> {
    ptr: *mut T,
    n: usize,
    step: usize,
}

struct CombiIterator2<T> {
    ptr: *mut T,
    n: usize,
    iter1: Option<CombiIterator1<T>>,
    iter2: Option<Box<CombiIterator2<T>>>,
}

struct CombiIterator3<T> {
    ptr: *mut T,
    n: usize,
    iter1: Option<CombiIterator2<T>>,
    iter2: Option<Box<CombiIterator3<T>>>,
}

impl<T> CombiIterator3<T> {
    pub fn new(slice: &mut [T]) -> Self {
        assert!(3 <= slice.len());

        let ptr = slice.as_mut_ptr();
        unsafe { Self::new_ptr(ptr, slice.len()) }
    }

    unsafe fn new_ptr(ptr: *mut T, n: usize) -> Self {
        assert!(n > 2);
        Self {
            ptr,
            n,
            iter1: Some(CombiIterator2::new_ptr(ptr.add(1), n - 1)),
            iter2: None,
        }
    }
}

impl<T> CombiIterator2<T> {
    pub fn new(slice: &mut [T]) -> Self {
        assert!(2 <= slice.len());

        let ptr = slice.as_mut_ptr();
        unsafe { Self::new_ptr(ptr, slice.len()) }
    }

    unsafe fn new_ptr(ptr: *mut T, n: usize) -> Self {
        assert!(n > 1);
        Self {
            ptr,
            n,
            iter1: Some(CombiIterator1::new_ptr(ptr.add(1), n - 1)),
            iter2: None,
        }
    }
}

impl<T> CombiIterator1<T> {
    pub fn new(slice: &mut [T]) -> Self {
        assert!(1 <= slice.len());

        let ptr = slice.as_mut_ptr();
        unsafe { Self::new_ptr(ptr, slice.len()) }
    }

    unsafe fn new_ptr(ptr: *mut T, n: usize) -> Self {
        assert!(n > 0);
        Self { ptr, n, step: 0 }
    }
}

impl<T> Iterator for CombiIterator3<T> {
    type Item = ();

    fn next(&mut self) -> Option<Self::Item> {
        if self.n == 3 {
            return None;
        }

        if let Some(iter) = &mut self.iter2 {
            if iter.next().is_none() {
                unsafe {
                    let slice = std::slice::from_raw_parts_mut(self.ptr, self.n);
                    slice.rotate_right(1);

                    self.iter2 = None;
                }

                None
            } else {
                Some(())
            }
        } else if let Some(iter) = &mut self.iter1 {
            if iter.next().is_none() {
                unsafe {
                    let slice = std::slice::from_raw_parts_mut(self.ptr, self.n);
                    slice.rotate_left(1);

                    self.iter2 = Some(Box::new(CombiIterator3::new_ptr(self.ptr, self.n - 1)));
                    self.iter1 = None;
                }
            }

            Some(())
        } else {
            None
        }
    }
}

impl<T> Iterator for CombiIterator2<T> {
    type Item = ();

    fn next(&mut self) -> Option<Self::Item> {
        if self.n == 2 {
            return None;
        }

        if let Some(iter) = &mut self.iter2 {
            if iter.next().is_none() {
                unsafe {
                    let slice = std::slice::from_raw_parts_mut(self.ptr, self.n);
                    slice.rotate_right(1);

                    self.iter2 = None;
                }

                None
            } else {
                Some(())
            }
        } else if let Some(iter) = &mut self.iter1 {
            if iter.next().is_none() {
                unsafe {
                    let slice = std::slice::from_raw_parts_mut(self.ptr, self.n);
                    slice.rotate_left(1);

                    self.iter2 = Some(Box::new(CombiIterator2::new_ptr(self.ptr, self.n - 1)));
                    self.iter1 = None;
                }
            }

            Some(())
        } else {
            None
        }
    }
}

impl<T> Iterator for CombiIterator1<T> {
    type Item = ();

    fn next(&mut self) -> Option<Self::Item> {
        self.step += 1;
        let Self { ptr, n, step } = *self;

        unsafe {
            let slice = std::slice::from_raw_parts_mut(ptr, n);

            slice.rotate_left(1);

            if step == slice.len() {
                None
            } else {
                Some(())
            }
        }
    }
}