permutohedron 0.1.5

use std::marker::PhantomData;

/// Heap's algorithm for generating permutations, recursive version.
///
/// The recursive algorithm supports slices of any size (even though
/// only a small number of elements is practical), and is generally
/// a bit faster than the iterative version.
pub fn heap_recursive<T, F>(xs: &mut [T], mut f: F) where F: FnMut(&mut [T])
{
    match xs.len() {
        0 | 1 => f(xs),
        2 => {
            // [1, 2], [2, 1]
            f(xs);
            xs.swap(0, 1);
            f(xs);
        }
        n => heap_unrolled_(n, xs, &mut f),
    }
}

// TODO: Find a more parallel version with less data dependencies:
// i.e. don't swap the same items (for example index 0) every time.

/// Unrolled version of heap's algorithm due to Sedgewick
fn heap_unrolled_<T>(n: usize, xs: &mut [T], f: &mut FnMut(&mut [T])) {
    debug_assert!(n >= 3);
    match n {
        3 => {
            // [1, 2, 3], [2, 1, 3], [3, 1, 2], [1, 3, 2], [2, 3, 1], [3, 2, 1]
            f(xs);
            xs.swap(0, 1);
            f(xs);
            xs.swap(0, 2);
            f(xs);
            xs.swap(0, 1);
            f(xs);
            xs.swap(0, 2);
            f(xs);
            xs.swap(0, 1);
            f(xs);
        }
        n => {
            for i in 0..n - 1 {
                heap_unrolled_(n - 1, xs, f);
                let j = if n % 2 == 0 { i } else { 0 };
                // One swap *between* each iteration.
                xs.swap(j, n - 1);
            }
            heap_unrolled_(n - 1, xs, f);
        }
    }
}

/// Maximum number of elements we can generate permutations for using
/// Heap's algorithm (iterative version).
pub const MAXHEAP: usize = 16;

/// Heap's algorithm for generating permutations.
///
/// An iterative method of generating all permutations of a sequence.
///
/// Note that for *n* elements there are *n!* (*n* factorial) permutations.
// lock the repr since it performs the best in this order..(?)
#[repr(C)]
pub struct Heap<'a, Data: 'a + ?Sized, T: 'a> {
    data: &'a mut Data,
    // c, and n: u8 is enough range.
    // u32 performs better for n, u8 for c.
    // n: == !0 at start, 0 after first permutation is emitted
    n: u32,
    // c[x] is the counter for the (x + 1) th location
    c: [u8; MAXHEAP - 1],
    _element: PhantomData<&'a mut T>
}

impl<'a, T, Data: ?Sized> Heap<'a, Data, T>
    where Data: AsMut<[T]>
{
    /// Create a new `Heap`.
    pub fn new(data: &'a mut Data) -> Self {
        assert!(data.as_mut().len() <= MAXHEAP);
        Heap {
            data: data,
            c: [0; MAXHEAP - 1],
            n: !0,
            _element: PhantomData,
        }
    }

    /// Return a reference to the inner data
    pub fn get(&self) -> &Data {
        self.data
    }

    /// Return a mutable reference to the inner data
    pub fn get_mut(&mut self) -> &mut Data {
        self.data
    }

    /// Reset the permutations walker, without changing the data. It allows
    /// generating permutations again with the current state as starting
    /// point.
    pub fn reset(&mut self) {
        self.n = !0;
        for c in &mut self.c[..] { *c = 0; }
    }

    /// Step the internal data into the next permutation and return
    /// a reference to it. Return `None` when all permutations
    /// have been visited.
    ///
    /// Note that for *n* elements there are *n!* (*n* factorial) permutations.
    pub fn next_permutation(&mut self) -> Option<&mut Data> {
        if self.n == !0 {
            self.n = 0;
            Some(self.data)
        } else {
            while 1 + (self.n as usize) < self.data.as_mut().len() {
                let n = self.n as u8;
                let nu = self.n as usize;
                let c = &mut self.c;
                if c[nu] <= n {
                    // `n` acts like the current length - 2 of the slice we are permuting
                    // `c[n]` acts like `i` in the recursive algorithm
                    let j = if nu % 2 == 0 { c[nu] as usize } else { 0 };
                    self.data.as_mut().swap(j, nu + 1);
                    c[nu] += 1;
                    self.n = 0;
                    return Some(self.data);
                } else {
                    c[nu] = 0;
                    self.n += 1;
                }
            }
            None
        }
    }
}

/// Iterate the permutations
///
/// **Note:** You can also generate the permutations lazily by using
/// `.next_permutation()`.
impl<'a, Data: ?Sized, T> Iterator for Heap<'a, Data, T>
    where Data: AsMut<[T]> + ToOwned,
{
    type Item = Data::Owned;
    fn next(&mut self) -> Option<Self::Item> {
        match self.next_permutation() {
            None => None,
            Some(perm) => Some(perm.to_owned()),
        }
    }
}

/// Compute *n!* (*n* factorial)
pub fn factorial(n: usize) -> usize {
    let mut prod = 1;
    for x in 1..n + 1 { prod *= x; }
    prod
}

#[test]
fn first_and_reset() {
    let mut data = [1, 2, 3];
    let mut heap = Heap::new(&mut data);
    let mut perm123 = vec![[1, 2, 3], [2, 1, 3], [3, 1, 2], [1, 3, 2], [2, 3, 1], [3, 2, 1]];
    assert_eq!(heap.by_ref().collect::<Vec<_>>(), perm123);

    // test reset
    heap.reset();
    // for the 1,2,3 case this happens to work out to the reverse order
    perm123.reverse();
    assert_eq!(heap.by_ref().collect::<Vec<_>>(), perm123);
}

#[test]
fn permutations_0_to_6() {
    let mut data = [0; 6];
    for n in 0..data.len() {
        let count = factorial(n);
        for (index, elt) in data.iter_mut().enumerate() {
            *elt = index + 1;
        }
        let mut permutations = Heap::new(&mut data[..n]).collect::<Vec<_>>();
        assert_eq!(permutations.len(), count);
        permutations.sort();
        permutations.dedup();
        assert_eq!(permutations.len(), count);
        // Each permutation contains all of 1 to n
        assert!(permutations.iter().all(|perm| perm.len() == n));
        assert!(permutations.iter().all(|perm| (1..n + 1).all(|i| perm.iter().position(|elt| *elt == i).is_some())));
    }
}

#[test]
fn count_permutations_iter() {
    let mut data = [0; 10];
    for n in 0..data.len() + 1 {
        let count = factorial(n);
        let mut permutations = Heap::new(&mut data[..n]);
        let mut i = 0;
        while let Some(_) = permutations.next_permutation() {
            i += 1;
        }
        assert_eq!(i, count);
        println!("{}! = {} ok", n, count);
    }
}

#[test]
fn count_permutations_recur() {
    let mut data = [0; 10];
    for n in 0..data.len() + 1 {
        let count = factorial(n);
        let mut i = 0;
        heap_recursive(&mut data[..n], |_| i += 1);
        assert_eq!(i, count);
        println!("{}! = {} ok", n, count);
    }
}

#[test]
fn permutations_0_to_6_recursive() {
    let mut data = [0; 6];
    for n in 0..data.len() {
        let count = factorial(n);
        for (index, elt) in data.iter_mut().enumerate() {
            *elt = index + 1;
        }
        let mut permutations = Vec::with_capacity(count);
        heap_recursive(&mut data[..n], |elt| permutations.push(elt.to_owned()));
        println!("{:?}", permutations);
        assert_eq!(permutations.len(), count);
        permutations.sort();
        permutations.dedup();
        assert_eq!(permutations.len(), count);
        // Each permutation contains all of 1 to n
        assert!(permutations.iter().all(|perm| perm.len() == n));
        assert!(permutations.iter().all(|perm| (1..n + 1).all(|i| perm.iter().position(|elt| *elt == i).is_some())));
    }
}