handle.hpp

Go to the documentation of this file.

// Copyright (c) 2018--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 "./allocators.hpp"
#include "./memcpy.hpp"
#include "../concepts.hpp"
#include "../macros.hpp"
 
#include <array>
#include <memory>
#include <type_traits>
#include <utility>
 
namespace nda::mem {


  static constexpr bool init_dcmplx = true;

  template <typename T, int Al>
  struct alignas(Al) aligner {
    T x;

    [[nodiscard]] T &get() noexcept { return x; }

    [[nodiscard]] T const &get() const noexcept { return x; }
  };

  struct do_not_initialize_t {};

  inline static constexpr do_not_initialize_t do_not_initialize{};

  struct init_zero_t {};

  inline static constexpr init_zero_t init_zero{};



  // Forward declaration.
  template <typename T, AddressSpace AdrSp = Host>
  struct handle_shared;

  template <typename T, Allocator A = mallocator<>>
  struct handle_heap {
    static_assert(std::is_nothrow_destructible_v<T>, "nda::mem::handle_heap requires the value_type to have a non-throwing destructor");
 
    private:
    // Pointer to the start of the actual data.
    T *_data = nullptr;
 
    // Size of the data (number of T elements). Invariant: size > 0 iif data != nullptr.
    size_t _size = 0;
 
    // Allocator to use.
#ifndef NDA_DEBUG_LEAK_CHECK
    static inline A allocator; // NOLINT (allocator is not specific to a single instance)
#else
    static inline leak_check<A> allocator; // NOLINT (allocator is not specific to a single instance)
#endif
 
    // For shared ownership (points to a blk_T_t).
    mutable std::shared_ptr<void> sptr;
 
    // Type of the memory block, i.e. a pointer to the data and its size.
    using blk_T_t = std::pair<T *, size_t>;
 
    // Release the handled memory (data pointer and size are not set to null here).
    static void destruct(blk_T_t b) noexcept {
      auto [data, size] = b;
 
      // do nothing if the data is null
      if (data == nullptr) return;
 
      // if needed, call the destructors of the objects stored
      if constexpr (A::address_space == Host and !(std::is_trivial_v<T> or nda::is_complex_v<T>)) {
        for (size_t i = 0; i < size; ++i) data[i].~T();
      }
 
      // deallocate the memory block
      allocator.deallocate({(char *)data, size * sizeof(T)});
    }
 
    // Deleter for the shared pointer.
    // NB: with a custom deleter, shared_ptr does NOT delete the managed object itself,
    // so we must free the blk_T_t node allocated by `new blk_T_t{...}` in get_sptr().
    static void deleter(void *p) noexcept {
      destruct(*((blk_T_t *)p));
      delete (blk_T_t *)p;
    }
 
    public:
    using value_type = T;

    using allocator_type = A;

    static constexpr auto address_space = allocator_type::address_space;

    std::shared_ptr<void> get_sptr() const {
      if (not sptr) sptr.reset(new blk_T_t{_data, _size}, deleter);
      return sptr;
    }

    ~handle_heap() noexcept {
      if (not sptr and not(is_null())) destruct({_data, _size});
    }

    handle_heap() = default;

    handle_heap(handle_heap &&h) noexcept : _data(h._data), _size(h._size), sptr(std::move(h.sptr)) {
      h._data = nullptr;
      h._size = 0;
    }

    handle_heap &operator=(handle_heap &&h) noexcept {
      // release current resources if they are not shared and not null
      if (not sptr and not(is_null())) destruct({_data, _size});
 
      // move the resources from the source handle
      _data = h._data;
      _size = h._size;
      sptr  = std::move(h.sptr);
 
      // reset the source handle to a null state
      h._data = nullptr;
      h._size = 0;
      return *this;
    }

    explicit handle_heap(handle_heap const &h) : handle_heap(h.size(), do_not_initialize) {
      if (is_null()) return;
      if constexpr (std::is_trivially_copyable_v<T>) {
        memcpy<address_space, address_space>(_data, h.data(), h.size() * sizeof(T));
      } else {
        for (size_t i = 0; i < _size; ++i) new (_data + i) T(h[i]);
      }
    }

    handle_heap &operator=(handle_heap const &h) {
      *this = handle_heap{h};
      return *this;
    }

    template <OwningHandle<value_type> H>
    explicit handle_heap(H const &h) : handle_heap(h.size(), do_not_initialize) {
      if (is_null()) return;
      if constexpr (std::is_trivially_copyable_v<T>) {
        memcpy<address_space, H::address_space>((void *)_data, (void *)h.data(), _size * sizeof(T));
      } else {
        static_assert(address_space == H::address_space,
                      "Constructing an nda::mem::handle_heap from a handle of a different address space requires a trivially copyable value_type");
        for (size_t i = 0; i < _size; ++i) new (_data + i) T(h[i]);
      }
    }

    template <Allocator AS>
    handle_heap &operator=(handle_heap<T, AS> const &h) {
      *this = handle_heap{h};
      return *this;
    }

    handle_heap(long size, do_not_initialize_t) {
      if (size == 0) return;
      auto b = allocator.allocate(size * sizeof(T));
      if (not b.ptr) throw std::bad_alloc{};
      _data = (T *)b.ptr;
      _size = size;
    }

    handle_heap(long size, init_zero_t) {
      if (size == 0) return;
      auto b = allocator.allocate_zero(size * sizeof(T));
      if (not b.ptr) throw std::bad_alloc{};
      _data = (T *)b.ptr;
      _size = size;
    }

    handle_heap(long size) {
      if (size == 0) return;
      blk_t b;
      if constexpr (is_complex_v<T> && init_dcmplx)
        b = allocator.allocate_zero(size * sizeof(T));
      else
        b = allocator.allocate(size * sizeof(T));
      if (not b.ptr) throw std::bad_alloc{};
      _data = (T *)b.ptr;
      _size = size;
 
      // call placement new for non trivial and non complex types
      if constexpr (!std::is_trivial_v<T> and !is_complex_v<T>) {
        for (size_t i = 0; i < size; ++i) new (_data + i) T();
      }
    }

    [[nodiscard]] T &operator[](long i) noexcept { return _data[i]; }

    [[nodiscard]] T const &operator[](long i) const noexcept { return _data[i]; }

    [[nodiscard]] bool is_null() const noexcept {
#ifdef NDA_DEBUG
      // check the invariants in debug mode
      EXPECTS((_data == nullptr) == (_size == 0));
#endif
      return _data == nullptr;
    }

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

    [[nodiscard]] long size() const noexcept { return _size; }
  };

  template <typename T, size_t Size>
  struct handle_stack {
    static_assert(std::is_copy_constructible_v<T>, "nda::mem::handle_stack requires the value_type to be copy constructible");
    static_assert(std::is_nothrow_destructible_v<T>, "nda::mem::handle_stack requires the value_type to have a non-throwing destructor");
 
    private:
    // Memory buffer on the stack to store the data.
    std::array<char, sizeof(T) * Size> buffer;
 
    public:
    using value_type = T;

    static constexpr auto address_space = Host;

    ~handle_stack() noexcept {
      if constexpr (!std::is_trivial_v<T>) {
        // explicitly call the destructor for each non-trivial object
        for (size_t i = 0; i < Size; ++i) data()[i].~T();
      }
    }

    handle_stack() = default;

    handle_stack(handle_stack &&h) noexcept { operator=(h); }

    handle_stack &operator=(handle_stack &&h) noexcept {
      operator=(h);
      return *this;
    }

    handle_stack(handle_stack const &h) noexcept { operator=(h); }

    handle_stack &operator=(handle_stack const &h) {
      for (size_t i = 0; i < Size; ++i) new (data() + i) T(h[i]);
      return *this;
    }

    handle_stack(long /*size*/, do_not_initialize_t) {}

    handle_stack(long /*size*/, init_zero_t) {
      static_assert(std::is_scalar_v<T> or is_complex_v<T>, "nda::mem::handle_stack can only be initialized to zero for scalar and complex types");
      for (size_t i = 0; i < Size; ++i) data()[i] = 0;
    }

    handle_stack(long /*size*/) : handle_stack{} {
      if constexpr (!std::is_trivial_v<T> and !is_complex_v<T>) {
        for (size_t i = 0; i < Size; ++i) new (data() + i) T();
      }
    }

    [[nodiscard]] T &operator[](long i) noexcept { return data()[i]; }

    [[nodiscard]] T const &operator[](long i) const noexcept { return data()[i]; }

    static constexpr bool is_null() noexcept { return false; }

    [[nodiscard]] T *data() const noexcept { return (T *)buffer.data(); }

    static constexpr long size() noexcept { return Size; }
  };

  template <typename T, size_t Size>
  struct handle_sso { // struct alignas(alignof(T)) handle_sso {
    static_assert(Size > 0, "Size == 0 makes no sense in nda::mem::handle_sso");
    static_assert(std::is_copy_constructible_v<T>, "nda::mem::handle_sso requires the value_type to be copy constructible");
    static_assert(std::is_nothrow_destructible_v<T>, "nda::mem::handle_sso requires the value_type to have a non-throwing destructor");
 
    private:
    // Memory buffer on the stack to store the data.
    std::array<char, sizeof(T) * Size> buffer;
 
    // Pointer to the start of the actual data.
    T *_data = nullptr;
 
    // Size of the data (number of T elements). Invariant: size > 0 iif data != nullptr.
    size_t _size = 0;
 
    // Release the handled memory.
    void clean() noexcept {
      if (is_null()) return;
      if constexpr (!std::is_trivial_v<T>) {
        // explicitly call the destructor for each non-trivial object
        for (size_t i = 0; i < _size; ++i) data()[i].~T();
      }
      if (on_heap()) mallocator<>::deallocate({(char *)_data, _size * sizeof(T)});
      _data = nullptr;
      _size = 0;
    }
 
    public:
    using value_type = T;

    static constexpr auto address_space = Host;

    handle_sso() {}; // NOLINT (user-defined constructor to avoid value initialization of the buffer)

    ~handle_sso() noexcept { clean(); }

    handle_sso(handle_sso &&h) noexcept : _size(h._size) {
      if (on_heap()) {
        _data = h._data;
      } else {
        if (_size != 0) {
          _data = (T *)buffer.data();
          for (size_t i = 0; i < _size; ++i) new (data() + i) T(h[i]);
        }
      }
      h._data = nullptr;
      h._size = 0;
    }

    handle_sso &operator=(handle_sso &&h) noexcept {
      clean();
      _size = h._size;
      if (on_heap()) {
        _data = h._data;
      } else {
        if (_size != 0) {
          _data = (T *)buffer.data();
          for (size_t i = 0; i < _size; ++i) new (_data + i) T(h[i]);
        }
      }
      h._data = nullptr;
      h._size = 0;
      return *this;
    }

    handle_sso(handle_sso const &h) : handle_sso(h.size(), do_not_initialize) {
      if (is_null()) return;
      if constexpr (std::is_trivially_copyable_v<T>) {
        memcpy<address_space, address_space>((void *)_data, (void *)h.data(), _size * sizeof(T));
      } else {
        for (size_t i = 0; i < _size; ++i) new (_data + i) T(h[i]);
      }
    }

    handle_sso &operator=(handle_sso const &h) {
      if (this == &h) return *this; // self assignment
      clean();
      _size = h._size;
      if (_size == 0) return *this;
      if (on_heap()) {
        auto b = mallocator<>::allocate(_size * sizeof(T));
        if (not b.ptr) throw std::bad_alloc{};
        _data = (T *)b.ptr;
      } else {
        _data = (T *)buffer.data();
      }
      for (size_t i = 0; i < _size; ++i) new (_data + i) T(h[i]);
      return *this;
    }

    template <OwningHandle<value_type> H>
    explicit handle_sso(H const &h) : handle_sso(h.size(), do_not_initialize) {
      if (is_null()) return;
      if constexpr (std::is_trivially_copyable_v<T>) {
        memcpy<address_space, H::address_space>((void *)_data, (void *)h.data(), _size * sizeof(T));
      } else {
        static_assert(address_space == H::address_space,
                      "Constructing an nda::mem::handle_sso from a handle of a different address space requires a trivially copyable value_type");
        for (size_t i = 0; i < _size; ++i) new (_data + i) T(h[i]);
      }
    }

    handle_sso(long size, do_not_initialize_t) {
      if (size == 0) return;
      _size = size;
      if (not on_heap()) {
        _data = (T *)buffer.data();
      } else {
        auto b = mallocator<>::allocate(size * sizeof(T));
        if (not b.ptr) throw std::bad_alloc{};
        _data = (T *)b.ptr;
      }
    }

    handle_sso(long size, init_zero_t) {
      static_assert(std::is_scalar_v<T> or is_complex_v<T>, "nda::mem::handle_sso can only be initialized to zero for scalar and complex types");
      if (size == 0) return;
      _size = size;
      if (not on_heap()) {
        _data = (T *)buffer.data();
        for (size_t i = 0; i < _size; ++i) data()[i] = 0;
      } else {
        auto b = mallocator<>::allocate_zero(size * sizeof(T)); //, alignof(T));
        if (not b.ptr) throw std::bad_alloc{};
        _data = (T *)b.ptr;
      }
    }

    handle_sso(long size) {
      if (size == 0) return;
      _size = size;
      if (not on_heap()) {
        _data = (T *)buffer.data();
      } else {
        blk_t b;
        if constexpr (is_complex_v<T> && init_dcmplx)
          b = mallocator<>::allocate_zero(size * sizeof(T));
        else
          b = mallocator<>::allocate(size * sizeof(T));
        if (not b.ptr) throw std::bad_alloc{};
        _data = (T *)b.ptr;
      }
 
      // call placement new for non trivial and non complex types
      if constexpr (!std::is_trivial_v<T> and !is_complex_v<T>) {
        for (size_t i = 0; i < size; ++i) new (_data + i) T();
      }
    }

    [[nodiscard]] T &operator[](long i) noexcept { return _data[i]; }

    [[nodiscard]] T const &operator[](long i) const noexcept { return _data[i]; }

    [[nodiscard]] bool on_heap() const { return _size > Size; }

    [[nodiscard]] bool is_null() const noexcept {
#ifdef NDA_DEBUG
      EXPECTS((_data == nullptr) == (_size == 0));
#endif
      return _data == nullptr;
    }

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

    [[nodiscard]] long size() const noexcept { return _size; }
  };

  template <typename T, AddressSpace AdrSp>
  struct handle_shared {
    static_assert(std::is_nothrow_destructible_v<T>, "nda::mem::handle_shared requires the value_type to have a non-throwing destructor");
 
    private:
    // Pointer to the start of the actual data.
    T *_data = nullptr;
 
    // Size of the data (number of T elements). Invariant: size > 0 iif data != 0.
    size_t _size = 0;
 
    // Type of the memory block, i.e. a pointer to the data and its size.
    using blk_t = std::pair<T *, size_t>;
 
    // For shared ownership (points to a blk_T_t).
    std::shared_ptr<void> sptr;
 
    public:
    using value_type = T;

    static constexpr auto address_space = AdrSp;

    handle_shared() = default;

    handle_shared(T *data, size_t size, void *foreign_handle, void (*foreign_decref)(void *)) noexcept
       : _data(data), _size(size), sptr{foreign_handle, foreign_decref} {}

    template <Allocator A>
    handle_shared(handle_heap<T, A> const &h) noexcept
      requires(A::address_space == address_space)
       : _data(h.data()), _size(h.size()) {
      if (not h.is_null()) sptr = h.get_sptr();
    }

    [[nodiscard]] T &operator[](long i) noexcept { return _data[i]; }

    [[nodiscard]] T const &operator[](long i) const noexcept { return _data[i]; }

    [[nodiscard]] bool is_null() const noexcept {
#ifdef NDA_DEBUG
      // Check the Invariants in Debug Mode
      EXPECTS((_data == nullptr) == (_size == 0));
#endif
      return _data == nullptr;
    }

    [[nodiscard]] long refcount() const noexcept { return sptr.use_count(); }

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

    [[nodiscard]] long size() const noexcept { return _size; }
  };

  template <typename T, AddressSpace AdrSp = Host>
  struct handle_borrowed {
    private:
    // Value type of the data with const removed.
    using T0 = std::remove_const_t<T>;
 
    // Parent handle (required for regular -> shared promotion in Python Converter).
    handle_heap<T0> const *_parent = nullptr;
 
    // Pointer to the start of the actual data.
    T *_data = nullptr;
 
    public:
    using value_type = T;

    static constexpr auto address_space = AdrSp;

    handle_borrowed() = default;

    handle_borrowed &operator=(handle_borrowed &&) = default;

    handle_borrowed(handle_borrowed const &) = default;

    handle_borrowed &operator=(handle_borrowed const &) = default;

    handle_borrowed(T *ptr) noexcept : _data(ptr) {}

    template <Handle H>
      requires(address_space == H::address_space and (std::is_const_v<value_type> or !std::is_const_v<typename H::value_type>)
               and std::is_same_v<const value_type, const typename H::value_type>)
    handle_borrowed(H const &h, long offset = 0) noexcept : _data(h.data() + offset) {
      if constexpr (std::is_same_v<H, handle_heap<T0>>) _parent = &h;
    }

    [[nodiscard]] T &operator[](long i) noexcept { return _data[i]; }

    [[nodiscard]] T const &operator[](long i) const noexcept { return _data[i]; }

    [[nodiscard]] bool is_null() const noexcept { return _data == nullptr; }

    [[nodiscard]] handle_heap<T0> const *parent() const { return _parent; }

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

 
} // namespace nda::mem