Skip to main content

nu_command/matrix/
reduce.rs

1use crate::matrix::MatrixValue;
2use nu_engine::ClosureEval;
3use nu_engine::command_prelude::*;
4use nu_protocol::engine::Closure;
5
6#[derive(Clone)]
7pub struct MatrixReduce;
8
9impl Command for MatrixReduce {
10    fn name(&self) -> &str {
11        "matrix reduce"
12    }
13
14    fn signature(&self) -> Signature {
15        Signature::build("matrix reduce")
16            .input_output_types(vec![(Type::Custom("matrix".into()), Type::Any)])
17            .required(
18                "closure",
19                SyntaxShape::Closure(Some(vec![SyntaxShape::Any, SyntaxShape::Number])),
20                "Reducing function. Arguments are the current element, then the accumulator (same order as `reduce`).",
21            )
22            .named(
23                "fold",
24                SyntaxShape::Any,
25                "The initial value for the accumulator",
26                Some('f'),
27            )
28            .category(Category::Filters)
29    }
30
31    fn description(&self) -> &str {
32        "Reduce all elements of a matrix to a single value."
33    }
34
35    fn search_terms(&self) -> Vec<&str> {
36        vec!["fold", "accumulate"]
37    }
38
39    fn run(
40        &self,
41        engine_state: &EngineState,
42        stack: &mut Stack,
43        call: &Call,
44        input: PipelineData,
45    ) -> Result<PipelineData, ShellError> {
46        let head = call.head;
47        let closure: Closure = call.req(engine_state, stack, 0)?;
48        let fold: Option<Value> = call.get_flag(engine_state, stack, "fold")?;
49        let matrix = MatrixValue::from_value(&input.into_value(head)?)?;
50
51        let mut iter = matrix.array.iter();
52        let mut acc = if let Some(fold_val) = fold {
53            fold_val
54        } else {
55            match iter.next() {
56                Some(&first) => Value::float(first, head),
57                None => {
58                    return Err(ShellError::Generic(
59                        nu_protocol::shell_error::generic::GenericError::new(
60                            "Empty matrix",
61                            "cannot reduce an empty matrix without --fold",
62                            head,
63                        ),
64                    ));
65                }
66            }
67        };
68
69        let mut closure_eval = ClosureEval::new(engine_state, stack, closure);
70
71        // Match `reduce`: first arg is the element, second is the accumulator.
72        for &val in iter {
73            engine_state.signals().check(&head)?;
74            let element = Value::float(val, head);
75            acc = closure_eval
76                .add_arg(element)?
77                .add_arg(acc.clone())?
78                .run_with_input(PipelineData::value(acc, None))?
79                .into_value(head)?;
80        }
81
82        Ok(acc.with_span(head).into_pipeline_data())
83    }
84
85    fn examples(&self) -> Vec<Example<'static>> {
86        vec![
87            Example {
88                description: "Sum all elements of a 2x2 matrix",
89                example: "[[1 2] [3 4]] | into matrix | matrix reduce --fold 0.0 {|e, acc| $acc + $e}",
90                result: Some(Value::test_float(10.0)),
91            },
92            Example {
93                description: "Product of all elements in a matrix",
94                example: "[[1 2] [3 4]] | into matrix | matrix reduce --fold 1.0 {|e, acc| $acc * $e}",
95                result: Some(Value::test_float(24.0)),
96            },
97            Example {
98                description: "Sum without an initial value (uses first element as starting accumulator)",
99                example: "[[1 2] [3 4]] | into matrix | matrix reduce {|e, acc| $acc + $e}",
100                result: Some(Value::test_float(10.0)),
101            },
102        ]
103    }
104}
105
106#[cfg(test)]
107mod test {
108    use super::*;
109
110    #[test]
111    fn test_examples() -> nu_test_support::Result {
112        nu_test_support::test().examples(MatrixReduce)
113    }
114}