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
109struct 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 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 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}