pool_allocator.hh

// This file is a part of Aurelia.
// Copyright (C) 2010  Valentin David
// Copyright (C) 2010  University of Bergen
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.
//
// You should have received a copy of the GNU General Public License
// along with this program.  If not, see <http://www.gnu.org/licenses/>.

#ifndef __POOL_ALLOCATOR_HH
# define __POOL_ALLOCATOR_HH

# include <aurelia-config.hh>

# include <atomic>
# include "../lockfree/simple_stack.hh"

# ifdef HAVE_VALGRIND_VALGRIND_H
#  pragma GCC diagnostic push
#  pragma GCC diagnostic ignored "-Wunused-but-set-variable"
#  pragma GCC diagnostic ignored "-Wunused-value"
#  include <valgrind/valgrind.h>
#  include <valgrind/memcheck.h>
# else
#  define VALGRIND_CREATE_MEMPOOL(X, Y, Z)
#  define VALGRIND_MAKE_MEM_NOACCESS(X, Y)
#  define VALGRIND_MEMPOOL_ALLOC(X, Y, Z)
#  define VALGRIND_MEMPOOL_FREE(X, Y)
#  define VALGRIND_MAKE_MEM_DEFINED(X, Y)
#  define VALGRIND_MAKE_MEM_UNDEFINED(X, Y)
# endif

# include <cstdlib>
# include <memory>

namespace aurelia {

  template <int SIZE>
  struct block_list {
  private:
    lockfree::simple_stack<void*> blocks;

  public:
    block_list() {}
    block_list(const block_list&) = delete;

    void* new_block() {
      void *p = malloc(SIZE);
      if (p == NULL)
        throw std::bad_alloc();
      blocks.push(p);
      return p;
    }

    ~block_list() {
      while (true) {
        lockfree::simple_stack<void*>::wrapped w = blocks.pop();
        if (w.end())
          return ;
        free(*w);
      }
    }
  };

  template <typename T,
            int POOL_SIZE = 4096>
  struct chunk_alloc {
  private:
    static const void* mempool() {
      static bool init = false;
      if (!init) {
        init = true;
        VALGRIND_CREATE_MEMPOOL(&typeid(chunk_alloc), 0, false);
      }
      return &typeid(chunk_alloc);
    }

    union block {
      typename std::aligned_storage<sizeof(T),
                                    std::alignment_of<T>::value>::type obj;
      block *next;
    };

    static block*& free() {
      thread_local block* _free = NULL;
      return _free;
    }

    static block_list<sizeof(block)*POOL_SIZE>& blocks() {
      static block_list<sizeof(block)*POOL_SIZE> _blocks;
      return _blocks;
    }

    static block* ensure_blocks() {
      block* f = free();
      if (f == NULL) {
        f = (block*)blocks().new_block();
        free() = f;
        for (int i = 0; i < (POOL_SIZE - 1); ++i) {
          f[i].next = f + i + 1;
        }
        f[POOL_SIZE-1].next = NULL;
        VALGRIND_MAKE_MEM_NOACCESS(f, sizeof(block)*POOL_SIZE);
      }
      return f;
    }

  public:
    static void *alloc() {
      block* n = ensure_blocks();
      VALGRIND_MEMPOOL_ALLOC(mempool(), n, sizeof(block));
      VALGRIND_MAKE_MEM_DEFINED(&(n->next), sizeof(block*));
      free() = n->next;
      VALGRIND_MAKE_MEM_UNDEFINED(&(n->next), sizeof(block*));
      return (void*)&(n->obj);
    }

    static size_t offset() {
      block n;
      return ((char*)&n.obj - (char*)&n);
    }

    static void del(void *p) {
      block *n = (block*)(((char*)p) - offset());
      n->next = free();
      VALGRIND_MEMPOOL_FREE(mempool(), n);
      free() = n;
    }
  };

  template <typename T,
            int POOL_SIZE = 4096,
            typename RealAlloc = chunk_alloc<T, POOL_SIZE> >
  struct pool_allocator {
    typedef size_t size_type;
    typedef T* pointer;
    typedef T& reference;
    typedef const T* const_pointer;
    typedef const T& const_reference;
    typedef T value_type;

    template <typename U>
    struct rebind {
      typedef pool_allocator<U, POOL_SIZE> other;
    };

    template <typename U, int OTHERSIZE>
    pool_allocator(const pool_allocator<U, OTHERSIZE>&) {}
    pool_allocator() {}

    pointer allocate(size_type n) {
      if (n == 1)
        return (pointer)RealAlloc::alloc();
      std::allocator<T> a;
      return a.allocate(n);
    }

    void deallocate(pointer p, size_type n) {
      if (n == 1)
        return RealAlloc::del(p);
      std::allocator<T> a;
      a.deallocate(p, n);
    }

    void destroy(pointer p) {
      p->~T();
    }

    template <typename... Args>
    void construct(pointer p, Args... args...) {
      new (p) T(std::forward<Args>(args)...);
    }
  };

}

# ifdef HAVE_VALGRIND_VALGRIND_H
#  pragma GCC diagnostic pop
# endif

#endif