alloc.h

   1 ////////////////////////////////////////////////////////////////////////////////
   2 //
   3 // Pythonc memory allocator
   4 //
   5 // Copyright 2011 Zach Wegner
   6 //
   7 // Licensed under the Apache License, Version 2.0 (the "License");
   8 // you may not use this file except in compliance with the License.
   9 // You may obtain a copy of the License at
  10 //
  11 //   http://www.apache.org/licenses/LICENSE-2.0
  12 //
  13 // Unless required by applicable law or agreed to in writing, software
  14 // distributed under the License is distributed on an "AS IS" BASIS,
  15 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  16 // See the License for the specific language governing permissions and
  17 // limitations under the License.
  18 //
  19 ////////////////////////////////////////////////////////////////////////////////
  20
  21 #include <stddef.h>
  22 #include <assert.h>
  23
  24 #define BLOCK_SIZE (1 << 14)
  25 #define CHUNK_SIZE (1 << 22)
  26
  27 typedef unsigned char byte;
  28
  29 uint32_t bitscan64(uint64_t r) {
  30    asm ("bsfq %0, %0" : "=r" (r) : "0" (r));
  31    return r;
  32 }
  33
  34 template <uint64_t obj_size>
  35 class arena_block {
  36 public:
  37     static const uint64_t capacity = BLOCK_SIZE - sizeof(void *);
  38     static const uint64_t n_objects = (capacity * 64 / (obj_size * 64 + 8)) & ~63;
  39     byte data[obj_size * n_objects];
  40     arena_block<obj_size> *next_block;
  41     uint64_t live_bits[n_objects / 64];
  42     byte padding[capacity - (obj_size * n_objects + n_objects / 8)];
  43
  44     void init() {
  45         assert(sizeof(arena_block<obj_size>) == BLOCK_SIZE);
  46         assert(sizeof(this->live_bits) * 8 == n_objects);
  47         mark_dead();
  48     }
  49     void mark_dead() {
  50         for (uint32_t t = 0; t < n_objects / 64; t++)
  51             this->live_bits[t] = 0;
  52     }
  53     void *get_obj() {
  54         for (uint32_t t = 0; t < n_objects / 64; t++) {
  55             uint64_t dead = ~this->live_bits[t];
  56             if (dead) {
  57                 uint32_t bit = bitscan64(dead);
  58                 uint32_t idx = t * 64 + bit;
  59                 this->live_bits[t] |= (1ull << bit);
  60                 byte *b = &this->data[idx * obj_size];
  61                 return (void *)b;
  62             }
  63         }
  64         return NULL;
  65     }
  66     bool mark_live(uint32_t idx) {
  67         uint32_t t = idx / 64;
  68         uint64_t bit = 1ull << (idx & 63);
  69         bool already_live = (this->live_bits[t] & bit) != 0ull;
  70         this->live_bits[t] |= bit;
  71         return already_live;
  72     }
  73 };
  74
  75 // Silly preprocessor stuff for "cleanly" handling multiple block sizes.
  76 #define FOR_EACH_OBJ_SIZE(x) x(16) x(24) x(32) x(40) x(56) x(64) x(72)
  77
  78 class arena {
  79 private:
  80 #define DECL_BLOCK(size) arena_block<size> *head_##size;
  81     FOR_EACH_OBJ_SIZE(DECL_BLOCK)
  82 #undef DECL_BLOCK
  83     byte *chunk_start, *chunk_end;
  84
  85 public:
  86     arena() {
  87         this->get_new_chunk();
  88 #define INIT_BLOCK(size) \
  89         this->head_##size = (arena_block<size> *)this->new_block(); \
  90         this->head_##size->init();
  91
  92         FOR_EACH_OBJ_SIZE(INIT_BLOCK)
  93 #undef INIT_BLOCK
  94     }
  95
  96     void get_new_chunk() {
  97         this->chunk_start = new byte[CHUNK_SIZE];
  98         this->chunk_end = chunk_start + CHUNK_SIZE;
  99         // Align the start of the chunk
 100         this->chunk_start = (byte *)(((uint64_t)this->chunk_start + BLOCK_SIZE - 1) & ~(BLOCK_SIZE - 1));
 101     }
 102
 103     void *new_block() {
 104         if (this->chunk_end - this->chunk_start < BLOCK_SIZE)
 105             this->get_new_chunk();
 106
 107         void *block = this->chunk_start;
 108         this->chunk_start += BLOCK_SIZE;
 109         return block;
 110     }
 111
 112     void *allocate(uint64_t bytes) {
 113         switch (bytes) {
 114 #define OBJ_CASE(size) \
 115             case size: { \
 116                 arena_block<size> *block = this->head_##size; \
 117                 void *p = block->get_obj(); \
 118                 if (!p) { \
 119                     block = (arena_block<size> *)this->new_block(); \
 120                     block->init(); \
 121                     block->next_block = this->head_##size; \
 122                     this->head_##size = block; \
 123                     p = block->get_obj(); \
 124                 } \
 125                 return p; \
 126             }
 127
 128             FOR_EACH_OBJ_SIZE(OBJ_CASE)
 129 #undef OBJ_CASE
 130             default:
 131                 printf("bad obj size %" PRIu64 "\n", bytes);
 132                 exit(1);
 133         }
 134         return NULL;
 135     }
 136     void mark_dead() {
 137 #define MARK_DEAD(size) \
 138         for (arena_block<size> *p = this->head_##size; p; p = p->next_block) \
 139             p->mark_dead();
 140
 141         FOR_EACH_OBJ_SIZE(MARK_DEAD)
 142 #undef MARK_DEAD
 143     }
 144     bool mark_live(void *object, uint64_t bytes) {
 145         uint32_t idx = ((uint64_t)object & (BLOCK_SIZE - 1)) / bytes;
 146         void *block = (void *)((uint64_t)object & ~(BLOCK_SIZE - 1));
 147         switch (bytes) {
 148 #define OBJ_CASE(size) case size: return ((arena_block<size> *)block)->mark_live(idx);
 149             FOR_EACH_OBJ_SIZE(OBJ_CASE)
 150 #undef OBJ_CASE
 151             default:
 152                 printf("bad obj size %" PRIu64 "\n", bytes);
 153                 exit(1);
 154         }
 155     }
 156 };
 157
 158 arena allocator[1];
 159
 160 void *operator new(size_t bytes, arena *a) {
 161     return a->allocate(bytes);
 162 }