From 4b0926f99b6fdb82b573644827a1fcca71d8566d Mon Sep 17 00:00:00 2001
From: Arthur Feeney <artfe@amazon.com>
Date: Fri, 21 Jan 2022 11:24:36 -0600
Subject: [PATCH] Prevent heap allocation in diagonal product

---
 Eigen/src/Core/DiagonalMatrix.h |  2 +-
 test/diagonalmatrices.cpp       | 17 ++++++++++++++++-
 2 files changed, 17 insertions(+), 2 deletions(-)

diff --git a/Eigen/src/Core/DiagonalMatrix.h b/Eigen/src/Core/DiagonalMatrix.h
index 453f69bc9..6cbea6d30 100644
--- a/Eigen/src/Core/DiagonalMatrix.h
+++ b/Eigen/src/Core/DiagonalMatrix.h
@@ -134,7 +134,7 @@ struct traits<DiagonalMatrix<Scalar_,SizeAtCompileTime,MaxSizeAtCompileTime> >
   typedef Matrix<Scalar_,SizeAtCompileTime,1,0,MaxSizeAtCompileTime,1> DiagonalVectorType;
   typedef DiagonalShape StorageKind;
   enum {
-    Flags = LvalueBit | NoPreferredStorageOrderBit
+    Flags = LvalueBit | NoPreferredStorageOrderBit | NestByRefBit
   };
 };
 }
diff --git a/test/diagonalmatrices.cpp b/test/diagonalmatrices.cpp
index 276beade0..a88c3c124 100644
--- a/test/diagonalmatrices.cpp
+++ b/test/diagonalmatrices.cpp
@@ -7,6 +7,12 @@
 // Public License v. 2.0. If a copy of the MPL was not distributed
 // with this file, You can obtain one at http://mozilla.org/MPL/2.0/.
 
+
+// discard stack allocation as that too bypasses malloc
+#define EIGEN_STACK_ALLOCATION_LIMIT 0
+// heap allocation will raise an assert if enabled at runtime
+#define EIGEN_RUNTIME_NO_MALLOC
+
 #include "main.h"
 using namespace std;
 template<typename MatrixType> void diagonalmatrices(const MatrixType& m)
@@ -56,6 +62,7 @@ template<typename MatrixType> void diagonalmatrices(const MatrixType& m)
   Index i = internal::random<Index>(0, rows-1);
   Index j = internal::random<Index>(0, cols-1);
   
+  internal::set_is_malloc_allowed(false);
   VERIFY_IS_APPROX( ((ldm1 * m1)(i,j))  , ldm1.diagonal()(i) * m1(i,j) );
   VERIFY_IS_APPROX( ((ldm1 * (m1+m2))(i,j))  , ldm1.diagonal()(i) * (m1+m2)(i,j) );
   VERIFY_IS_APPROX( ((m1 * rdm1)(i,j))  , rdm1.diagonal()(j) * m1(i,j) );
@@ -65,6 +72,7 @@ template<typename MatrixType> void diagonalmatrices(const MatrixType& m)
   VERIFY_IS_APPROX( (((v1+v2).asDiagonal() * (m1+m2))(i,j))  , (v1+v2)(i) * (m1+m2)(i,j) );
   VERIFY_IS_APPROX( ((m1 * (rv1+rv2).asDiagonal())(i,j))  , (rv1+rv2)(j) * m1(i,j) );
   VERIFY_IS_APPROX( (((m1+m2) * (rv1+rv2).asDiagonal())(i,j))  , (rv1+rv2)(j) * (m1+m2)(i,j) );
+  internal::set_is_malloc_allowed(true);
   
   if(rows>1)
   {
@@ -84,7 +92,14 @@ template<typename MatrixType> void diagonalmatrices(const MatrixType& m)
   big.block(i,j,rows,cols) = m1;
   big.block(i,j,rows,cols) = big.block(i,j,rows,cols) * rv1.asDiagonal();
   VERIFY_IS_APPROX((big.block(i,j,rows,cols)) , m1 * rv1.asDiagonal() );
-  
+
+  // products do not allocate memory
+  MatrixType res(rows, cols);
+  internal::set_is_malloc_allowed(false);
+  res.noalias() = ldm1 * m;
+  res.noalias() = m * rdm1;
+  res.noalias() = ldm1 * m * rdm1;
+  internal::set_is_malloc_allowed(true);  
   
   // scalar multiple
   VERIFY_IS_APPROX(LeftDiagonalMatrix(ldm1*s1).diagonal(), ldm1.diagonal() * s1);