count_vars.hpp

Go to the documentation of this file.

#ifndef STAN_MATH_REV_CORE_COUNT_VARS_HPP
#define STAN_MATH_REV_CORE_COUNT_VARS_HPP
 
#include <stan/math/rev/meta.hpp>
#include <stan/math/rev/core/var.hpp>
#include <stan/math/prim/functor/apply.hpp>
 
#include <utility>
#include <vector>
 
namespace stan {
namespace math {
 
namespace internal {
template <typename VecVar, require_std_vector_vt<is_var, VecVar>* = nullptr,
          typename... Pargs>
inline size_t count_vars_impl(size_t count, VecVar&& x, Pargs&&... args);
 
template <typename VecContainer,
          require_std_vector_st<is_var, VecContainer>* = nullptr,
          require_std_vector_vt<is_container, VecContainer>* = nullptr,
          typename... Pargs>
inline size_t count_vars_impl(size_t count, VecContainer&& x, Pargs&&... args);
 
template <typename EigT, require_eigen_vt<is_var, EigT>* = nullptr,
          typename... Pargs>
inline size_t count_vars_impl(size_t count, EigT&& x, Pargs&&... args);
 
template <typename... Pargs>
inline size_t count_vars_impl(size_t count, const var& x, Pargs&&... args);
 
template <typename Arith, require_arithmetic_t<scalar_type_t<Arith>>* = nullptr,
          typename... Pargs>
inline size_t count_vars_impl(size_t count, Arith& x, Pargs&&... args);
 
template <typename Tuple, require_tuple_t<Tuple>* = nullptr, typename... Pargs>
inline size_t count_vars_impl(size_t count, Tuple&& x, Pargs&&... args);
 
inline size_t count_vars_impl(size_t count);
template <typename VecVar, require_std_vector_vt<is_var, VecVar>*,
          typename... Pargs>
inline size_t count_vars_impl(size_t count, VecVar&& x, Pargs&&... args) {
  return count_vars_impl(count + x.size(), std::forward<Pargs>(args)...);
}
 
template <typename VecContainer, require_std_vector_st<is_var, VecContainer>*,
          require_std_vector_vt<is_container, VecContainer>*, typename... Pargs>
inline size_t count_vars_impl(size_t count, VecContainer&& x, Pargs&&... args) {
  for (auto&& x_iter : x) {
    count = count_vars_impl(count, x_iter);
  }
  return count_vars_impl(count, std::forward<Pargs>(args)...);
}
 
template <typename EigT, require_eigen_vt<is_var, EigT>*, typename... Pargs>
inline size_t count_vars_impl(size_t count, EigT&& x, Pargs&&... args) {
  return count_vars_impl(count + x.size(), std::forward<Pargs>(args)...);
}
 
template <typename... Pargs>
inline size_t count_vars_impl(size_t count, const var& x, Pargs&&... args) {
  return count_vars_impl(count + 1, std::forward<Pargs>(args)...);
}
 
template <typename Arith, require_arithmetic_t<scalar_type_t<Arith>>*,
          typename... Pargs>
inline size_t count_vars_impl(size_t count, Arith& x, Pargs&&... args) {
  return count_vars_impl(count, std::forward<Pargs>(args)...);
}
 
inline size_t count_vars_impl(size_t count, std::basic_ostream<char>*&) {
  return count;
}
 
template <typename Tuple, require_tuple_t<Tuple>*, typename... Pargs>
inline size_t count_vars_impl(size_t count, Tuple&& x, Pargs&&... args) {
  count = stan::math::apply(
      [count](auto&&... tuple_args) {
        return count_vars_impl(count, tuple_args...);
      },
      std::forward<Tuple>(x));
  return count_vars_impl(count, std::forward<Pargs>(args)...);
}
 
inline size_t count_vars_impl(size_t count) { return count; }
 
}  // namespace internal
 
template <typename... Pargs>
inline size_t count_vars(Pargs&&... args) {
  return internal::count_vars_impl(0, std::forward<Pargs>(args)...);
}
 
}  // namespace math
}  // namespace stan
 
#endif