holder.hpp

Go to the documentation of this file.

#ifndef STAN_MATH_PRIM_FUN_HOLDER_HPP
#define STAN_MATH_PRIM_FUN_HOLDER_HPP
 
#include <stan/math/prim/meta/is_eigen.hpp>
#include <stan/math/prim/meta/is_plain_type.hpp>
#include <stan/math/prim/fun/Eigen.hpp>
#include <memory>
#include <type_traits>
#include <utility>
// This was implemented following the tutorial on edding new expressions to
// Eigen: https://eigen.tuxfamily.org/dox/TopicNewExpressionType.html
 
namespace stan {
namespace math {
 
template <class ArgType, typename... Ptrs>
class Holder;
 
}  // namespace math
}  // namespace stan
 
namespace Eigen {
namespace internal {
 
template <class ArgType, typename... Ptrs>
struct traits<stan::math::Holder<ArgType, Ptrs...>> {
  typedef typename ArgType::StorageKind StorageKind;
  typedef typename traits<ArgType>::XprKind XprKind;
  typedef typename ArgType::StorageIndex StorageIndex;
  typedef typename ArgType::Scalar Scalar;
  enum {
    // Possible flags are documented here:
    // https://eigen.tuxfamily.org/dox/group__flags.html
    Flags = (ArgType::Flags
             & (RowMajorBit | LvalueBit | LinearAccessBit | DirectAccessBit
                | PacketAccessBit | NoPreferredStorageOrderBit))
            | NestByRefBit,
    RowsAtCompileTime = ArgType::RowsAtCompileTime,
    ColsAtCompileTime = ArgType::ColsAtCompileTime,
    MaxRowsAtCompileTime = ArgType::MaxRowsAtCompileTime,
    MaxColsAtCompileTime = ArgType::MaxColsAtCompileTime,
    InnerStrideAtCompileTime = ArgType::InnerStrideAtCompileTime,
    OuterStrideAtCompileTime = ArgType::OuterStrideAtCompileTime
  };
};
 
}  // namespace internal
}  // namespace Eigen
 
namespace stan {
namespace internal {
template <typename T>
struct is_holder : std::false_type {};
template <typename ArgType, typename... Ptrs>
struct is_holder<stan::math::Holder<ArgType, Ptrs...>> : std::true_type {};
}  // namespace internal
 
template <typename T>
struct is_holder : internal::is_holder<std::decay_t<T>> {};
 
template <typename T>
inline constexpr bool is_holder_v = is_holder<T>::value;
 
template <typename T>
using require_holder_t = require_t<is_holder<T>>;
 
namespace math {
 
template <
    typename F, typename... Args,
    require_not_plain_type_t<std::invoke_result_t<F, Args&&...>>* = nullptr>
inline auto make_holder(F&& func, Args&&... args);
 
template <typename F, typename... Args,
          require_plain_type_t<std::invoke_result_t<F, Args&&...>>* = nullptr>
inline auto make_holder(F&& func, Args&&... args);
template <typename ArgType, typename... Ptrs>
class Holder
    : public Eigen::internal::dense_xpr_base<Holder<ArgType, Ptrs...>>::type {
 public:
  typedef typename Eigen::internal::ref_selector<Holder<ArgType, Ptrs...>>::type
      Nested;
  typename Eigen::internal::ref_selector<ArgType>::non_const_type m_arg;
  std::tuple<std::unique_ptr<Ptrs>...> m_unique_ptrs;
  explicit Holder(ArgType&& arg, Ptrs*... pointers)
      : m_arg(std::forward<ArgType>(arg)),
        m_unique_ptrs(std::unique_ptr<Ptrs>(pointers)...) {}
 
  // we need to explicitly default copy and move constructors as we are
  // defining copy and move assignment operators
  Holder(const Holder<ArgType, Ptrs...>&) = default;
  Holder(Holder<ArgType, Ptrs...>&&) = default;
 
  // all these functions just call the same on the argument
  Eigen::Index rows() const { return m_arg.rows(); }
  Eigen::Index cols() const { return m_arg.cols(); }
  Eigen::Index innerStride() const { return m_arg.innerStride(); }
  Eigen::Index outerStride() const { return m_arg.outerStride(); }
  auto* data() { return m_arg.data(); }
  const auto* data() const { return m_arg.data(); }
 
  const auto& coeffRef(Eigen::Index row, Eigen::Index col) const {
    return m_arg.coeffRef(row, col);
  }
  const auto& coeffRef(Eigen::Index index) const {
    return m_arg.coeffRef(index);
  }
 
  template <typename T, require_eigen_t<T>* = nullptr>
  inline Holder<ArgType, Ptrs...>& operator=(const T& other) {
    m_arg = other;
    return *this;
  }
 
  // copy and move assignment operators need to be separately overloaded,
  // otherwise defaults will be used.
  inline Holder<ArgType, Ptrs...>& operator=(
      const Holder<ArgType, Ptrs...>& other) {
    m_arg = other;
    return *this;
  }
  inline Holder<ArgType, Ptrs...>& operator=(Holder<ArgType, Ptrs...>&& other) {
    m_arg = std::move(other.m_arg);
    return *this;
  }
};
 
template <typename T, require_holder_t<T>* = nullptr>
inline auto operator-(T&& h) {
  return make_holder([](auto&& arg) { return -arg; }, std::forward<T>(h).m_arg);
}
template <typename T, require_holder_t<T>* = nullptr>
inline auto operator+(T&& h) {
  return make_holder([](auto&& arg) { return arg; }, std::forward<T>(h).m_arg);
}
 
template <typename T, typename Other, require_holder_t<T>* = nullptr,
          require_holder_t<Other>* = nullptr>
inline auto operator-(T&& h, Other&& other) {
  return make_holder(
      [](auto&& arg, auto&& other_) {
        return arg - std::forward<decltype(other_)>(other_);
      },
      std::forward<T>(h).m_arg, std::forward<Other>(other));
}
template <typename T, typename Other, require_holder_t<T>* = nullptr,
          require_holder_t<Other>* = nullptr>
inline auto operator+(T&& h, Other&& other) {
  return make_holder(
      [](auto&& arg, auto&& other_) {
        return arg + std::forward<decltype(other_)>(other_);
      },
      std::forward<T>(h).m_arg, std::forward<Other>(other));
}
template <typename T, typename Other, require_holder_t<T>* = nullptr,
          require_holder_t<Other>* = nullptr>
inline auto operator*(T&& h, Other&& other) {
  return make_holder(
      [](auto&& arg, auto&& other_) {
        return arg * std::forward<decltype(other_)>(other_);
      },
      std::forward<T>(h).m_arg, std::forward<Other>(other));
}
template <typename T, typename Other, require_holder_t<T>* = nullptr,
          require_holder_t<Other>* = nullptr>
inline auto operator/(T&& h, Other&& other) {
  return make_holder(
      [](auto&& arg, auto&& other_) {
        return arg / std::forward<decltype(other_)>(other_);
      },
      std::forward<T>(h).m_arg, std::forward<Other>(other));
}
 
}  // namespace math
}  // namespace stan
 
namespace Eigen {
namespace internal {
 
template <typename ArgType, typename... Ptrs>
struct evaluator<stan::math::Holder<ArgType, Ptrs...>>
    : evaluator_base<stan::math::Holder<ArgType, Ptrs...>> {
  using PlainObjectType = stan::math::Holder<ArgType, Ptrs...>;
  using XprType = stan::math::Holder<ArgType, Ptrs...>;
  using ArgTypeNestedCleaned = typename remove_all<ArgType>::type;
  using CoeffReturnType = typename XprType::CoeffReturnType;
  using Scalar = typename XprType::Scalar;
  enum {
    IsRowMajor = XprType::IsRowMajor,
    IsColMajor = !IsRowMajor,
    IsVectorAtCompileTime = XprType::IsVectorAtCompileTime,
    RowsAtCompileTime = XprType::RowsAtCompileTime,
    ColsAtCompileTime = XprType::ColsAtCompileTime,
 
    CoeffReadCost = evaluator<ArgTypeNestedCleaned>::CoeffReadCost,
    Flags = evaluator<ArgTypeNestedCleaned>::Flags,
    Alignment = evaluator<ArgTypeNestedCleaned>::Alignment,
  };
  enum {
    // We do not need to know the outer stride for vectors
    OuterStrideAtCompileTime
    = IsVectorAtCompileTime
          ? 0
          : (IsRowMajor ? ColsAtCompileTime : RowsAtCompileTime)
  };
 
  evaluator<ArgTypeNestedCleaned> m_argImpl;
 
  explicit evaluator(const XprType& xpr) : m_argImpl(xpr.m_arg) {}
  explicit evaluator(XprType&& xpr)
      : m_argImpl(std::forward<XprType>(xpr).m_arg) {}
 
  // all these functions just call the same on the argument
  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE CoeffReturnType coeff(Index row,
                                                              Index col) const {
    return m_argImpl.coeff(row, col);
  }
  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE CoeffReturnType
  coeff(Index index) const {
    return m_argImpl.coeff(index);
  }
 
  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE Scalar& coeffRef(Index row, Index col) {
    return m_argImpl.coeffRef(row, col);
  }
  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE Scalar& coeffRef(Index index) {
    return m_argImpl.coeffRef(index);
  }
 
  template <int LoadMode, typename PacketType>
  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE PacketType packet(Index row,
                                                          Index col) const {
    return m_argImpl.template packet<LoadMode, PacketType>(row, col);
  }
  template <int LoadMode, typename PacketType>
  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE PacketType packet(Index index) const {
    return m_argImpl.template packet<LoadMode, PacketType>(index);
  }
 
  template <int StoreMode, typename PacketType>
  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void writePacket(Index row, Index col,
                                                         const PacketType& x) {
    return m_argImpl.template writePacket<StoreMode, PacketType>(row, col, x);
  }
  template <int StoreMode, typename PacketType>
  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void writePacket(Index index,
                                                         const PacketType& x) {
    return m_argImpl.template writePacket<StoreMode, PacketType>(index, x);
  }
 
  template <int LoadMode, typename PacketType>
  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE PacketType
  packetSegment(Index row, Index col, Index begin, Index count) const {
    return m_argImpl.template packetSegment<LoadMode, PacketType>(row, col,
                                                                  begin, count);
  }
 
  template <int LoadMode, typename PacketType>
  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE PacketType
  packetSegment(Index index, Index begin, Index count) const {
    return m_argImpl.template packetSegment<LoadMode, PacketType>(index, begin,
                                                                  count);
  }
 
  template <int StoreMode, typename PacketType>
  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void writePacketSegment(
      Index row, Index col, const PacketType& x, Index begin, Index count) {
    return m_argImpl.template writePacketSegment<StoreMode, PacketType>(
        row, col, x, begin, count);
  }
 
  template <int StoreMode, typename PacketType>
  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void writePacketSegment(
      Index index, const PacketType& x, Index begin, Index count) {
    return m_argImpl.template writePacketSegment<StoreMode, PacketType>(
        index, x, begin, count);
  }
};
 
}  // namespace internal
}  // namespace Eigen
 
namespace stan {
namespace math {
 
template <typename T, typename... Ptrs,
          std::enable_if_t<sizeof...(Ptrs) >= 1>* = nullptr>
inline Holder<T, Ptrs...> holder(T&& arg, Ptrs*... pointers) {
  return Holder<T, Ptrs...>(std::forward<T>(arg), pointers...);
}
// trivial case with no pointers constructs no holder object
template <typename T>
inline decltype(auto) holder(T&& arg) {
  if constexpr (std::is_rvalue_reference<T&&>::value) {
    return std::decay_t<T>(std::forward<T>(arg));
  } else {
    return std::forward<T>(arg);
  }
}
 
namespace internal {
// the function holder_handle_element is also used in holder_cl
template <typename T>
inline auto holder_handle_element(T& a, T*& res) {
  res = &a;
  return std::make_tuple();
}
template <typename T,
          std::enable_if_t<!(Eigen::internal::traits<std::decay_t<T>>::Flags
                             & Eigen::NestByRefBit)>* = nullptr>
inline auto holder_handle_element(T&& a, std::remove_reference_t<T>*& res) {
  res = &a;
  return std::make_tuple();
}
 
template <typename T, require_t<std::is_rvalue_reference<T&&>>* = nullptr,
          std::enable_if_t<
              static_cast<bool>(Eigen::internal::traits<std::decay_t<T>>::Flags&
                                    Eigen::NestByRefBit)>* = nullptr>
inline auto holder_handle_element(T&& a, T*& res) {
  res = new T(std::move(a));
  return std::make_tuple(res);
}
template <typename T, require_t<std::is_rvalue_reference<T&&>>* = nullptr,
          require_not_eigen_t<T>* = nullptr>
inline auto holder_handle_element(T&& a, T*& res) {
  res = new T(std::move(a));
  return std::make_tuple(res);
}
 
template <typename T, std::size_t... Is, typename... Args>
inline auto make_holder_impl_construct_object(
    T&& expr, std::index_sequence<Is...>, const std::tuple<Args*...>& ptrs) {
  return holder(std::forward<T>(expr), std::get<Is>(ptrs)...);
}
 
template <typename F, std::size_t... Is, typename... Args>
inline auto make_holder_impl(F&& func, std::index_sequence<Is...>,
                             Args&&... args) {
  std::tuple<std::remove_reference_t<Args>*...> res;
  auto ptrs = std::tuple_cat(
      holder_handle_element(std::forward<Args>(args), std::get<Is>(res))...);
  return make_holder_impl_construct_object(
      std::forward<F>(func)(*std::get<Is>(res)...),
      std::make_index_sequence<std::tuple_size<decltype(ptrs)>::value>(), ptrs);
}
 
}  // namespace internal
 
template <typename F, typename... Args,
          require_not_plain_type_t<std::invoke_result_t<F, Args&&...>>*>
inline auto make_holder(F&& func, Args&&... args) {
  if constexpr (is_var_matrix_v<std::invoke_result_t<F, Args&&...>>) {
    return std::forward<F>(func)(std::forward<Args>(args)...);
  } else {
    return internal::make_holder_impl(
        std::forward<F>(func), std::make_index_sequence<sizeof...(Args)>(),
        std::forward<Args>(args)...);
  }
}
 
template <typename F, typename... Args,
          require_plain_type_t<std::invoke_result_t<F, Args&&...>>*>
inline auto make_holder(F&& func, Args&&... args) {
  return std::forward<F>(func)(std::forward<Args>(args)...);
}
 
}  // namespace math
}  // namespace stan
 
#endif