Line data    Source code

       1             : #pragma once
       2             : 
       3             : #include <panoc-alm/util/problem.hpp>
       4             : 
       5             : #include <algorithm>
       6             : #include <cassert>
       7             : #include <stack>
       8             : #include <utility>
       9             : #include <vector>
      10             : 
      11             : namespace pa {
      12             : 
      13           4 : class vec_allocator {
      14             :   private:
      15             :     size_t num_vec;
      16             :     Eigen::Index n;
      17             :     std::vector<real_t> storage;
      18           8 :     struct findstack : std::stack<real_t *, std::vector<real_t *>> {
      19             :         using stack::stack;
      20             : 
      21      240036 :         auto begin() { return this->c.begin(); }
      22             :         auto begin() const { return this->c.begin(); }
      23      480072 :         auto end() { return this->c.end(); }
      24             :         auto end() const { return this->c.end(); }
      25             : 
      26             :         auto cbegin() const { return this->c.cbegin(); }
      27             :         auto cend() const { return this->c.cend(); }
      28             : 
      29             :         auto rbegin() { return this->c.rbegin(); }
      30             :         auto rbegin() const { return this->c.rbegin(); }
      31             :         auto rend() { return this->c.rend(); }
      32             :         auto rend() const { return this->c.rend(); }
      33             : 
      34             :         auto crbegin() const { return this->c.crbegin(); }
      35             :         auto crend() const { return this->c.crend(); }
      36             :     } stack;
      37           4 :     size_t highwatermark = 0;
      38             : 
      39             :   public:
      40           8 :     vec_allocator(size_t num_vec, Eigen::Index n)
      41           4 :         : num_vec(num_vec), n(n), storage(num_vec * n, NaN) {
      42          63 :         for (auto it = storage.begin(); it < storage.end(); it += n)
      43          59 :             stack.push(&*it);
      44           4 :     }
      45             : 
      46             :     vec_allocator(const vec_allocator &) = delete;
      47             :     vec_allocator(vec_allocator &&)      = delete;
      48             :     vec_allocator &operator=(const vec_allocator &) = delete;
      49             :     vec_allocator &operator=(vec_allocator &&) = delete;
      50             : 
      51             :     struct alloc_raii_wrapper {
      52             :         using mvec = Eigen::Map<vec>;
      53             :         mvec v;
      54             :         vec_allocator *alloc;
      55             : 
      56             :         alloc_raii_wrapper(real_t *dptr, Eigen::Index n, vec_allocator *alloc)
      57             :             : v{dptr, n}, alloc{alloc} {}
      58           1 :         alloc_raii_wrapper(mvec &&v, vec_allocator *alloc)
      59           1 :             : v{std::move(v)}, alloc{alloc} {}
      60           1 :         ~alloc_raii_wrapper() {
      61           1 :             assert(alloc);
      62           1 :             alloc->free(v);
      63           1 :         }
      64             :         alloc_raii_wrapper(const alloc_raii_wrapper &) = delete;
      65             :         alloc_raii_wrapper(alloc_raii_wrapper &&o)
      66             :             : v{std::exchange(o.v, {nullptr, 0})}, //
      67             :               alloc{std::exchange(o.alloc, nullptr)} {}
      68             :         alloc_raii_wrapper &operator=(const alloc_raii_wrapper &) = delete;
      69             :         alloc_raii_wrapper &operator=(alloc_raii_wrapper &&o) {
      70             :             this->v     = std::exchange(o.v, {nullptr, 0});
      71             :             this->alloc = std::exchange(o.alloc, nullptr);
      72             :             return *this;
      73             :         }
      74             : 
      75             :         alloc_raii_wrapper &operator=(crvec v) {
      76             :             this->v = v;
      77             :             return *this;
      78             :         }
      79             :         operator crvec() const { return v; }
      80             :         operator rvec() { return v; }
      81             :     };
      82             : 
      83      120019 :     auto alloc() {
      84      120019 :         if (stack.empty())
      85           1 :             throw std::bad_alloc();
      86      120018 :         auto dptr = stack.top();
      87      120018 :         stack.pop();
      88      120018 :         highwatermark = std::max(used_space(), highwatermark);
      89      120018 :         return Eigen::Map<vec>(dptr, n);
      90      120018 :     }
      91             : 
      92           2 :     alloc_raii_wrapper alloc_raii() { return {alloc(), this}; }
      93             : 
      94      120018 :     void free(rvec v) {
      95      120018 :         auto dptr = v.data();
      96      120018 :         assert(dptr >= &*storage.begin());
      97      120018 :         assert(dptr <= &*storage.end() - n);
      98      240036 :         assert(std::find(stack.begin(), stack.end(), dptr) == stack.end() &&
      99             :                "double free");
     100      120018 :         stack.push(dptr);
     101      120018 :     }
     102             : 
     103             :     template <class... Vecs>
     104      100010 :     void free(rvec first, Vecs &&...vecs) {
     105      100010 :         free(first);
     106      100010 :         free(vecs...);
     107      100010 :     }
     108             : 
     109      260046 :     size_t size() const { return stack.size(); }
     110      240038 :     size_t used_space() const { return num_vec - size(); }
     111             : 
     112       20002 :     Eigen::Index vector_size() const { return n; }
     113             : 
     114           8 :     size_t highwater() const { return highwatermark; }
     115             : };
     116             : 
     117             : } // namespace pa