Аллокатор для объектов маленького размера.

Написан под влиянием книги Александреску.

#include <vector>
#include <algorithm>
#include <iostream>

using std::vector;

const unsigned char kBlocksInChunk = 128;
const size_t kMaxSizeOfSmallObject = 32;
const size_t kBlockSize = 64;

struct Chunk {
  unsigned char* data_;
  unsigned char first_available_block_;
  unsigned char available_blocks_count_;

  void Init(size_t block_size, unsigned char blocks_count) {
    data_ = static_cast<unsigned char*>(::operator new(block_size *
                                                       blocks_count));
    first_available_block_ = 0;
    available_blocks_count_ = blocks_count;
    unsigned char* tmp_data = data_;
    for (unsigned char i = 0; i != blocks_count; ++i) {
      *tmp_data = i;
      tmp_data += block_size;
    }
  }

  void* Allocate(size_t block_size) {
    if (available_blocks_count_ == 0) {
      return NULL;
    }
    unsigned char * result = data_ + first_available_block_ * block_size;
    first_available_block_ = *result;
    --available_blocks_count_;
    return result;
  }

  void Deallocate(void* p, size_t block_size) {
    // Pointer not in chunk.
    if (p < data_) {
      return;
    }
    unsigned char * to_release = static_cast<unsigned char*>(p);
    // Alignment check.
    if ((to_release - data_) % block_size != 0) {
      return;
    }
    *to_release = first_available_block_;
    first_available_block_ = static_cast<unsigned char>((to_release - data_) /
                                                        block_size);
    ++available_blocks_count_;
  }

  void Release() {
    ::operator delete(data_);
  }

  bool operator ==(const Chunk& other) const {
    return (data_ == other.data_) &&
        (first_available_block_ == other.first_available_block_) &&
        (available_blocks_count_ == other.available_blocks_count_);
  }
};

class FixedAllocator {
private:
  friend class SmallObjAllocator;

  size_t block_size_;
  unsigned char blocks_num_;
  vector<Chunk> chunks_;
  Chunk* alloc_chunk_;
  Chunk* dealloc_chunk_;

public:
  FixedAllocator(size_t block_size, unsigned char block_num):
    block_size_(block_size), blocks_num_(block_num), alloc_chunk_(NULL),
    dealloc_chunk_(NULL) {}

  void* Allocate() {
    if (alloc_chunk_ == NULL ||
        alloc_chunk_->available_blocks_count_ == 0) {
      // Seek for another chunk.
      bool is_found = false;
      for (vector<Chunk>::iterator it = chunks_.begin();
           it != chunks_.end(); ++it) {
        if (it->available_blocks_count_ > 0) {
          // New chunk is found.
          alloc_chunk_ = &(*it);
          is_found = true;
          break;
        }
      }
      if (!is_found) {
        // Add new chunk.
        Chunk new_chunk;
        new_chunk.Init(block_size_, blocks_num_);
        chunks_.push_back(new_chunk);
        alloc_chunk_ = &(chunks_.back());
        dealloc_chunk_ = &(chunks_.back());
      }
    }
    return alloc_chunk_->Allocate(block_size_);
  }

  void Deallocate(void* p) {
    Chunk* found_chunk = NULL;
    if (dealloc_chunk_ != NULL && p >= dealloc_chunk_->data_ &&
        p < dealloc_chunk_->data_ + block_size_ * blocks_num_) {
      found_chunk = dealloc_chunk_;
    } else {
      // Find chunk.
      for (vector<Chunk>::iterator it = chunks_.begin();
           it != chunks_.end(); ++it) {
        if (p >= it->data_ && p < it->data_ + block_size_ * blocks_num_) {
          // Chunk is found.
          found_chunk = &(*it);
          break;
        }
      }
    }
    found_chunk->Deallocate(p, block_size_);
    if (found_chunk->available_blocks_count_ == blocks_num_) {
      // Chunk is free. Erase it.
      vector<Chunk>::iterator to_delete = std::find(chunks_.begin(),
                                                    chunks_.end(),
                                                    *found_chunk);
      to_delete->Release();
      chunks_.erase(to_delete);
    }
  }
};

class SmallObjAllocator {
public:
  SmallObjAllocator(size_t chunk_size, size_t max_obj_size):
    chunk_size_(chunk_size), max_obj_size_(max_obj_size) {}

  void* Allocate(size_t sz) {
    if (sz > max_obj_size_) {
      return ::operator new(sz);
    }
    // Find FixedAllocator.
    FixedAllocator* found_allocator = NULL;
    for (vector<FixedAllocator>::iterator it = pool_.begin();
         it != pool_.end(); ++it) {
      if (it->block_size_ == sz) {
        found_allocator = &(*it);
        break;
      }
    }
    if (found_allocator == NULL) {
      // FixedAllocator is not found.
      // Add new FixedAllocator.
      pool_.push_back(FixedAllocator(chunk_size_, kBlocksInChunk));
      found_allocator = &(pool_.back());
    }

    return found_allocator->Allocate();
  }

  void Deallocate(void* p, size_t sz) {
    if (sz > max_obj_size_) {
      ::operator delete(p);
      return;
    }
    // Find FixedAllocator.
    FixedAllocator* found_allocator = NULL;
    for (vector<FixedAllocator>::iterator it = pool_.begin();
         it != pool_.end(); ++it) {
      if (it->block_size_ == sz) {
        found_allocator = &(*it);
        break;
      }
    }
    if (found_allocator == NULL) {
      // FixedAllocator not found!!!
      return;
    }
    found_allocator->Deallocate(p);
  }

private:
  vector<FixedAllocator> pool_;
  size_t chunk_size_;
  size_t max_obj_size_;
};

class SmallObjAllocatorSingleton {
public:
  static SmallObjAllocator* getInstance() {
    static SmallObjAllocator* instance_ = NULL; // The only instance.
    if (instance_ == NULL) {
      instance_ = new SmallObjAllocator(kBlockSize, kMaxSizeOfSmallObject);
    }
    return instance_;
  }
private:
  SmallObjAllocatorSingleton(); // No new objects.
  SmallObjAllocatorSingleton(const SmallObjAllocatorSingleton&); // No copies.
};

class SmallObject {
public:
  static void* operator new(size_t sz) {
    return SmallObjAllocatorSingleton::getInstance()->Allocate(sz);
  }

  static void operator delete(void* p, size_t sz) {
    SmallObjAllocatorSingleton::getInstance()->Deallocate(p, sz);
  }

  virtual ~SmallObject() {}
};

// STL Allocator
template <class T> class small_obj_allocator {
private:
  SmallObjAllocator soa;
public:
  typedef T value_type;
  typedef size_t size_type;
  typedef ptrdiff_t difference_type;
  typedef T* pointer;
  typedef const T* const_pointer;
  typedef T& reference;
  typedef const T& const_reference;

  pointer address(reference r) const {
    return &r;
  }
  const_pointer address(const_reference r) const {
    return &r;
  }

  small_obj_allocator() throw():
    soa(SmallObjAllocator(kBlockSize, kMaxSizeOfSmallObject)) {}

  template<class U> small_obj_allocator(const small_obj_allocator<U>& other)
  throw() {
    soa = other.soa;
  }

  ~small_obj_allocator() throw() {}

  // Allocate memory for n objects T:
  pointer allocate(size_type n) {
    soa.Allocate(n * sizeof(T));
  }

  void deallocate(pointer p, size_type n) {
    soa.Deallocate(p, n * sizeof(T));
  }

  void construct(pointer p, const T& val) {
    new(static_cast<void*>(p)) T(val);
  }

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

  size_type max_size() const throw() {
    return 100 * 1024 * 1024; // 100 MB
  }

  template<class U> struct rebind {
    typedef small_obj_allocator<U> other;
  };
};

// Derivative from SmallObject example:
class ExampleObject: public SmallObject {
public:
  ExampleObject(int i): x(i) {
  }

  ~ExampleObject() {}

  int get_x() {
    return x;
  }

  void set_x(int i) {
    x = i;
  }

private:
  int x;
};

int main () {
  // ExampleObject test:
  vector<ExampleObject> vec;
  for (int i = 0; i < 10000; ++i) {
    vec.push_back(ExampleObject(i));
  }
  bool test_passed = true;
  int i = 0;
  for (vector<ExampleObject>::iterator it = vec.begin();
       it != vec.end(); ++it, ++i) {
    if (it->get_x() != i) {
      std::cout << "ExampleObject test failed :(" << std::endl;
      test_passed = false;
      break;
    }
  }
  if (test_passed) {
    std::cout << "ExampleObject test passed." << std::endl;
  }

  // STL allocator test:
  vector<int, small_obj_allocator<int> > vec2;
  for (i = 0; i < 10000; ++i) {
    vec2.push_back(i);
  }
  test_passed = true;
  i = 0;
  for (vector<int, small_obj_allocator<int> >::iterator it = vec2.begin();
       it != vec2.end(); ++it, ++i) {
    if (*it != i) {
      std::cout << "STL allocator test failed :(" << std::endl;
      test_passed = false;
      break;
    }
  }
  if (test_passed) {
    std::cout << "STL allocator test passed." << std::endl;
  }
  return 0;
}