Skip to main content

nu_command/filters/
permutations.rs

1use nu_engine::command_prelude::*;
2use nu_protocol::{ListStream, Signals};
3
4#[derive(Clone)]
5pub struct Permutations;
6
7impl Command for Permutations {
8    fn name(&self) -> &str {
9        "permutations"
10    }
11
12    fn signature(&self) -> Signature {
13        Signature::build("permutations")
14            .input_output_types(vec![
15                (
16                    Type::List(Box::new(Type::Any)),
17                    Type::List(Box::new(Type::List(Box::new(Type::Any)))),
18                ),
19                (Type::table(), Type::table()),
20            ])
21            .category(Category::Filters)
22    }
23
24    fn description(&self) -> &str {
25        "Generates all permutations of the input list."
26    }
27
28    fn search_terms(&self) -> Vec<&str> {
29        vec!["arrangements", "orderings", "combinatorics", "factorial"]
30    }
31
32    fn run(
33        &self,
34        engine_state: &EngineState,
35        _stack: &mut Stack,
36        call: &Call,
37        mut input: PipelineData,
38    ) -> Result<PipelineData, ShellError> {
39        let head = call.head;
40        let metadata = input.take_metadata();
41
42        let vals: Vec<Value> = input.into_iter().collect();
43        let signals = engine_state.signals().clone();
44
45        if vals.is_empty() {
46            return Ok(PipelineData::Value(Value::list(vec![], head), metadata));
47        }
48
49        let iter = PermutationsIter::new(vals, signals.clone(), head);
50        let stream = ListStream::new(iter, head, signals);
51        Ok(PipelineData::ListStream(stream, metadata))
52    }
53
54    fn examples(&self) -> Vec<Example<'static>> {
55        vec![
56            Example {
57                example: "[1 2 3] | permutations",
58                description: "Generate all permutations",
59                result: Some(Value::test_list(vec![
60                    Value::test_list(vec![
61                        Value::test_int(1),
62                        Value::test_int(2),
63                        Value::test_int(3),
64                    ]),
65                    Value::test_list(vec![
66                        Value::test_int(2),
67                        Value::test_int(1),
68                        Value::test_int(3),
69                    ]),
70                    Value::test_list(vec![
71                        Value::test_int(3),
72                        Value::test_int(1),
73                        Value::test_int(2),
74                    ]),
75                    Value::test_list(vec![
76                        Value::test_int(1),
77                        Value::test_int(3),
78                        Value::test_int(2),
79                    ]),
80                    Value::test_list(vec![
81                        Value::test_int(2),
82                        Value::test_int(3),
83                        Value::test_int(1),
84                    ]),
85                    Value::test_list(vec![
86                        Value::test_int(3),
87                        Value::test_int(2),
88                        Value::test_int(1),
89                    ]),
90                ])),
91            },
92            Example {
93                example: "[1 2] | permutations",
94                description: "Generate permutations of two elements",
95                result: Some(Value::test_list(vec![
96                    Value::test_list(vec![Value::test_int(1), Value::test_int(2)]),
97                    Value::test_list(vec![Value::test_int(2), Value::test_int(1)]),
98                ])),
99            },
100            Example {
101                example: "[] | permutations",
102                description: "Empty list yields no permutations",
103                result: Some(Value::test_list(vec![])),
104            },
105        ]
106    }
107}
108
109/// A streaming iterator that yields all permutations of a list using Heap's algorithm.
110///
111/// Heap's algorithm generates each permutation by swapping a single pair of
112/// elements from the previous permutation, making it well-suited for streaming
113/// (O(1) amortized per permutation). The `c` array tracks the current state
114/// of the nested loop counters.
115struct PermutationsIter {
116    values: Vec<Value>,
117    c: Vec<usize>,
118    i: usize,
119    n: usize,
120    first: bool,
121    done: bool,
122    signals: Signals,
123    span: Span,
124}
125
126impl PermutationsIter {
127    fn new(values: Vec<Value>, signals: Signals, span: Span) -> Self {
128        let n = values.len();
129        let done = n == 0;
130        Self {
131            c: vec![0; n],
132            i: 0,
133            n,
134            values,
135            first: true,
136            done,
137            signals,
138            span,
139        }
140    }
141}
142
143impl Iterator for PermutationsIter {
144    type Item = Value;
145
146    fn next(&mut self) -> Option<Self::Item> {
147        if self.done {
148            return None;
149        }
150
151        if self.signals.interrupted() {
152            return None;
153        }
154
155        if self.first {
156            self.first = false;
157            return Some(Value::list(self.values.clone(), self.span));
158        }
159
160        // Heap's algorithm: generate the next permutation by swapping.
161        while self.i < self.n {
162            if self.c[self.i] < self.i {
163                if self.i.is_multiple_of(2) {
164                    self.values.swap(0, self.i);
165                } else {
166                    self.values.swap(self.c[self.i], self.i);
167                }
168                self.c[self.i] += 1;
169                self.i = 0;
170                return Some(Value::list(self.values.clone(), self.span));
171            } else {
172                self.c[self.i] = 0;
173                self.i += 1;
174            }
175        }
176
177        self.done = true;
178        None
179    }
180}
181
182#[cfg(test)]
183mod test {
184    use super::Permutations;
185    use nu_test_support::prelude::*;
186
187    #[test]
188    fn test_examples() -> nu_test_support::Result {
189        nu_test_support::test().examples(Permutations)
190    }
191
192    #[test]
193    fn permutations_n3() -> Result {
194        let result: Value = test().run("[1 2 3] | permutations")?;
195        assert_eq!(
196            result,
197            Value::test_list(vec![
198                Value::test_list(vec![
199                    Value::test_int(1),
200                    Value::test_int(2),
201                    Value::test_int(3)
202                ]),
203                Value::test_list(vec![
204                    Value::test_int(2),
205                    Value::test_int(1),
206                    Value::test_int(3)
207                ]),
208                Value::test_list(vec![
209                    Value::test_int(3),
210                    Value::test_int(1),
211                    Value::test_int(2)
212                ]),
213                Value::test_list(vec![
214                    Value::test_int(1),
215                    Value::test_int(3),
216                    Value::test_int(2)
217                ]),
218                Value::test_list(vec![
219                    Value::test_int(2),
220                    Value::test_int(3),
221                    Value::test_int(1)
222                ]),
223                Value::test_list(vec![
224                    Value::test_int(3),
225                    Value::test_int(2),
226                    Value::test_int(1)
227                ]),
228            ])
229        );
230        Ok(())
231    }
232
233    #[test]
234    fn permutations_n2() -> Result {
235        let result: Value = test().run("[1 2] | permutations")?;
236        assert_eq!(
237            result,
238            Value::test_list(vec![
239                Value::test_list(vec![Value::test_int(1), Value::test_int(2)]),
240                Value::test_list(vec![Value::test_int(2), Value::test_int(1)]),
241            ])
242        );
243        Ok(())
244    }
245
246    #[test]
247    fn permutations_n1() -> Result {
248        let result: Value = test().run("[42] | permutations")?;
249        assert_eq!(
250            result,
251            Value::test_list(vec![Value::test_list(vec![Value::test_int(42)])])
252        );
253        Ok(())
254    }
255
256    #[test]
257    fn permutations_empty() -> Result {
258        test()
259            .run("[] | permutations")
260            .expect_value_eq(Value::test_list(vec![]))
261    }
262
263    #[test]
264    fn permutations_n4_count() -> Result {
265        // 4! = 24 permutations
266        let result: Value = test().run("[1 2 3 4] | permutations | length")?;
267        assert_eq!(result, Value::test_int(24));
268        Ok(())
269    }
270
271    #[test]
272    fn permutations_strings() -> Result {
273        let result: Value = test().run("[a b] | permutations | get 0 | str join '-'")?;
274        assert_eq!(result, Value::test_string("a-b"));
275        Ok(())
276    }
277}