Implement evaluator for Diagonal.

3 files changed, 73 insertions, 1 deletions

diff --git a/Eigen/src/Core/CoreEvaluators.h b/Eigen/src/Core/CoreEvaluators.h
index 1c6c06493..1ef82e4be 100644
--- a/Eigen/src/Core/CoreEvaluators.h
+++ b/Eigen/src/Core/CoreEvaluators.h
@@ -881,6 +881,52 @@ protected:
 };
 
 
+// -------------------- Diagonal --------------------
+
+template<typename ArgType, int DiagIndex>
+struct evaluator_impl<Diagonal<ArgType, DiagIndex> >
+{
+  typedef Diagonal<ArgType, DiagIndex> DiagonalType;
+
+  evaluator_impl(const DiagonalType& diagonal) 
+    : m_argImpl(diagonal.nestedExpression()),
+      m_index(diagonal.index())
+  { }
+ 
+  typedef typename DiagonalType::Index Index;
+  typedef typename DiagonalType::Scalar Scalar;
+  typedef typename DiagonalType::CoeffReturnType CoeffReturnType;
+
+  CoeffReturnType coeff(Index row, Index) const
+  {
+    return m_argImpl.coeff(row + rowOffset(), row + colOffset());
+  }
+
+  CoeffReturnType coeff(Index index) const
+  {
+    return m_argImpl.coeff(index + rowOffset(), index + colOffset());
+  }
+
+  Scalar& coeffRef(Index row, Index)
+  {
+    return m_argImpl.coeffRef(row + rowOffset(), row + colOffset());
+  }
+
+  Scalar& coeffRef(Index index)
+  {
+    return m_argImpl.coeffRef(index + rowOffset(), index + colOffset());
+  }
+
+protected:
+  typename evaluator<ArgType>::type m_argImpl;
+  Index m_index; // TODO: Don't use if known at compile time
+
+private:
+  EIGEN_STRONG_INLINE Index rowOffset() const { return m_index>0 ? 0 : -m_index; }
+  EIGEN_STRONG_INLINE Index colOffset() const { return m_index>0 ? m_index : 0; }
+};
+
+
 } // namespace internal
 
 #endif // EIGEN_COREEVALUATORS_H
diff --git a/Eigen/src/Core/Diagonal.h b/Eigen/src/Core/Diagonal.h
index e807a49e4..3c9f517d1 100644
--- a/Eigen/src/Core/Diagonal.h
+++ b/Eigen/src/Core/Diagonal.h
@@ -133,6 +133,17 @@ template<typename MatrixType, int DiagIndex> class Diagonal
       return m_matrix.coeff(index+rowOffset(), index+colOffset());
     }
 
+    const typename internal::remove_all<typename MatrixType::Nested>::type& 
+    nestedExpression() const 
+    {
+      return m_matrix;
+    }
+
+    int index() const
+    {
+      return m_index.value();
+    }
+
   protected:
     const typename MatrixType::Nested m_matrix;
     const internal::variable_if_dynamic<Index, DiagIndex> m_index;
diff --git a/test/evaluators.cpp b/test/evaluators.cpp
index fa545c7a5..5a123f0ad 100644
--- a/test/evaluators.cpp
+++ b/test/evaluators.cpp
@@ -197,7 +197,22 @@ void test_evaluators()
   VERIFY_IS_APPROX_EVALUATOR(arr2, arr1.reverse());
   VERIFY_IS_APPROX_EVALUATOR(arr2, arr1.colwise().reverse());
   VERIFY_IS_APPROX_EVALUATOR(arr2, arr1.rowwise().reverse());
-
   arr2.reverse() = arr1;
   VERIFY_IS_APPROX(arr2, arr1.reverse());
+
+  // test Diagonal
+  VERIFY_IS_APPROX_EVALUATOR(vec1, mat1.diagonal());
+  vec1.resize(5);
+  VERIFY_IS_APPROX_EVALUATOR(vec1, mat1.diagonal(1));
+  VERIFY_IS_APPROX_EVALUATOR(vec1, mat1.diagonal<-1>());
+  vec1.setRandom();
+
+  mat2 = mat1;
+  copy_using_evaluator(mat1.diagonal(1), vec1);
+  mat2.diagonal(1) = vec1;
+  VERIFY_IS_APPROX(mat1, mat2);
+
+  copy_using_evaluator(mat1.diagonal<-1>(), mat1.diagonal(1));
+  mat2.diagonal<-1>() = mat2.diagonal(1);
+  VERIFY_IS_APPROX(mat1, mat2);
 }