Eigen/src/Core/NullaryWrapper.h

#ifndef EIGEN_NULLARYWRAPPER_H
#define EIGEN_NULLARYWRAPPER_H

namespace Eigen {
namespace internal {

template<typename Scalar,typename NullaryOp,
         bool has_nullary = has_nullary_operator<NullaryOp>::value,
         bool has_unary   = has_unary_operator<NullaryOp>::value,
         bool has_binary  = has_binary_operator<NullaryOp>::value>
struct nullary_wrapper
{
  template <typename Index>
  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE Scalar operator()(const NullaryOp& op, Index i, Index j) const { return op(i,j); }
  template <typename Index>
  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE Scalar operator()(const NullaryOp& op, Index i) const { return op(i); }

  template <typename T, typename Index> EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE T packetOp(const NullaryOp& op, Index i, Index j) const { return op.packetOp(i,j); }
  template <typename T, typename Index> EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE T packetOp(const NullaryOp& op, Index i) const { return op.packetOp(i); }
};

template<typename Scalar,typename NullaryOp>
struct nullary_wrapper<Scalar,NullaryOp,true,false,false>
{
  template <typename Index>
  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE Scalar operator()(const NullaryOp& op, Index=0, Index=0) const { return op(); }
  template <typename T, typename Index> EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE T packetOp(const NullaryOp& op, Index=0, Index=0) const { return op.packetOp(); }
};

template<typename Scalar,typename NullaryOp>
struct nullary_wrapper<Scalar,NullaryOp,false,false,true>
{
  template <typename Index>
  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE Scalar operator()(const NullaryOp& op, Index i, Index j=0) const { return op(i,j); }
  template <typename T, typename Index> EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE T packetOp(const NullaryOp& op, Index i, Index j=0) const { return op.packetOp(i,j); }
};

// We need the following specialization for vector-only functors assigned to a runtime vector,
// for instance, using linspace and assigning a RowVectorXd to a MatrixXd or even a row of a MatrixXd.
// In this case, i==0 and j is used for the actual iteration.
template<typename Scalar,typename NullaryOp>
struct nullary_wrapper<Scalar,NullaryOp,false,true,false>
  : nullary_wrapper<Scalar,NullaryOp,false,true,true> // to get the identity wrapper
{
  typedef nullary_wrapper<Scalar,NullaryOp,false,true,true> base;
  using base::operator();
  using base::packetOp;
  template <typename Index>
  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE Scalar operator()(const NullaryOp& op, Index i, Index j) const {
    eigen_assert(i==0 || j==0);
    return op(i+j);
  }
  template <typename T, typename Index> EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE T packetOp(const NullaryOp& op, Index i, Index j) const {
    eigen_assert(i==0 || j==0);
    return op.packetOp(i+j);
  }
};

template<typename Scalar,typename NullaryOp>
struct nullary_wrapper<Scalar,NullaryOp,false,false,false> {};

}  // namespace Eigen
}  // namespace internal

#endif  // EIGEN_NULLARYWRAPPER_H