map.hpp

Go to the documentation of this file.

// Copyright (c) 2019-2024 Simons Foundation
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0.txt
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
// Authors: Thomas Hahn, Olivier Parcollet, Nils Wentzell
 
#pragma once
 
#include "./concepts.hpp"
#include "./layout/range.hpp"
#include "./macros.hpp"
#include "./traits.hpp"
 
#include <cstddef>
#include <utility>
#include <tuple>
 
namespace nda {
 
  // Forward declarations.
  template <typename F, Array... A>
  struct expr_call;
 
  template <class F>
  struct mapped;
 
  namespace detail {
 
    // Implementation of the nda::get_algebra trait for function call expressions.
    template <typename... Char>
    constexpr char _impl_find_common_algebra(char x0, Char... x) {
      return (((x == x0) && ...) ? x0 : 'N');
    }
 
  } // namespace detail
 
  template <typename F, Array... As>
  constexpr char get_algebra<expr_call<F, As...>> = detail::_impl_find_common_algebra(get_algebra<As>...);
 
  template <typename F, Array... As>
  struct expr_call {
    F f;
 
    std::tuple<const As...> a;
 
    private:
    // Implementation of the function call operator.
    template <size_t... Is, typename... Args>
    [[gnu::always_inline]] [[nodiscard]] auto _call(std::index_sequence<Is...>, Args const &...args) const {
      // if args contains a range, we need to return an expr_call on the resulting slice
      if constexpr ((is_range_or_ellipsis<Args> or ... or false)) {
        return mapped<F>{f}(std::get<Is>(a)(args...)...);
      } else {
        return f(std::get<Is>(a)(args...)...);
      }
    }
 
    // Implementation of the subscript operator.
    template <size_t... Is, typename Arg>
    [[gnu::always_inline]] auto _call_bra(std::index_sequence<Is...>, Arg const &arg) const {
      return f(std::get<Is>(a)[arg]...);
    }
 
    public:
    template <typename... Args>
    auto operator()(Args const &...args) const {
      return _call(std::make_index_sequence<sizeof...(As)>{}, args...);
    }
 
    template <typename Arg>
    auto operator[](Arg const &arg) const {
      return _call_bra(std::make_index_sequence<sizeof...(As)>{}, arg);
    }
 
    // FIXME copy needed for the && case only. Overload ?
    [[nodiscard]] auto shape() const { return std::get<0>(a).shape(); }
 
    [[nodiscard]] long size() const { return std::get<0>(a).size(); }
  };
 
  template <class F>
  struct mapped {
    F f;
 
    template <Array A0, Array... As>
    expr_call<F, A0, As...> operator()(A0 &&a0, As &&...as) const {
      EXPECTS(((as.shape() == a0.shape()) && ...)); // same shape
      return {f, {std::forward<A0>(a0), std::forward<As>(as)...}};
    }
 
    template <Scalar T0, Scalar... Ts>
    auto operator()(T0 t0, Ts... ts) const {
      return f(t0, ts...);
    }
  };
 
  template <class F>
  mapped<F> map(F f) {
    return {std::move(f)};
  }
 
} // namespace nda