n_ary_state.h

Go to the documentation of this file.

#pragma once

#include <cmath>
#include <stdexcept>
#include <type_traits>

#include <Eigen/Dense>

#include "drake/common/drake_assert.h"
#include "drake/core/Vector.h"

namespace drake {

using Drake::toEigen;  // TODO(maddog) ...until Drake-->drake fully.

template <class UnitVector>
class NAryState {
 public:
  // The Drake::Vector concept has no explicit way of recovering the
  // ScalarType upon which the template was specialized, but through
  // the miracle of decltype(), it can be done.  std::decay<> is used
  // to strip any CV or reference qualifiers off the type.
  using UnitScalar =
      typename std::decay<decltype(toEigen(UnitVector())(0))>::type;
  // Required by Drake::Vector concept.
  static const int RowsAtCompileTime = Eigen::Dynamic;
  // Required by Drake::Vector concept.
  typedef Eigen::Matrix<UnitScalar, RowsAtCompileTime, 1> EigenType;

  NAryState() : unit_size_(unit_size()), count_(UnitCountFromRows(0)) {}

  explicit NAryState(int count)
      : unit_size_(unit_size()),
        // Ensure correct internal count_ (i.e., -1 if UnitVector's size
        // is zero).
        count_(UnitCountFromRows(RowsFromUnitCount(count))),
        combined_vector_(EigenType::Constant(RowsFromUnitCount(count), 1,
                                             NAN)) {}

  int count() const { return count_; }

  void Append(const UnitVector& unit) {
    if (unit_size_ == 0) {
      // NOP --- in particular, count_ should remain at -1.
      DRAKE_ASSERT(count_ == -1);
      return;
    }
    // Enlarge combined_vector_ by size of the converted unit.
    combined_vector_.conservativeResize(combined_vector_.rows() + unit_size_,
                                        Eigen::NoChange);
    // Copy unit's Eigen-rep into the tail-end of enlarged combined_vector_.
    combined_vector_.bottomRows(unit_size_) = toEigen(unit);
    // Keep track of the total unit count.
    ++count_;
  }

  UnitVector get(int pos) const {
    if (unit_size_ > 0) {
      if ((pos < 0) || (pos >= count_)) {
        throw std::out_of_range("Position pos exceeds range [0, count()].");
      }
    }
    const int row0 = pos * unit_size_;
    return UnitVector(combined_vector_.block(row0, 0,
                                             unit_size_, 1));
  }

  void set(int pos, const UnitVector& unit) {
    if (unit_size_ > 0) {
      if ((pos < 0) || (pos >= count_)) {
        throw std::out_of_range("Position pos exceeds range [0, count()].");
      }
    }
    const int row0 = pos * unit_size_;
    combined_vector_.block(row0, 0, unit_size_, 1) = toEigen(unit);
  }

  // Required by Drake::Vector concept.
  template <typename Derived>
  // NOLINTNEXTLINE(runtime/explicit)
  NAryState(const Eigen::MatrixBase<Derived>& initial)
      : unit_size_(unit_size()),
        count_(UnitCountFromRows(initial.rows())),
        combined_vector_(initial) {
  }

  // Required by Drake::Vector concept.
  template <typename Derived>
  NAryState& operator=(const Eigen::MatrixBase<Derived>& rhs) {
    count_ = UnitCountFromRows(rhs.rows());
    combined_vector_ = rhs;
    return *this;
  }

  // Required by Drake::Vector concept.
  int size() const { return combined_vector_.rows(); }

  // Required by Drake::Vector concept.
  friend EigenType toEigen(const NAryState<UnitVector>& vec) {
    return vec.combined_vector_;
  }

  static int unit_size() { return UnitVector().size(); }

  static int UnitCountFromRows(int rows) {
    const int us { unit_size() };
    if (us > 0) {
      if ((rows % us) != 0) {
        throw std::domain_error("Row count not a multiple of non-null unit.");
      }
      return rows / us;
    }
    return -1;
  }

  static int RowsFromUnitCount(int count) {
    const int us { unit_size() };
    if (count >= 0) {
      return count * us;
    }
    if (us != 0) {
      throw std::domain_error("Negative count for non-null unit.");
    }
    return 0;
  }

 private:
  int unit_size_;
  // count_ < 0 indicates "not counted", i.e., UnitVector is a null vector.
  int count_;
  EigenType combined_vector_;
};

}  // namespace drake