A tiny, malloc-based GC that uses a radix tree to store GC metadata (1.5% - 3% overhead)

radix_tree_gc.c

#include <assert.h>
#include <stdbool.h>
#include <stdlib.h>
#include <stdio.h>
#include <stdint.h>
#include <string.h>

#define BITMAP_NUM_BITS 64
#define LG_BITMAP_NUM_BITS 6

#define BITMAP_NUM_BITS_MASK (BITMAP_NUM_BITS - 1)

void bitmap_set(uint64_t* bitmap, size_t index) {
  size_t i = index >> LG_BITMAP_NUM_BITS;
  bitmap[i] |= (((uint64_t)1) << (index & BITMAP_NUM_BITS_MASK));
}

#define VADDR_BITS 48
#define LSB_BITS 18
#define MSB_BITS (VADDR_BITS - LSB_BITS)
#define HEIGHT 3
#define OBJECT_SIZE 16
#define LG_OBJECT_SIZE 4
#define NUM_OBJECTS_PER_LEAF ((1<<LSB_BITS)/OBJECT_SIZE)
#define NUM_OBJECTS_WORDS_PER_LEAF (NUM_OBJECTS_PER_LEAF/64)

typedef struct RNode_ RNode;
typedef struct RLeaf_ RLeaf;

struct RNode_ {
  union {
    RNode* nodes;
    RLeaf* leaf;
  };
};

struct RLeaf_ {
  uint64_t objects[NUM_OBJECTS_WORDS_PER_LEAF];
  uint64_t marks[NUM_OBJECTS_WORDS_PER_LEAF];
  uintptr_t prefix;
  RLeaf* next;
};

typedef struct {
  uintptr_t bits;
  uintptr_t bitshift;
} RLevel;

typedef struct {
  RNode nodes[1<<(MSB_BITS/HEIGHT)];
  RLevel levels[HEIGHT + 1];
  RLeaf *leaves;
  size_t num_leaves;
  size_t metadata_bytes;
} RTree;

static void tree_init(RTree *tree) {
  memset(tree->nodes, 0, sizeof(tree->nodes));
  uintptr_t msb_bits = MSB_BITS;
  uintptr_t accumbits = 0;
  uintptr_t bits;
  for (int i = HEIGHT; i >= 0; --i) {
    if (i < HEIGHT) {
      bits = MSB_BITS/HEIGHT;
      if (msb_bits > bits) {
        msb_bits -= bits;
      } else {
        bits = msb_bits;
      }
    } else {
      bits = LSB_BITS;
    }
    tree->levels[i].bits = bits;
    tree->levels[i].bitshift = accumbits;
    accumbits += bits;
  }
  tree->leaves = NULL;
  tree->num_leaves = 0;
  tree->metadata_bytes = 0;
}

static uintptr_t index_for_level(RLevel level, uintptr_t key) {
  return (key >> level.bitshift) & ((((uintptr_t)1) << level.bits) - 1);
}

static RLeaf* find_leaf(RTree *tree, uintptr_t key, bool create_if_absent) {
  RNode *nodes = tree->nodes;
  uintptr_t index;
  for (int i = 0; i < HEIGHT - 1; ++i) {
    index = index_for_level(tree->levels[i], key);
    assert(index < (1 << tree->levels[i].bits));
    if (nodes[index].nodes == NULL) {
      if (create_if_absent) {
        size_t size = sizeof(RNode) * (1 << tree->levels[i].bits);
        tree->metadata_bytes += size;
        nodes[index].nodes = malloc(size);
        memset(nodes[index].nodes, 0, size);
      } else {
        abort();
        return NULL;
      }
    }
    nodes = nodes[index].nodes;
  }
  index = index_for_level(tree->levels[HEIGHT - 1], key);
  assert(index < (1 << tree->levels[HEIGHT - 1].bits));
  if (nodes[index].leaf == NULL) {
    if (create_if_absent) {
      size_t size = sizeof(RLeaf);
      tree->metadata_bytes += size;
      tree->num_leaves++;
      RLeaf *leaf = malloc(size);
      memset(leaf, 0, size);
      nodes[index].leaf = leaf;
      leaf->prefix = key & ~(((uintptr_t)1 << tree->levels[HEIGHT].bits) - 1);
      leaf->next = tree->leaves;
      tree->leaves = leaf;
    } else {
      abort();
    }
  }
  return nodes[index].leaf; 
}

typedef struct {
  RTree tree;
} GC;

void gc_init(GC* gc) {
  tree_init(&gc->tree);
}

void* gc_alloc(GC* gc, size_t size) {
  if (size < OBJECT_SIZE) {
    size = OBJECT_SIZE;
  }
  void* ptr = malloc(size);
  uintptr_t key = (uintptr_t)ptr;
  RLeaf* leaf = find_leaf(&gc->tree, key, true);
  if (leaf == NULL) abort();
  uintptr_t index = index_for_level(gc->tree.levels[HEIGHT], key) >> LG_OBJECT_SIZE;
  bitmap_set(leaf->objects, index);
  return ptr;
}

void gc_mark(GC* gc, void* ptr) {
  uintptr_t key = (uintptr_t)ptr;
  RLeaf* leaf = find_leaf(&gc->tree, key, false);
  if (leaf == NULL) return;
  uintptr_t index = index_for_level(gc->tree.levels[HEIGHT], key) >> LG_OBJECT_SIZE;
  bitmap_set(leaf->marks, index);
}

static inline char* to_binary(size_t value, char* buf) {
  for (int b = 0; b < 64; ++b) {
      buf[b] = ((value >> (63 - b)) & 0x1) ? '1' : '0';
  }
  buf[64] = '\0';
  return buf;
}

void gc_sweep(GC* gc) {
  RLeaf* leaf = gc->tree.leaves;
  while (leaf) {
    for (size_t i = 0; i < NUM_OBJECTS_WORDS_PER_LEAF; ++i) {
      int64_t word = (int64_t)leaf->objects[i];
      size_t bitoff = 0;
      size_t bit;
      while (bitoff < 64 && (bit = (size_t)__builtin_ffsl(word >> bitoff)) > 0) {
        size_t bitpos = (bit + bitoff - 1);
        if (!((leaf->marks[i] >> bitpos) & 0x1)) {
          size_t index = (i << LG_BITMAP_NUM_BITS) + bitpos;
          void* ptr = (void*)(leaf->prefix + (index << LG_OBJECT_SIZE));
          printf("free %p\n", ptr);
          free(ptr);
        }
        bitoff += bit;
      }
    }
    for (size_t i = 0; i < NUM_OBJECTS_WORDS_PER_LEAF; ++i) {
      leaf->objects[i] &= leaf->marks[i];
    }
    for (size_t i = 0; i < NUM_OBJECTS_WORDS_PER_LEAF; ++i) {
      leaf->marks[i] = 0;
    }
    leaf = leaf->next;
  }
}

void test_mark_and_sweep(size_t num_allocs) {
  GC gc;
  gc_init(&gc);

  for (size_t i = 0; i < num_allocs; ++i) {
    void* ptr = gc_alloc(&gc, 16);
    printf("alloc %p\n", ptr);
    if (i % 2 == 0) {
      gc_mark(&gc, ptr);
    }
  }
  gc_sweep(&gc);
  gc_sweep(&gc);
}

void test_overhead(size_t num_allocs) {
  RTree tree;
  tree_init(&tree);
  size_t num_bytes = 16;
  for (size_t i = 0; i < num_allocs; ++i) {
    uintptr_t p = (uintptr_t)malloc(num_bytes);
    //uintptr_t p = i * 16;
    find_leaf(&tree, p, true);
  }
  printf("%f %f %f (num_allocs: %zu, num_leaves: %zu, objects_per_page: %d)\n",
         (double)tree.metadata_bytes / (1024 * 1024),
         (double)(num_allocs * num_bytes) / (1024 * 1024),
         (double)tree.metadata_bytes / (num_allocs * num_bytes),
         num_allocs, tree.num_leaves,
         ((1 << LSB_BITS))/OBJECT_SIZE);
}

void test_overshift() {
  uint64_t val = (uint64_t)-6149102341220990976;
  char buf[65];
  printf("%s\n", to_binary(val, buf));
  printf("%s\n", to_binary(val << 62, buf));
  printf("%s\n", to_binary(val << 64, buf));
}

int main(int argc, char** argv) {
  for (size_t i = 0; i < 26; ++i) {
    test_overhead((1 << i));
  }
  test_mark_and_sweep(16);
}