CwiseTernaryOp.h

// This file is part of Eigen, a lightweight C++ template library
// for linear algebra.
//
// Copyright (C) 2008-2014 Gael Guennebaud <gael.guennebaud@inria.fr>
// Copyright (C) 2006-2008 Benoit Jacob <jacob.benoit.1@gmail.com>
// Copyright (C) 2016 Eugene Brevdo <ebrevdo@gmail.com>
//
// This Source Code Form is subject to the terms of the Mozilla
// 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/.
// SPDX-License-Identifier: MPL-2.0
 
#ifndef EIGEN_CWISE_TERNARY_OP_H
#define EIGEN_CWISE_TERNARY_OP_H
 
// IWYU pragma: private
#include "./InternalHeaderCheck.h"
 
namespace Eigen {
 
namespace internal {
template <typename TernaryOp, typename Arg1, typename Arg2, typename Arg3>
struct traits<CwiseTernaryOp<TernaryOp, Arg1, Arg2, Arg3>> {
  // we must not inherit from traits<Arg1> since it has
  // the potential to cause problems with MSVC
  using Ancestor = remove_all_t<Arg1>;
  using XprKind = typename traits<Ancestor>::XprKind;
  enum {
    RowsAtCompileTime = traits<Ancestor>::RowsAtCompileTime,
    ColsAtCompileTime = traits<Ancestor>::ColsAtCompileTime,
    MaxRowsAtCompileTime = traits<Ancestor>::MaxRowsAtCompileTime,
    MaxColsAtCompileTime = traits<Ancestor>::MaxColsAtCompileTime
  };
 
  // even though we require Arg1, Arg2, and Arg3 to have the same scalar type
  // (see CwiseTernaryOp constructor),
  // we still want to handle the case when the result type is different.
  using Scalar = typename result_of<TernaryOp(const typename Arg1::Scalar&, const typename Arg2::Scalar&,
                                              const typename Arg3::Scalar&)>::type;
 
  using StorageKind = typename internal::traits<Arg1>::StorageKind;
  using StorageIndex = typename internal::traits<Arg1>::StorageIndex;
 
  using Arg1Nested = typename Arg1::Nested;
  using Arg2Nested = typename Arg2::Nested;
  using Arg3Nested = typename Arg3::Nested;
  using Arg1Nested_ = std::remove_reference_t<Arg1Nested>;
  using Arg2Nested_ = std::remove_reference_t<Arg2Nested>;
  using Arg3Nested_ = std::remove_reference_t<Arg3Nested>;
  enum { Flags = Arg1Nested_::Flags & RowMajorBit };
};
}  // end namespace internal
 
template <typename TernaryOp, typename Arg1, typename Arg2, typename Arg3, typename StorageKind>
class CwiseTernaryOpImpl;

template <typename TernaryOp, typename Arg1Type, typename Arg2Type, typename Arg3Type>
class CwiseTernaryOp : public CwiseTernaryOpImpl<TernaryOp, Arg1Type, Arg2Type, Arg3Type,
                                                 typename internal::traits<Arg1Type>::StorageKind>,
                       internal::no_assignment_operator {
 public:
  using Arg1 = internal::remove_all_t<Arg1Type>;
  using Arg2 = internal::remove_all_t<Arg2Type>;
  using Arg3 = internal::remove_all_t<Arg3Type>;
 
  // require the sizes to match
  EIGEN_STATIC_ASSERT_SAME_MATRIX_SIZE(Arg1, Arg2)
  EIGEN_STATIC_ASSERT_SAME_MATRIX_SIZE(Arg1, Arg3)
 
  // The index types should match
  EIGEN_STATIC_ASSERT((std::is_same<typename internal::traits<Arg1Type>::StorageKind,
                                    typename internal::traits<Arg2Type>::StorageKind>::value),
                      STORAGE_KIND_MUST_MATCH)
  EIGEN_STATIC_ASSERT((std::is_same<typename internal::traits<Arg1Type>::StorageKind,
                                    typename internal::traits<Arg3Type>::StorageKind>::value),
                      STORAGE_KIND_MUST_MATCH)
 
  using Base = typename CwiseTernaryOpImpl<TernaryOp, Arg1Type, Arg2Type, Arg3Type,
                                           typename internal::traits<Arg1Type>::StorageKind>::Base;
  EIGEN_GENERIC_PUBLIC_INTERFACE(CwiseTernaryOp)
 
  using Arg1Nested = typename internal::ref_selector<Arg1Type>::type;
  using Arg2Nested = typename internal::ref_selector<Arg2Type>::type;
  using Arg3Nested = typename internal::ref_selector<Arg3Type>::type;
  using Arg1Nested_ = std::remove_reference_t<Arg1Nested>;
  using Arg2Nested_ = std::remove_reference_t<Arg2Nested>;
  using Arg3Nested_ = std::remove_reference_t<Arg3Nested>;
 
  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE CwiseTernaryOp(const Arg1& a1, const Arg2& a2, const Arg3& a3,
                                                       const TernaryOp& func = TernaryOp())
      : m_arg1(a1), m_arg2(a2), m_arg3(a3), m_functor(func) {
    eigen_assert(a1.rows() == a2.rows() && a1.cols() == a2.cols() && a1.rows() == a3.rows() && a1.cols() == a3.cols());
  }
 
  EIGEN_DEVICE_FUNC constexpr EIGEN_STRONG_INLINE Index rows() const {
    // return the fixed size type if available to enable compile time
    // optimizations
    EIGEN_IF_CONSTEXPR (internal::traits<internal::remove_all_t<Arg1Nested>>::RowsAtCompileTime == Dynamic &&
                        internal::traits<internal::remove_all_t<Arg2Nested>>::RowsAtCompileTime == Dynamic) {
      return m_arg3.rows();
    } else EIGEN_IF_CONSTEXPR (internal::traits<internal::remove_all_t<Arg1Nested>>::RowsAtCompileTime == Dynamic &&
                               internal::traits<internal::remove_all_t<Arg3Nested>>::RowsAtCompileTime == Dynamic) {
      return m_arg2.rows();
    } else {
      return m_arg1.rows();
    }
  }
  EIGEN_DEVICE_FUNC constexpr EIGEN_STRONG_INLINE Index cols() const {
    // return the fixed size type if available to enable compile time
    // optimizations
    EIGEN_IF_CONSTEXPR (internal::traits<internal::remove_all_t<Arg1Nested>>::ColsAtCompileTime == Dynamic &&
                        internal::traits<internal::remove_all_t<Arg2Nested>>::ColsAtCompileTime == Dynamic) {
      return m_arg3.cols();
    } else EIGEN_IF_CONSTEXPR (internal::traits<internal::remove_all_t<Arg1Nested>>::ColsAtCompileTime == Dynamic &&
                               internal::traits<internal::remove_all_t<Arg3Nested>>::ColsAtCompileTime == Dynamic) {
      return m_arg2.cols();
    } else {
      return m_arg1.cols();
    }
  }

  EIGEN_DEVICE_FUNC constexpr const Arg1Nested_& arg1() const { return m_arg1; }
  EIGEN_DEVICE_FUNC constexpr const Arg2Nested_& arg2() const { return m_arg2; }
  EIGEN_DEVICE_FUNC constexpr const Arg3Nested_& arg3() const { return m_arg3; }
  EIGEN_DEVICE_FUNC constexpr const TernaryOp& functor() const { return m_functor; }
 
 protected:
  Arg1Nested m_arg1;
  Arg2Nested m_arg2;
  Arg3Nested m_arg3;
  const TernaryOp m_functor;
};
 
// Generic API dispatcher
template <typename TernaryOp, typename Arg1, typename Arg2, typename Arg3, typename StorageKind>
class CwiseTernaryOpImpl : public internal::generic_xpr_base<CwiseTernaryOp<TernaryOp, Arg1, Arg2, Arg3>>::type {
 public:
  using Base = typename internal::generic_xpr_base<CwiseTernaryOp<TernaryOp, Arg1, Arg2, Arg3>>::type;
};
 
}  // end namespace Eigen
 
#endif  // EIGEN_CWISE_TERNARY_OP_H