memory.hpp

Go to the documentation of this file.

#ifndef INCLUDE_RICCATI_MEMORY_ARENA_ALLOC_HPP
#define INCLUDE_RICCATI_MEMORY_ARENA_ALLOC_HPP
 
#include <riccati/macros.hpp>
#include <riccati/utils.hpp>
#include <Eigen/Dense>
#include <stdint.h>
#include <cstdlib>
#include <cstddef>
#include <sstream>
#include <stdexcept>
#include <vector>
 
namespace riccati {
 
namespace internal {
template <unsigned int Alignment, typename T>
RICCATI_ALWAYS_INLINE bool is_aligned(T* ptr) {
  return (reinterpret_cast<uintptr_t>(ptr) % Alignment) == 0U;
}
 
constexpr size_t DEFAULT_INITIAL_NBYTES = 1 << 16;  // 64KB
 
RICCATI_ALWAYS_INLINE unsigned char* eight_byte_aligned_malloc(size_t size) {
  unsigned char* ptr = static_cast<unsigned char*>(malloc(size));
  if (!ptr) {
    return ptr;  // malloc failed to alloc
  }
  if (!is_aligned<8U>(ptr)) {
    [](auto* ptr) RICCATI_COLD_PATH {
      std::stringstream s;
      s << "invalid alignment to 8 bytes, ptr="
        << reinterpret_cast<uintptr_t>(ptr) << std::endl;
      throw std::runtime_error(s.str());
    }(ptr);
  }
  return ptr;
}
 
}  // namespace internal
 
class arena_alloc {
 public:
  using byte_t = unsigned char;
  std::vector<byte_t*> blocks_;  // storage for blocks,
                                 // may be bigger than cur_block_
  std::vector<size_t> sizes_;    // could store initial & shift for others
  size_t cur_block_;             // index into blocks_ for next alloc
  byte_t* cur_block_end_;        // ptr to cur_block_ptr_ + sizes_[cur_block_]
  byte_t* next_loc_;             // ptr to next available spot in cur
                                 // block
 
  RICCATI_NO_INLINE byte_t* move_to_next_block(size_t len) {
    byte_t* result;
    ++cur_block_;
    // Find the next block (if any) containing at least len bytes.
    while ((cur_block_ < blocks_.size()) && (sizes_[cur_block_] < len)) {
      ++cur_block_;
    }
    // Allocate a new block if necessary.
    if (unlikely(cur_block_ >= blocks_.size())) {
      // New block should be max(2*size of last block, len) bytes.
      size_t newsize = sizes_.back() * 2;
      if (newsize < len) {
        newsize = len;
      }
      blocks_.push_back(internal::eight_byte_aligned_malloc(newsize));
      if (!blocks_.back()) {
        throw std::bad_alloc();
      }
      sizes_.push_back(newsize);
    }
    result = blocks_[cur_block_];
    // Get the object's state back in order.
    next_loc_ = result + len;
    cur_block_end_ = result + sizes_[cur_block_];
    return result;
  }
 
 public:
  RICCATI_NO_INLINE explicit arena_alloc(size_t initial_nbytes)
      : blocks_(1, internal::eight_byte_aligned_malloc(initial_nbytes)),
        sizes_(1, initial_nbytes),
        cur_block_(0),
        cur_block_end_(blocks_[0] + initial_nbytes),
        next_loc_(blocks_[0]) {
    if (unlikely(!blocks_[0])) {
      throw std::bad_alloc();  // no msg allowed in bad_alloc ctor
    }
  }
  RICCATI_NO_INLINE arena_alloc()
      : blocks_(1, internal::eight_byte_aligned_malloc(
                       internal::DEFAULT_INITIAL_NBYTES)),
        sizes_(1, internal::DEFAULT_INITIAL_NBYTES),
        cur_block_(0),
        cur_block_end_(blocks_[0] + internal::DEFAULT_INITIAL_NBYTES),
        next_loc_(blocks_[0]) {
    if (unlikely(!blocks_[0])) {
      throw std::bad_alloc();  // no msg allowed in bad_alloc ctor
    }
  }
  arena_alloc(const arena_alloc&) = delete;
  arena_alloc& operator=(const arena_alloc&) = delete;
  arena_alloc(arena_alloc&&) = delete;
  arena_alloc& operator=(arena_alloc&&) = delete;
 
  RICCATI_NO_INLINE ~arena_alloc() {
    // free ALL blocks
    for (auto& block : blocks_) {
      if (block) {
        free(block);
      }
    }
  }
 
  RICCATI_ALWAYS_INLINE void* alloc(size_t len) {
    size_t pad = len % 8 == 0 ? 0 : 8 - len % 8;
 
    // Typically, just return and increment the next location.
    byte_t* result = next_loc_;
    next_loc_ += len + pad;
    // Occasionally, we have to switch blocks.
    if (unlikely(next_loc_ >= cur_block_end_)) {
      result = move_to_next_block(len);
    }
    return reinterpret_cast<void*>(result);
  }
 
  template <typename T>
  RICCATI_ALWAYS_INLINE T* alloc_array(size_t n) {
    return static_cast<T*>(alloc(n * sizeof(T)));
  }
 
  RICCATI_ALWAYS_INLINE void recover_all() {
    if (unlikely(blocks_.size() > 1)) {
      std::size_t sum = 0;
      for (size_t i = 1; i < blocks_.size(); ++i) {
        sum += sizes_[i];
        free(blocks_[i]);
      }
      blocks_.clear();
      sizes_.clear();
      blocks_.push_back(internal::eight_byte_aligned_malloc(sum));
      sizes_.push_back(sum);
    }
    cur_block_ = 0;
    next_loc_ = blocks_[0];
    cur_block_end_ = next_loc_ + sizes_[0];
  }
 
  inline void free_all() {
    // frees all BUT the first (index 0) block
    for (size_t i = 1; i < blocks_.size(); ++i) {
      if (blocks_[i]) {
        free(blocks_[i]);
      }
    }
    sizes_.resize(1);
    blocks_.resize(1);
    recover_all();
  }
 
  inline size_t bytes_allocated() const {
    size_t sum = 0;
    for (size_t i = 0; i <= cur_block_; ++i) {
      sum += sizes_[i];
    }
    return sum;
  }
 
  inline bool in_stack(const void* ptr) const {
    for (size_t i = 0; i < cur_block_; ++i) {
      if (ptr >= blocks_[i] && ptr < blocks_[i] + sizes_[i]) {
        return true;
      }
    }
    if (ptr >= blocks_[cur_block_] && ptr < next_loc_) {
      return true;
    }
    return false;
  }
};
 
template <typename T, typename ArenaType>
struct arena_allocator {
  ArenaType* alloc_;
  bool owns_alloc_{false};
  using value_type = T;
  RICCATI_NO_INLINE explicit arena_allocator(ArenaType* alloc,
                                             bool owns_alloc = false)
      : alloc_(alloc), owns_alloc_(owns_alloc) {}
  RICCATI_NO_INLINE arena_allocator()
      : alloc_(new ArenaType{}), owns_alloc_(true) {}
 
  RICCATI_NO_INLINE arena_allocator(const arena_allocator& rhs)
      : alloc_(rhs.alloc_), owns_alloc_(false) {};
  template <typename U, typename UArena>
  RICCATI_NO_INLINE arena_allocator(const arena_allocator<U, UArena>& rhs)
      : alloc_(rhs.alloc_), owns_alloc_(false) {}
 
  ~arena_allocator() {
    if (owns_alloc_) {
      delete alloc_;
    }
  }
 
  template <typename T_ = T>
  RICCATI_ALWAYS_INLINE T_* allocate(std::size_t n) {
    return alloc_->template alloc_array<T_>(n);
  }
 
  RICCATI_ALWAYS_INLINE void recover_memory() noexcept {
    alloc_->recover_all();
  }
 
  void deallocate(T* /*p*/, std::size_t /*n*/) noexcept {}
 
  constexpr bool operator==(const arena_allocator&) const noexcept {
    return true;
  }
  constexpr bool operator!=(const arena_allocator&) const noexcept {
    return false;
  }
};
 
struct dummy_allocator {
  std::vector<void*> ptrs_;
  template <typename T>
  inline T* allocate(std::size_t n) {
    void* ptr = std::malloc(n * sizeof(T));
    ptrs_.push_back(ptr);
    return reinterpret_cast<T*>(ptr);
  }
  constexpr inline void recover_memory() noexcept {}
  template <typename T>
  constexpr inline void deallocate(T* p, std::size_t n) noexcept {}
  constexpr bool operator==(const dummy_allocator&) const noexcept {
    return true;
  }
  constexpr bool operator!=(const dummy_allocator&) const noexcept {
    return false;
  }
  ~dummy_allocator() {
    for (auto ptr : ptrs_) {
      std::free(ptr);
    }
  }
};
 
template <typename Expr>
RICCATI_ALWAYS_INLINE auto eval(dummy_allocator& alloc, Expr&& expr) {
  return eval(std::forward<Expr>(expr));
}
 
}  // namespace riccati
#endif