Skip to main content

marigold_impl/
combinations.rs

1use async_trait::async_trait;
2use futures::stream::Stream;
3use futures::stream::StreamExt;
4use itertools::Combinations;
5use tracing::instrument;
6
7#[async_trait]
8pub trait Combinable<T> {
9    async fn combinations(
10        self,
11        k: usize,
12    ) -> futures::stream::Iter<Combinations<std::vec::IntoIter<T>>>;
13}
14
15/// This is a glue trait to allow streams to use Combinable in itertools.
16/// The current implementation eagerly consumes the parent stream.
17#[async_trait]
18impl<T, SInput> Combinable<T> for SInput
19where
20    SInput: Stream<Item = T> + Send,
21    T: Clone + Send + std::fmt::Debug,
22{
23    #[instrument(skip(self))]
24    async fn combinations(
25        self,
26        k: usize,
27    ) -> futures::stream::Iter<Combinations<std::vec::IntoIter<T>>> {
28        use itertools::Itertools;
29
30        let combinations_iterable = self.collect::<Vec<_>>().await.into_iter().combinations(k);
31        futures::stream::iter(combinations_iterable)
32    }
33}
34
35#[cfg(test)]
36mod tests {
37    use super::Combinable;
38    use futures::stream::StreamExt;
39
40    #[tokio::test]
41    async fn combinations_basic() {
42        assert_eq!(
43            futures::stream::iter(vec![1, 2, 3])
44                .combinations(2)
45                .await
46                .collect::<Vec<_>>()
47                .await,
48            vec![vec![1, 2], vec![1, 3], vec![2, 3]]
49        );
50    }
51
52    #[tokio::test]
53    async fn combinations_k_zero() {
54        // k=0 yields exactly one combination: the empty set
55        let result: Vec<Vec<i32>> = futures::stream::iter(vec![1, 2, 3])
56            .combinations(0)
57            .await
58            .collect::<Vec<_>>()
59            .await;
60        assert_eq!(result, vec![Vec::<i32>::new()]);
61    }
62
63    #[tokio::test]
64    async fn combinations_k_equals_length() {
65        // k = len yields exactly one combination containing all elements
66        let result: Vec<Vec<i32>> = futures::stream::iter(vec![1, 2, 3])
67            .combinations(3)
68            .await
69            .collect::<Vec<_>>()
70            .await;
71        assert_eq!(result, vec![vec![1, 2, 3]]);
72    }
73
74    #[tokio::test]
75    async fn combinations_k_exceeds_length() {
76        // k > len yields nothing
77        let result: Vec<Vec<i32>> = futures::stream::iter(vec![1, 2])
78            .combinations(5)
79            .await
80            .collect::<Vec<_>>()
81            .await;
82        assert!(result.is_empty());
83    }
84
85    #[tokio::test]
86    async fn combinations_empty_stream() {
87        let result: Vec<Vec<i32>> = futures::stream::iter(Vec::<i32>::new())
88            .combinations(2)
89            .await
90            .collect::<Vec<_>>()
91            .await;
92        assert!(result.is_empty());
93    }
94
95    #[tokio::test]
96    async fn combinations_count() {
97        // C(5,2) = 10
98        let result: Vec<Vec<i32>> = futures::stream::iter(vec![1, 2, 3, 4, 5])
99            .combinations(2)
100            .await
101            .collect::<Vec<_>>()
102            .await;
103        assert_eq!(result.len(), 10);
104    }
105}