typedefs.h

#pragma once
 
#include <array>
#include <atomic>
#include <cstddef>
#include <cstring>
#include <functional>
#include <limits>
#include <memory>
#include <optional>
#include <stdexcept>
#include <string>
#include <string_view>
#include <unordered_map>
#include <vector>
 
#include <ucp/api/ucp.h>
 
namespace ucxx {
 
class Buffer;
class Request;
class RequestAm;

typedef enum {
  UCXX_LOG_LEVEL_FATAL,       /* Immediate termination */
  UCXX_LOG_LEVEL_ERROR,       /* Error is returned to the user */
  UCXX_LOG_LEVEL_WARN,        /* Something's wrong, but we continue */
  UCXX_LOG_LEVEL_DIAG,        /* Diagnostics, silent adjustments or internal error handling */
  UCXX_LOG_LEVEL_INFO,        /* Information */
  UCXX_LOG_LEVEL_DEBUG,       /* Low-volume debugging */
  UCXX_LOG_LEVEL_TRACE,       /* High-volume debugging */
  UCXX_LOG_LEVEL_TRACE_REQ,   /* Every send/receive request */
  UCXX_LOG_LEVEL_TRACE_DATA,  /* Data sent/received on the transport */
  UCXX_LOG_LEVEL_TRACE_ASYNC, /* Asynchronous progress engine */
  UCXX_LOG_LEVEL_TRACE_FUNC,  /* Function calls */
  UCXX_LOG_LEVEL_TRACE_POLL,  /* Polling functions */
  UCXX_LOG_LEVEL_LAST,        /* Last level barrier, not an actual level */
  UCXX_LOG_LEVEL_PRINT        /* Temporary output */
} ucxx_log_level_t;

enum class TransferDirection { Send = 0, Receive };

enum Tag : ucp_tag_t {};

enum TagMask : ucp_tag_t {};

static constexpr TagMask TagMaskFull{std::numeric_limits<std::underlying_type_t<TagMask>>::max()};

typedef std::unordered_map<std::string, std::string> ConfigMap;

typedef std::function<void(ucs_status_t, std::shared_ptr<void>)> RequestCallbackUserFunction;

typedef std::shared_ptr<void> RequestCallbackUserData;

typedef RequestCallbackUserFunction EndpointCloseCallbackUserFunction;

typedef RequestCallbackUserData EndpointCloseCallbackUserData;

typedef std::function<std::shared_ptr<Buffer>(size_t)> AmAllocatorType;

typedef std::function<void(std::shared_ptr<Request>, ucp_ep_h)> AmReceiverCallbackType;

static constexpr size_t AmReceiverCallbackOwnerMaxLen = 63;

static constexpr size_t AmReceiverCallbackOwnerStorageSize = AmReceiverCallbackOwnerMaxLen + 1;

class AmReceiverCallbackOwnerType {
 public:
  AmReceiverCallbackOwnerType() = default;

  AmReceiverCallbackOwnerType(const char* s)  // NOLINT(runtime/explicit)
  {
    if (s == nullptr) return;
    const size_t len = std::strlen(s);
    if (len > AmReceiverCallbackOwnerMaxLen)
      throw std::invalid_argument(
        "AmReceiverCallbackOwnerType: owner name exceeds "
        "maximum length of " +
        std::to_string(AmReceiverCallbackOwnerMaxLen) + " characters");
    std::memcpy(_data.data(), s, len);
  }

  AmReceiverCallbackOwnerType(const std::string& s)  // NOLINT(runtime/explicit)
    : AmReceiverCallbackOwnerType(s.c_str())
  {
  }

  [[nodiscard]] const char* data() const noexcept { return _data.data(); }

  [[nodiscard]] char* data() noexcept { return _data.data(); }

  static constexpr size_t storageSize() noexcept { return AmReceiverCallbackOwnerStorageSize; }

  bool operator==(const AmReceiverCallbackOwnerType& other) const noexcept
  {
    return _data == other._data;
  }

  bool operator!=(const AmReceiverCallbackOwnerType& other) const noexcept
  {
    return !(*this == other);
  }
 
 private:
  std::array<char, AmReceiverCallbackOwnerStorageSize> _data{};
};

typedef uint64_t AmReceiverCallbackIdType;

class AmReceiverCallbackInfo {
 public:
  AmReceiverCallbackOwnerType owner;  
  AmReceiverCallbackIdType id;        
 
  AmReceiverCallbackInfo() = delete;

  AmReceiverCallbackInfo(const AmReceiverCallbackOwnerType owner, AmReceiverCallbackIdType id);
};

enum class AmSendMemoryTypePolicy {
  FallbackToHost = 0,  
  ErrorOnUnsupported,  
};

struct AmSendParams {
  uint32_t flags{UCP_AM_SEND_FLAG_REPLY};              
  ucp_datatype_t datatype{ucp_dt_make_contig(1)};      
  ucs_memory_type_t memoryType{UCS_MEMORY_TYPE_HOST};  
  AmSendMemoryTypePolicy memoryTypePolicy{
    AmSendMemoryTypePolicy::FallbackToHost};  
  std::optional<AmReceiverCallbackInfo> receiverCallbackInfo{
    std::nullopt};                      
  std::vector<std::byte> userHeader{};  

  void setUserHeader(const void* data, size_t size)
  {
    if (size > 0 && data == nullptr)
      throw std::invalid_argument(
        "AmSendParams::setUserHeader received null data with non-zero size");
    userHeader.resize(size);
    if (size > 0) memcpy(userHeader.data(), data, size);
  }

  void setUserHeader(std::string_view data) { setUserHeader(data.data(), data.size()); }
};

typedef const std::string SerializedRemoteKey;

struct AmReceiverCallbackOwnerHash {
  size_t operator()(const AmReceiverCallbackOwnerType& o) const noexcept
  {
    return std::hash<std::string_view>{}(
      std::string_view(o.data(), AmReceiverCallbackOwnerStorageSize));
  }
};
 
}  // namespace ucxx