Omega_h_future.hpp

#ifndef OMEGA_H_Future_HPP
#define OMEGA_H_Future_HPP
 
#include <functional>
#include <vector>
 
#include <Omega_h_mpi.h>
#include "Omega_h_array.hpp"
#include "Omega_h_fail.hpp"
 
namespace Omega_h {
 
template <typename T>
class Future {
 public:
#if OMEGA_H_MPI_NEEDS_HOST_COPY
  using sendbuf_type = HostRead<T>;
  using recvbuf_type = HostWrite<T>;
#else
  using sendbuf_type = Read<T>;
  using recvbuf_type = Write<T>;
#endif
 
  using callback_type = std::function<Read<T>(recvbuf_type)>;
  using additional_callback_type = std::function<Read<T>(Read<T>)>;
 
#ifdef OMEGA_H_USE_MPI
  using requests_type = std::vector<MPI_Request>;
#else
  using requests_type = std::vector<int>;
#endif  // OMEGA_H_USE_MPI
 
#if OMEGA_H_MPI_NEEDS_HOST_COPY
  Future(HostRead<T> sendbuf, HostWrite<T> recvbuf,
      const requests_type&& requests, const callback_type callback);
#else
  Future(Read<T> sendbuf, Write<T> recvbuf, const requests_type&& requests = {},
      const callback_type callback = [](recvbuf_type buf) -> Read<T> {
        return buf;
      });
 
  explicit Future(Write<T> recvbuf);
 
#endif
  explicit Future(Read<T> buf);
 
  void add_callback(const additional_callback_type& callback) {
    auto cur_cb = callback_;
    callback_ = [=](recvbuf_type recvbuf) -> Read<T> {
      return callback(cur_cb(recvbuf));
    };
  }
 
  bool completed();
 
  Read<T> get();
 
 private:
  enum class Status {
    waiting,    
    completed,  
    consumed    
  };
 
  sendbuf_type sendbuf_;
  recvbuf_type recvbuf_;
  callback_type callback_;
  requests_type requests_;
  Status status_;
};
 
#if OMEGA_H_MPI_NEEDS_HOST_COPY
#define OMEGA_H_EXPL_INST_DECL(T)                                              \
  extern template Future<T>::Future(HostRead<T> sendbuf, HostWrite<T> recvbuf, \
      const requests_type&&, const callback_type callback);
#else
#define OMEGA_H_EXPL_INST_DECL(T)                                              \
  extern template Future<T>::Future(Read<T> sendbuf, Write<T> recvbuf,         \
      const requests_type&& requests, const callback_type callback);           \
  extern template Future<T>::Future(Write<T> recvbuf);
#endif
OMEGA_H_EXPL_INST_DECL(I8)
OMEGA_H_EXPL_INST_DECL(I32)
OMEGA_H_EXPL_INST_DECL(I64)
OMEGA_H_EXPL_INST_DECL(Real)
#undef OMEGA_H_EXPL_INST_DECL
 
#define OMEGA_H_EXPL_INST_DECL(T)                                              \
  extern template Future<T>::Future(Read<T> buf);                              \
  extern template bool Future<T>::completed();                                 \
  extern template Read<T> Future<T>::get();
 
OMEGA_H_EXPL_INST_DECL(I8)
OMEGA_H_EXPL_INST_DECL(I32)
OMEGA_H_EXPL_INST_DECL(I64)
OMEGA_H_EXPL_INST_DECL(Real)
#undef OMEGA_H_EXPL_INST_DECL
 
}  // namespace Omega_h
 
#endif  // OMEGA_H_Future_HPP