stack_alloc.hpp

Go to the documentation of this file.

#ifndef STAN__MEMORY__STACK_ALLOC_HPP
#define STAN__MEMORY__STACK_ALLOC_HPP

#include <cstdlib>
#include <cstddef>
#include <sstream>
#include <stdexcept>
#if defined(_MSC_VER)
    #include <msinttypes.h>  // Microsoft Visual Studio lacks compliant stdint.h
#else
    #include <stdint.h> // FIXME: replace with cstddef?
#endif
#include <vector>
#include <stan/meta/likely.hpp>

namespace stan { 

  namespace memory { 

    template <typename T>
    bool is_aligned(T* ptr, unsigned int bytes_aligned) {
      return (reinterpret_cast<uintptr_t>(ptr) % bytes_aligned) == 0U;
    }


    namespace {
      const size_t DEFAULT_INITIAL_NBYTES = 1 << 16; // 64KB


      // FIXME: enforce alignment
      // big fun to inline, but only called twice
      inline char* eight_byte_aligned_malloc(size_t size) {
        char* ptr = static_cast<char*>(malloc(size));
        if (!ptr) return ptr; // malloc failed to alloc
        if (!is_aligned(ptr,8U)) {
          std::stringstream s;
          s << "invalid alignment to 8 bytes, ptr=" 
            << reinterpret_cast<uintptr_t>(ptr) 
            << std::endl;
          throw std::runtime_error(s.str());
        }
        return ptr;
      }
    }

    class stack_alloc {
    private: 
      std::vector<char*> 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
      char* cur_block_end_;       // ptr to cur_block_ptr_ + sizes_[cur_block_]
      char* next_loc_;            // ptr to next available spot in cur
                                  // block
      // next three for keeping track of nested allocations on top of stack:
      std::vector<size_t> nested_cur_blocks_;
      std::vector<char*> nested_next_locs_;
      std::vector<char*> nested_cur_block_ends_;
      

      char* move_to_next_block(size_t len) {
        char* 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(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:

      stack_alloc(size_t initial_nbytes = DEFAULT_INITIAL_NBYTES) :
        blocks_(1, 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 (!blocks_[0])
          throw std::bad_alloc();  // no msg allowed in bad_alloc ctor
      }

      ~stack_alloc() { 
        // free ALL blocks
        for (size_t i = 0; i < blocks_.size(); ++i)
          if (blocks_[i])
            free(blocks_[i]);
      }

      inline void* alloc(size_t len) {
        // Typically, just return and increment the next location.
        char* result = next_loc_;
        next_loc_ += len;
        // Occasionally, we have to switch blocks.
        if (unlikely(next_loc_ >= cur_block_end_))
          result = move_to_next_block(len);
        return (void*)result;
      }

      inline void recover_all() {
        cur_block_ = 0;
        next_loc_ = blocks_[0];
        cur_block_end_ = next_loc_ + sizes_[0];
      }

      inline void start_nested() {
        nested_cur_blocks_.push_back(cur_block_);
        nested_next_locs_.push_back(next_loc_);
        nested_cur_block_ends_.push_back(cur_block_end_);
      }

      inline void recover_nested() {
        if (unlikely(nested_cur_blocks_.empty()))
          recover_all();

        cur_block_ = nested_cur_blocks_.back();
        nested_cur_blocks_.pop_back();

        next_loc_ = nested_next_locs_.back();
        nested_next_locs_.pop_back();
        
        cur_block_end_ = nested_cur_block_ends_.back();
        nested_cur_block_ends_.pop_back();
      }
    
      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();
      }

      size_t bytes_allocated() {
        size_t sum = 0;
        for (size_t i = 0; i <= cur_block_; ++i) {
          sum += sizes_[i];
        }
        return sum;
      }

    };

  }
}
#endif