map.hpp

Go to the documentation of this file.

// Copyright (c) 2019--present, The Simons Foundation
// This file is part of TRIQS/nda and is licensed under the Apache License, Version 2.0.
// SPDX-License-Identifier: Apache-2.0
// See LICENSE in the root of this distribution for details.

 
#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>
  inline constexpr bool is_expression<expr_call<F, As...>> = true;


  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