rotgen/include/rotgen/functions/extract.hpp

//==================================================================================================
/*
  ROTGEN - Runtime Overlay for Eigen
  Copyright : CODE RECKONS
  SPDX-License-Identifier: BSL-1.0
*/
//==================================================================================================
#pragma once
#include <rotgen/detail/helpers.hpp>

namespace rotgen
{
  namespace detail
  {
    template<concepts::entity Entity>
    void validate_extract([[maybe_unused]] Entity& e,
                          [[maybe_unused]] Index i0,
                          [[maybe_unused]] Index j0,
                          [[maybe_unused]] Index ni,
                          [[maybe_unused]] Index nj)
    {
      assert(i0 >= 0 && "block extraction uses negative row index.");
      assert(j0 >= 0 && "block extraction uses negative col index.");
      assert(i0 + ni <= e.rows() && "block extraction rows is out of range.");
      assert(j0 + nj <= e.cols() && "block extraction cols is out of range.");
    }
  }

  //======================== EXTRACT ========================
  template<concepts::entity Entity>
  auto extract(Entity& e, Index i0, Index j0, Index ni, Index nj)
  {
    detail::validate_extract(e, i0, j0, ni, nj);
    if constexpr (requires { e.as_map(); })
      return extract(e.as_map(), i0, j0, ni, nj);
    else return block<detail::propagate_const<Entity>>(e, i0, j0, ni, nj);
  }

  template<Index NI, Index NJ, concepts::entity Entity>
  requires(NI != -1 && NJ != -1)
  auto extract(Entity& e, Index i0, Index j0)
  {
    detail::validate_extract(e, i0, j0, NI, NJ);
    if constexpr (requires { e.as_map(); })
      return extract<NI, NJ>(e.as_map(), i0, j0);
    else return block<detail::propagate_const<Entity>, NI, NJ>(e, i0, j0);
  }

  template<Index NI, Index NJ, concepts::entity Entity>
  requires((NI != -1) != (NJ != -1))
  auto extract(Entity& e, Index i0, Index j0, Index ni, Index nj)
  {
    detail::validate_extract(e, i0, j0, ni, nj);
    if constexpr (requires { e.as_map(); })
      return extract<NI, NJ>(e.as_map(), i0, j0, ni, nj);
    else
      return block<detail::propagate_const<Entity>, NI, NJ>(e, i0, j0, ni, nj);
  }

  //======================== TOP LEFT CORNER ========================
  template<concepts::entity Entity>
  auto topLeftCorner(Entity& e, Index ni, Index nj)
  {
    return extract(e, 0, 0, ni, nj);
  }

  template<Index NI, Index NJ, concepts::entity Entity>
  auto topLeftCorner(Entity& e)
  {
    return extract<NI, NJ>(e, 0, 0);
  }

  //======================== TOP RIGHT CORNER ========================
  template<concepts::entity Entity>
  auto topRightCorner(Entity& e, Index ni, Index nj)
  {
    return extract(e, 0, e.cols() - nj, ni, nj);
  }

  template<Index NI, Index NJ, concepts::entity Entity>
  auto topRightCorner(Entity& e)
  {
    return extract<NI, NJ>(e, 0, e.cols() - NJ);
  }

  //======================== BOTTOM LEFT CORNER ========================
  template<concepts::entity Entity>
  auto bottomLeftCorner(Entity& e, Index ni, Index nj)
  {
    return extract(e, e.rows() - ni, 0, ni, nj);
  }

  template<Index NI, Index NJ, concepts::entity Entity>
  auto bottomLeftCorner(Entity& e)
  {
    return extract<NI, NJ>(e, e.rows() - NI, 0);
  }

  //======================== BOTTOM RIGHT CORNER ========================
  template<concepts::entity Entity>
  auto bottomRightCorner(Entity& e, Index ni, Index nj)
  {
    return extract(e, e.rows() - ni, e.cols() - nj, ni, nj);
  }

  template<Index NI, Index NJ, concepts::entity Entity>
  auto bottomRightCorner(Entity& e)
  {
    return extract<NI, NJ>(e, e.rows() - NI, e.cols() - NJ);
  }

  //======================== TOP ROWS ========================
  template<concepts::entity Entity> auto topRows(Entity& e, Index ni)
  {
    if constexpr (Entity::ColsAtCompileTime == -1)
      return extract(e, 0, 0, ni, e.cols());
    else return extract<-1, Entity::ColsAtCompileTime>(e, 0, 0, ni, e.cols());
  }

  template<Index NI, concepts::entity Entity> auto topRows(Entity& e)
  {
    if constexpr (Entity::ColsAtCompileTime == -1)
      return extract<NI, -1>(e, 0, 0, NI, e.cols());
    else return extract<NI, Entity::ColsAtCompileTime>(e, 0, 0);
  }

  //======================== MIDDLE ROWS ========================
  template<concepts::entity Entity>
  auto middleRows(Entity& e, Index i0, Index ni)
  {
    if constexpr (Entity::ColsAtCompileTime == -1)
      return extract(e, i0, 0, ni, e.cols());
    else return extract<-1, Entity::ColsAtCompileTime>(e, i0, 0, ni, e.cols());
  }

  template<Index NI, concepts::entity Entity>
  auto middleRows(Entity& e, Index i0)
  {
    if constexpr (Entity::ColsAtCompileTime == -1)
      return extract<NI, -1>(e, i0, 0, NI, e.cols());
    else return extract<NI, Entity::ColsAtCompileTime>(e, i0, 0);
  }

  //======================== BOTTOM ROWS ========================
  template<concepts::entity Entity> auto bottomRows(Entity& e, Index ni)
  {
    if constexpr (Entity::ColsAtCompileTime == -1)
      return extract(e, e.rows() - ni, 0, ni, e.cols());
    else
      return extract<-1, Entity::ColsAtCompileTime>(e, e.rows() - ni, 0, ni,
                                                    e.cols());
  }

  template<Index NI, concepts::entity Entity> auto bottomRows(Entity& e)
  {
    if constexpr (Entity::ColsAtCompileTime == -1)
      return extract<NI, -1>(e, e.rows() - NI, 0, NI, e.cols());
    else return extract<NI, Entity::ColsAtCompileTime>(e, e.rows() - NI, 0);
  }

  //======================== LEFT COLS ========================
  template<concepts::entity Entity> auto leftCols(Entity& e, Index nj)
  {
    if constexpr (Entity::RowsAtCompileTime == -1)
      return extract(e, 0, 0, e.rows(), nj);
    else return extract<Entity::RowsAtCompileTime, -1>(e, 0, 0, e.rows(), nj);
  }

  template<Index NJ, concepts::entity Entity> auto leftCols(Entity& e)
  {
    if constexpr (Entity::RowsAtCompileTime == -1)
      return extract<-1, NJ>(e, 0, 0, e.rows(), NJ);
    else return extract<Entity::RowsAtCompileTime, NJ>(e, 0, 0);
  }

  //======================== MIDDLE COLS ========================
  template<concepts::entity Entity>
  auto middleCols(Entity& e, Index j0, Index nj)
  {
    if constexpr (Entity::RowsAtCompileTime == -1)
      return extract(e, 0, j0, e.rows(), nj);
    else return extract<Entity::RowsAtCompileTime, -1>(e, 0, j0, e.rows(), nj);
  }

  template<Index NJ, concepts::entity Entity>
  auto middleCols(Entity& e, Index j0)
  {
    if constexpr (Entity::RowsAtCompileTime == -1)
      return extract<-1, NJ>(e, 0, j0, e.rows(), NJ);
    else return extract<Entity::RowsAtCompileTime, NJ>(e, 0, j0);
  }

  //======================== RIGHT COLS ========================
  template<concepts::entity Entity> auto rightCols(Entity& e, Index nj)
  {
    if constexpr (Entity::RowsAtCompileTime == -1)
      return extract(e, 0, e.cols() - nj, e.rows(), nj);
    else
      return extract<Entity::RowsAtCompileTime, -1>(e, 0, e.cols() - nj,
                                                    e.rows(), nj);
    ;
  }

  template<Index NJ, concepts::entity Entity> auto rightCols(Entity& e)
  {
    if constexpr (Entity::RowsAtCompileTime == -1)
      return extract<-1, NJ>(e, 0, e.cols() - NJ, e.rows(), NJ);
    else return extract<Entity::RowsAtCompileTime, NJ>(e, 0, e.cols() - NJ);
  }

  //======================== ROW ========================
  template<concepts::entity Entity> auto row(Entity& e, Index i0)
  {
    if constexpr (Entity::ColsAtCompileTime == -1)
      return extract<1, -1>(e, i0, 0, 1, e.cols());
    else return extract<1, Entity::ColsAtCompileTime>(e, i0, 0);
  }

  //======================== COL ========================
  template<concepts::entity Entity> auto col(Entity& e, Index j0)
  {
    if constexpr (Entity::RowsAtCompileTime == -1)
      return extract<-1, 1>(e, 0, j0, e.rows(), 1);
    else return extract<Entity::RowsAtCompileTime, 1>(e, 0, j0);
  }

  //======================== VECTOR HEAD ========================
  template<concepts::entity Entity>
  auto head(Entity& e, Index n)
  requires(Entity::RowsAtCompileTime == 1 || Entity::ColsAtCompileTime == 1)
  {
    if constexpr (Entity::RowsAtCompileTime == 1)
      return extract<1, Dynamic>(e, 0, 0, 1, n);
    else if constexpr (Entity::ColsAtCompileTime == 1)
      return extract<Dynamic, 1>(e, 0, 0, n, 1);
  }

  template<Index N, concepts::entity Entity>
  auto head(Entity& e)
  requires(Entity::RowsAtCompileTime == 1 || Entity::ColsAtCompileTime == 1)
  {
    if constexpr (Entity::RowsAtCompileTime == 1) return extract<1, N>(e, 0, 0);
    else if constexpr (Entity::ColsAtCompileTime == 1)
      return extract<N, 1>(e, 0, 0);
  }

  //======================== VECTOR TAIL ========================
  template<concepts::entity Entity>
  auto tail(Entity& e, Index n)
  requires(Entity::RowsAtCompileTime == 1 || Entity::ColsAtCompileTime == 1)
  {
    if constexpr (Entity::RowsAtCompileTime == 1)
      return extract<1, Dynamic>(e, 0, e.cols() - n, 1, n);
    else if constexpr (Entity::ColsAtCompileTime == 1)
      return extract<Dynamic, 1>(e, e.rows() - n, 0, n, 1);
  }

  template<Index N, concepts::entity Entity>
  auto tail(Entity& e)
  requires(Entity::RowsAtCompileTime == 1 || Entity::ColsAtCompileTime == 1)
  {
    if constexpr (Entity::RowsAtCompileTime == 1)
      return extract<1, N>(e, 0, e.cols() - N);
    else if constexpr (Entity::ColsAtCompileTime == 1)
      return extract<N, 1>(e, e.rows() - N, 0);
  }

  //======================== VECTOR SEGMENT ========================
  template<concepts::entity Entity>
  auto segment(Entity& e, Index s, Index n)
  requires(Entity::RowsAtCompileTime == 1 || Entity::ColsAtCompileTime == 1)
  {
    if constexpr (Entity::RowsAtCompileTime == 1)
      return extract<1, Dynamic>(e, 0, s, 1, n);
    else if constexpr (Entity::ColsAtCompileTime == 1)
      return extract<Dynamic, 1>(e, s, 0, n, 1);
  }

  template<Index N, concepts::entity Entity>
  auto segment(Entity& e, Index s)
  requires(Entity::RowsAtCompileTime == 1 || Entity::ColsAtCompileTime == 1)
  {
    if constexpr (Entity::RowsAtCompileTime == 1) return extract<1, N>(e, 0, s);
    else if constexpr (Entity::ColsAtCompileTime == 1)
      return extract<N, 1>(e, s, 0);
  }
}