libreda-logic 0.0.3

// SPDX-FileCopyrightText: 2022 Thomas Kramer <code@tkramer.ch>
//
// SPDX-License-Identifier: AGPL-3.0-or-later

//! Iterate over all permutations of a sorted collection.

use streaming_iterator::StreamingIterator;

use super::permutation_gray_code::{permutation_swaps, PermutationSwaps};

/// Trait for collections which can be permuted by swapping pairs of elements.
pub trait Permutable {
    /// Get the number of elements in the collection.
    fn len(&self) -> usize;

    /// Swap two elements.
    fn swap(&mut self, i: usize, j: usize);
}

/// Iterator over permutations.
pub struct PermutationIter<T> {
    state: T,
    swaps: PermutationSwaps,
    first: bool, // Used only for streaming iterator.
    done: bool,
}

impl<T> Iterator for PermutationIter<T>
where
    T: Permutable + Clone,
{
    type Item = T;

    fn next(&mut self) -> Option<Self::Item> {
        if self.done {
            None
        } else {
            let output = self.state.clone();
            if let Some((i, j)) = self.swaps.next() {
                self.state.swap(i, j);
            } else {
                self.done = true;
            }
            Some(output)
        }
    }

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

/// Iterate over all permutations of the `elements`.
/// In each step only two elements are swapped.
pub fn permutations<T>(elements: T) -> PermutationIter<T>
where
    T: Permutable + Clone,
{
    let n = elements.len();

    PermutationIter {
        state: elements,
        swaps: permutation_swaps(n),
        first: true,
        done: false,
    }
}

#[test]
pub fn test_permutations() {
    #[derive(Copy, Clone, Hash, Eq, PartialEq)]
    struct Elements {
        elements: [u32; 4],
    }

    impl Permutable for Elements {
        fn len(&self) -> usize {
            self.elements.len()
        }

        fn swap(&mut self, i: usize, j: usize) {
            self.elements.swap(i, j);
        }
    }

    let elements = Elements {
        elements: [1, 2, 3, 4],
    };

    let all_permutations: std::collections::HashSet<_> = permutations(elements).collect();

    assert_eq!(all_permutations.len(), 1 * 2 * 3 * 4);
}

impl<T> StreamingIterator for PermutationIter<T>
where
    T: Permutable,
{
    type Item = T;

    fn advance(&mut self) {
        if self.first {
            // Skip the first 'advance' to output the original element at the beginning.
            self.first = false
        } else if let Some((i, j)) = self.swaps.next() {
            self.state.swap(i, j);
        } else {
            self.done = true;
        }
    }

    fn get(&self) -> Option<&Self::Item> {
        (!self.done).then_some(&self.state)
    }

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

#[test]
pub fn test_permutation_stream() {
    #[derive(Debug, Copy, Clone, Hash, Eq, PartialEq)]
    struct Elements {
        elements: [u32; 4],
    }

    impl Permutable for Elements {
        fn len(&self) -> usize {
            self.elements.len()
        }

        fn swap(&mut self, i: usize, j: usize) {
            self.elements.swap(i, j);
        }
    }

    let elements = Elements {
        elements: [1, 2, 3, 4],
    };

    let perm = permutations(elements);
    let mut perm_stream = permutations(elements);

    for e in perm {
        assert_eq!(StreamingIterator::next(&mut perm_stream), Some(&e));
    }

    StreamingIterator::next(&mut perm_stream);
    assert!(perm_stream.is_done());
}