From 15a63cb8927f729ee4f9ea4f356cdb6a118b6f61 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Violeta=20Hern=C3=A1ndez?= <eric.ivan.hdz@gmail.com>
Date: Fri, 9 Apr 2021 23:43:59 -0500
Subject: [PATCH] Memory improvements, extra comments.

The result of `multiplier ^ 2` is now written into a single buffer.
---
 src/linalg/pow.rs | 49 +++++++++++++++++++++++++++++++++++++----------
 1 file changed, 39 insertions(+), 10 deletions(-)

diff --git a/src/linalg/pow.rs b/src/linalg/pow.rs
index 27cab313..5ed3ae7f 100644
--- a/src/linalg/pow.rs
+++ b/src/linalg/pow.rs
@@ -10,18 +10,38 @@ impl<N: ComplexField, D> MatrixN<N, D>
 where
     D: DimMin<D, Output = D>,
     DefaultAllocator: Allocator<N, D, D>,
+    DefaultAllocator: Allocator<N, D>,
 {
-    /// Attempts to raise this matrix to an integer power in-place. Returns
-    /// `false` and leaves `self` untouched if the power is negative and the
-    /// matrix is non-invertible.
+    /// Computes the square of this matrix and writes it into a given buffer.
+    fn square_buf(&mut self, buf: &mut Self) {
+        // We unroll the first iteration to avoid new_uninitialized.
+        let mut aux_col = self.column(0).clone_owned();
+        aux_col = &*self * aux_col;
+        buf.column_mut(0).copy_from(&aux_col);
+
+        // We multiply the matrix by its i-th column,
+        for i in 1..self.ncols() {
+            aux_col.copy_from(&self.column(i));
+            aux_col = &*self * aux_col;
+            self.column_mut(i).copy_from(&aux_col);
+        }
+    }
+
+    /// Attempts to raise this matrix to an integral power `e` in-place. If this
+    /// matrix is non-invertible and `e` is negative, it leaves this matrix
+    /// untouched and returns `false`. Otherwise, it returns `true` and
+    /// overwrites this matrix with the result.
     pub fn pow_mut<T: PrimInt + DivAssign>(&mut self, mut e: T) -> bool {
         let zero = T::zero();
 
+        // A matrix raised to the zeroth power is just the identity.
         if e == zero {
             self.fill_with_identity();
             return true;
         }
 
+        // If e is negative, we compute the inverse matrix, then raise it to the
+        // power of -e.
         if e < zero {
             if !self.try_inverse_mut() {
                 return false;
@@ -30,22 +50,31 @@ where
 
         let one = T::one();
         let two = T::from(2u8).unwrap();
-        let mut multiplier = self.clone();
 
-        while e != zero {
+        // We use the buffer to hold the result of multiplier ^ 2, thus avoiding
+        // extra allocations.
+        let mut multiplier = self.clone();
+        let mut buf = self.clone();
+
+        // Exponentiation by squares.
+        loop {
             if e % two == one {
                 *self *= &multiplier;
             }
 
             e /= two;
-            multiplier *= multiplier.clone();
-        }
+            multiplier.square_buf(&mut buf);
+            multiplier.copy_from(&buf);
 
-        true
+            if e == zero {
+                return true;
+            }
+        }
     }
 
-    /// Raise this matrix to an integer power. Returns `None` only if the power
-    /// is negative and the matrix is non-invertible.
+    /// Attempts to raise this matrix to an integral power `e`. If this matrix
+    /// is non-invertible and `e` is negative, it returns `None`. Otherwise, it
+    /// returns the result as a new matrix. Uses exponentiation by squares.
     pub fn pow<T: PrimInt + DivAssign>(&self, e: T) -> Option<Self> {
         let mut clone = self.clone();