Implement iter::Sum for DMatrix (#552)

Fixes #514.
dimforge · Feb 23, 2019 · 20e9c6f · 20e9c6f
1 parent 9d08fdc
commit 20e9c6f
Showing 1 changed file with 61 additions and 1 deletion.
diff --git a/src/base/ops.rs b/src/base/ops.rs
@@ -11,7 +11,7 @@ use base::allocator::{Allocator, SameShapeAllocator, SameShapeC, SameShapeR};
 use base::constraint::{
     AreMultipliable, DimEq, SameNumberOfColumns, SameNumberOfRows, ShapeConstraint,
 };
-use base::dimension::{Dim, DimMul, DimName, DimProd};
+use base::dimension::{Dim, DimMul, DimName, DimProd, Dynamic};
 use base::storage::{ContiguousStorageMut, Storage, StorageMut};
 use base::{DefaultAllocator, Matrix, MatrixMN, MatrixN, MatrixSum, Scalar};
 
@@ -384,6 +384,36 @@ where
     }
 }
 
+impl<N, C: Dim> iter::Sum for MatrixMN<N, Dynamic, C>
+where
+    N: Scalar + ClosedAdd + Zero,
+    DefaultAllocator: Allocator<N, Dynamic, C>,
+{
+    /// # Example
+    /// ```
+    /// # use nalgebra::DVector;
+    /// assert_eq!(vec![DVector::repeat(3, 1.0f64),
+    ///                 DVector::repeat(3, 1.0f64),
+    ///                 DVector::repeat(3, 1.0f64)].into_iter().sum::<DVector<f64>>(),
+    ///            DVector::repeat(3, 1.0f64) + DVector::repeat(3, 1.0f64) + DVector::repeat(3, 1.0f64));
+    /// ```
+    ///
+    /// # Panics
+    /// Panics if the iterator is empty:
+    /// ```should_panic
+    /// # use std::iter;
+    /// # use nalgebra::DMatrix;
+    /// iter::empty::<DMatrix<f64>>().sum::<DMatrix<f64>>(); // panics!
+    /// ```
+    fn sum<I: Iterator<Item = MatrixMN<N, Dynamic, C>>>(mut iter: I) -> MatrixMN<N, Dynamic, C> {
+        if let Some(first) = iter.next() {
+            iter.fold(first, |acc, x| acc + x)
+        } else {
+            panic!("Cannot compute `sum` of empty iterator.")
+        }
+    }
+}
+
 impl<'a, N, R: DimName, C: DimName> iter::Sum<&'a MatrixMN<N, R, C>> for MatrixMN<N, R, C>
 where
     N: Scalar + ClosedAdd + Zero,
@@ -394,6 +424,36 @@ where
     }
 }
 
+impl<'a, N, C: Dim> iter::Sum<&'a MatrixMN<N, Dynamic, C>> for MatrixMN<N, Dynamic, C>
+where
+    N: Scalar + ClosedAdd + Zero,
+    DefaultAllocator: Allocator<N, Dynamic, C>,
+{
+    /// # Example
+    /// ```
+    /// # use nalgebra::DVector;
+    /// let v = &DVector::repeat(3, 1.0f64);
+    ///
+    /// assert_eq!(vec![v, v, v].into_iter().sum::<DVector<f64>>(),
+    ///            v + v + v);
+    /// ```
+    ///
+    /// # Panics
+    /// Panics if the iterator is empty:
+    /// ```should_panic
+    /// # use std::iter;
+    /// # use nalgebra::DMatrix;
+    /// iter::empty::<&DMatrix<f64>>().sum::<DMatrix<f64>>(); // panics!
+    /// ```
+    fn sum<I: Iterator<Item = &'a MatrixMN<N, Dynamic, C>>>(mut iter: I) -> MatrixMN<N, Dynamic, C> {
+        if let Some(first) = iter.next() {
+            iter.fold(first.clone(), |acc, x| acc + x)
+        } else {
+            panic!("Cannot compute `sum` of empty iterator.")
+        }
+    }
+}
+
 /*
  *
  * Multiplication