marigold_impl/
permutations.rs

1use async_trait::async_trait;
2use futures::stream::Stream;
3use futures::stream::StreamExt;
4use itertools::Permutations;
5use tracing::instrument;
6
7#[async_trait]
8pub trait Permutable<T: Clone> {
9    async fn permutations(
10        self,
11        k: usize,
12    ) -> futures::stream::Iter<Permutations<std::vec::IntoIter<T>>>;
13
14    async fn permutations_with_replacement(
15        self,
16        k: usize,
17    ) -> futures::stream::Iter<itertools::structs::MultiProduct<std::vec::IntoIter<T>>>;
18}
19
20/// This is a glue trait to allow streams to use Permutable in itertools.
21/// The current implementation eagerly consumes the parent stream.
22#[async_trait]
23impl<T, SInput> Permutable<T> for SInput
24where
25    SInput: Stream<Item = T> + Send,
26    T: Clone + Send + std::fmt::Debug,
27{
28    #[instrument(skip(self))]
29    async fn permutations(
30        self,
31        k: usize,
32    ) -> futures::stream::Iter<Permutations<std::vec::IntoIter<T>>> {
33        use itertools::Itertools;
34
35        let permutations_iterable = self.collect::<Vec<_>>().await.into_iter().permutations(k);
36        futures::stream::iter(permutations_iterable)
37    }
38
39    #[instrument(skip(self))]
40    async fn permutations_with_replacement(
41        self,
42        k: usize,
43    ) -> futures::stream::Iter<itertools::structs::MultiProduct<std::vec::IntoIter<T>>> {
44        use itertools::Itertools;
45
46        let items = self.collect::<Vec<_>>().await;
47        let iterators: Vec<std::vec::IntoIter<T>> =
48            (0..k).map(|_| items.clone().into_iter()).collect();
49        let product = iterators.into_iter().multi_cartesian_product();
50        futures::stream::iter(product)
51    }
52}
53
54#[cfg(test)]
55mod tests {
56    use super::Permutable;
57    use futures::stream::StreamExt;
58
59    #[tokio::test]
60    async fn permutations() {
61        assert_eq!(
62            futures::stream::iter(vec![1, 2, 3])
63                .permutations(2)
64                .await
65                .collect::<Vec<_>>()
66                .await,
67            vec![
68                vec![1, 2],
69                vec![1, 3],
70                vec![2, 1],
71                vec![2, 3],
72                vec![3, 1],
73                vec![3, 2]
74            ]
75        );
76    }
77
78    #[tokio::test]
79    async fn permutations_with_replacement() {
80        assert_eq!(
81            futures::stream::iter(vec![0, 1, 2])
82                .permutations_with_replacement(2)
83                .await
84                .collect::<Vec<_>>()
85                .await,
86            vec![
87                vec![0, 0],
88                vec![0, 1],
89                vec![0, 2],
90                vec![1, 0],
91                vec![1, 1],
92                vec![1, 2],
93                vec![2, 0],
94                vec![2, 1],
95                vec![2, 2],
96            ]
97        );
98    }
99}
marigold_impl/permutations.rs

marigold_impl/
permutations.rs
