Add sum and product to arrays (#966)

5 files changed, 72 insertions, 0 deletions

diff --git a/docs/src/reference/types.md b/docs/src/reference/types.md
index f5f7e6e4..24756422 100644
--- a/docs/src/reference/types.md
+++ b/docs/src/reference/types.md
@@ -630,6 +630,20 @@ Folds all items into a single value using an accumulator function.
   and one for an item.
 - returns: any
 
+### sum()
+Sums all items (works for any types that can be added).
+
+- default: any (named)
+  If set and the array is empty, sum will return this.
+- returns: any
+
+### product()
+Calculates the product all items (works for any types that can be multiplied)
+
+- default: any (named)
+  If set and the array is empty, sum will return this.
+- returns: any
+
 ### any()
 Whether the given function returns `{true}` for any item in the array.
 
diff --git a/src/eval/array.rs b/src/eval/array.rs
index b04fdab8..1166ce94 100644
--- a/src/eval/array.rs
+++ b/src/eval/array.rs
@@ -24,6 +24,7 @@ macro_rules! __array {
 
 #[doc(inline)]
 pub use crate::__array as array;
+use crate::eval::ops::{add, mul};
 #[doc(hidden)]
 pub use ecow::eco_vec;
 
@@ -199,6 +200,40 @@ impl Array {
         Ok(acc)
     }
 
+    /// Calculates the sum of the array's items
+    pub fn sum(&self, default: Option<Value>, span: Span) -> SourceResult<Value> {
+        let mut acc = self
+            .first()
+            .map(|x| x.clone())
+            .or_else(|_| {
+                default.ok_or_else(|| {
+                    eco_format!("cannot calculate sum of empty array with no default")
+                })
+            })
+            .at(span)?;
+        for i in self.iter().skip(1) {
+            acc = add(acc, i.clone()).at(span)?;
+        }
+        Ok(acc)
+    }
+
+    /// Calculates the product of the array's items
+    pub fn product(&self, default: Option<Value>, span: Span) -> SourceResult<Value> {
+        let mut acc = self
+            .first()
+            .map(|x| x.clone())
+            .or_else(|_| {
+                default.ok_or_else(|| {
+                    eco_format!("cannot calculate product of empty array with no default")
+                })
+            })
+            .at(span)?;
+        for i in self.iter().skip(1) {
+            acc = mul(acc, i.clone()).at(span)?;
+        }
+        Ok(acc)
+    }
+
     /// Whether any item matches.
     pub fn any(&self, vm: &mut Vm, func: Func) -> SourceResult<bool> {
         for item in self.iter() {
diff --git a/src/eval/methods.rs b/src/eval/methods.rs
index 56d1c7b7..29b729cb 100644
--- a/src/eval/methods.rs
+++ b/src/eval/methods.rs
@@ -105,6 +105,8 @@ pub fn call(
             "fold" => {
                 array.fold(vm, args.expect("initial value")?, args.expect("function")?)?
             }
+            "sum" => array.sum(args.named("default")?, span)?,
+            "product" => array.product(args.named("default")?, span)?,
             "any" => Value::Bool(array.any(vm, args.expect("function")?)?),
             "all" => Value::Bool(array.all(vm, args.expect("function")?)?),
             "flatten" => Value::Array(array.flatten()),
diff --git a/tests/ref/compiler/array.png b/tests/ref/compiler/array.png
index a7c52f1a..a96dfe64 100644
--- a/tests/ref/compiler/array.png
+++ b/tests/ref/compiler/array.png
diff --git a/tests/typ/compiler/array.typ b/tests/typ/compiler/array.typ
index c9e85ed7..c52160b0 100644
--- a/tests/typ/compiler/array.typ
+++ b/tests/typ/compiler/array.typ
@@ -167,6 +167,27 @@
 #(1, 2, 3).fold(0, () => none)
 
 ---
+// Test the `sum` method.
+#test(().sum(default: 0), 0)
+#test(().sum(default: []), [])
+#test((1, 2, 3).sum(), 6)
+
+---
+// Error: 2-10 cannot calculate sum of empty array with no default
+#().sum()
+
+---
+// Test the `product` method.
+#test(().product(default: 0), 0)
+#test(().product(default: []), [])
+#test(([ab], 3).product(), [ab]*3)
+#test((1, 2, 3).product(), 6)
+
+---
+// Error: 2-14 cannot calculate product of empty array with no default
+#().product()
+
+---
 // Test the `rev` method.
 #test(range(3).rev(), (2, 1, 0))