Simple implementation of Clojure's PersistentVector in C with Boehm GC (bdw-gc).

persistent_vector.c

#include <stdio.h>
#include <assert.h>
#include <stdbool.h>
#include <string.h>
#include <sys/time.h>
#include <gc.h>

#define MIN(x, y) \
  (((x) < (y)) ? (x) : (y))

#define MAX(x, y) \
  (((x) > (y)) ? (x) : (y))

typedef enum {
  PV_NIL,
  PV_BOOL,
  PV_INTEGER,
  PV_ARRAY,
  PV_VECTOR_NODE,
  PV_PERSISTENT_VECTOR,
  PV_TRANSIENT_VECTOR,
} PV_TYPE;

typedef struct {
  PV_TYPE type;
} Object;

typedef struct {
  PV_TYPE type;
  long val;
} Integer;

Integer *Integer_new(long n) {
  Integer *i = GC_malloc_atomic(sizeof(Integer));
  i->type = PV_INTEGER;
  i->val = n;
  return i;
}

long Integer_value(Integer *i) {
  assert(i->type == PV_INTEGER);
  return i->val;
}

typedef struct {
  PV_TYPE type;
  unsigned count;
  void *ary[0];
} Array;

Array *Array_new(unsigned size) {
  Array *a = GC_malloc(sizeof(Array)+(sizeof(void*)*size));
  a->type = PV_ARRAY;
  a->count = size;
  return a;
}

void *Array_get(Array *a, unsigned idx) {
  assert(idx < a->count);
  return a->ary[idx];
}

void Array_set(Array *a, unsigned idx, void *val) {
  assert(idx < a->count);
  a->ary[idx] = val;
}

Array *Array_copyn(Array *dest, Array* source, unsigned count) {
  count = MIN(dest->count, count);
  count = MIN(source->count, count);

  memcpy(dest->ary, source->ary, count*sizeof(void*));
  dest->count = MAX(count, dest->count); // Don't change count if source was smaller
  return dest;
}

Array *Array_copy_resize(Array *source, unsigned count) {
  return Array_copyn(Array_new(count), source, count);
}

Array *Array_copy(Array *dest, Array* source) {
  return Array_copyn(dest, source, dest->count);
}

Array *Array_clone(Array *source) {
  return Array_copyn(Array_new(source->count), source, source->count);
}

typedef struct {
  PV_TYPE type;
  bool edit;
  Array *array;
} VectorNode;

VectorNode *VectorNode_new(bool edit, Array *array) {
  VectorNode *vn = GC_malloc(sizeof(VectorNode));
  vn->type = PV_VECTOR_NODE;
  
  if (!array)
    array = Array_new(32);

  vn->array = array;

  return vn;
}

typedef struct {
  PV_TYPE type;
  unsigned count;
  unsigned shift;
  VectorNode *root;
  Array *tail;
} PersistentVector;

PersistentVector *PersistentVector_new(unsigned count, unsigned shift, VectorNode *root, Array *tail) {
  PersistentVector *pv = GC_malloc(sizeof(PersistentVector));
  pv->type = PV_PERSISTENT_VECTOR;
  pv->count = count;
  pv->shift = shift;
  pv->root = root;
  pv->tail = tail;
  return pv;
}

unsigned tailoff(PersistentVector *pv) {
  if (pv->count < 32) return 0;
  return ((pv->count - 1) >> 5) << 5;
}

Array *arrayFor(PersistentVector *pv, unsigned i) {
  assert(i < pv->count);

  if (i >= tailoff(pv))
    return pv->tail;

  VectorNode *node = pv->root;
  int level;
  for (level = (int)pv->shift; level > 0; level -= 5)
    node = (VectorNode*)Array_get(node->array, (i >> level) & 0x01f);
  return node->array;
}

void *PersistentVector_nth(PersistentVector *pv, unsigned i, void *nf) {
  if (nf && i >= pv->count)
    return nf;
  return Array_get(arrayFor(pv, i), i & 0x01f);
}

VectorNode *newPath(bool edit, unsigned level, VectorNode *node) {
  if (level == 0)
    return node;
  VectorNode *ret = VectorNode_new(edit, NULL);
  Array_set(ret->array, 0, newPath(edit, level - 5, node));
  return ret;
}

VectorNode *pushTail(PersistentVector *pv, unsigned level, VectorNode *parent, VectorNode *tailnode) {
  unsigned subidx = ((pv->count - 1) >> level) & 0x01f;
  VectorNode *ret = VectorNode_new(parent->edit, Array_clone(parent->array));
  VectorNode *nodeToInsert;
  if (level == 5)
    nodeToInsert = tailnode;
  else {
    VectorNode *child = Array_get(parent->array, subidx);
    if (child)
      nodeToInsert = pushTail(pv, level-5, (VectorNode*)Array_get(parent->array, subidx), tailnode);
    else
      nodeToInsert = newPath(pv->root->edit, level-5, tailnode);
  }
  Array_set(ret->array, subidx, nodeToInsert);
  return ret;
}

PersistentVector *PersistentVector_cons(PersistentVector *pv, void *val) {
  // Room in tail?
  if (pv->count - tailoff(pv) < 32) {
    Array *newTail = Array_copy_resize(pv->tail, pv->tail->count+1);
    Array_set(newTail, newTail->count-1, val);
    return PersistentVector_new(pv->count+1, pv->shift, pv->root, newTail);
  }

  //Full tail push into tree
  VectorNode *newroot;
  VectorNode *tailnode = VectorNode_new(pv->root->edit, Array_clone(pv->tail));

  unsigned newshift = pv->shift;
  //Overflow root?
  if ((pv->count >> 5) > (1 << pv->shift)) {
    newroot = VectorNode_new(pv->root->edit, NULL);
    Array_set(newroot->array, 0, pv->root);
    Array_set(newroot->array, 1, newPath(pv->root->edit, pv->shift, tailnode));
    newshift += 5;
  } else
    newroot = pushTail(pv, pv->shift, pv->root, tailnode);

  Array *newTail = Array_new(1);
  Array_set(newTail, 0, val);
  return PersistentVector_new(pv->count+1, newshift, newroot, newTail);
}

VectorNode *doAssoc(unsigned level, VectorNode *node, unsigned i, void *val) {
  VectorNode *ret = VectorNode_new(node->edit, Array_clone(node->array));
  if (level == 0)
    Array_set(ret->array, i & 0x01f, val);
  else {
    unsigned subidx = (i >> level) & 0x01f;
    Array_set(ret->array, subidx, doAssoc(level-5, (VectorNode*)Array_get(node->array, subidx), i, val));
  }
  return ret;
}

PersistentVector *PersistentVector_assocN(PersistentVector *pv, unsigned i, void *val) {
  assert(i <= pv->count); // <= is ok here

  if (i == pv->count) return PersistentVector_cons(pv, val);

  if (i >= tailoff(pv)) {
    Array *newTail = Array_clone(pv->tail);
    Array_set(newTail, i & 0x01f, val);
    return PersistentVector_new(pv->count, pv->shift, pv->root, newTail);
  }
  return PersistentVector_new(pv->count, pv->shift, doAssoc(pv->shift, pv->root, i, val), pv->tail);
}

static VectorNode *EMPTY_NODE = NULL;
static PersistentVector *EMPTY = NULL;
static Object *NIL = { PV_NIL };

int main() {
  GC_INIT();
  //GC_enable_incremental(); // This is slow as shit.
  GC_expand_hp(512L*1024*1024);

  EMPTY_NODE = VectorNode_new(false, NULL);
  EMPTY = PersistentVector_new(0, 5, EMPTY_NODE, Array_new(0));

  #define VECSIZE 1048609
  //#define VECSIZE 32801
  //#define VECSIZE 1057
  //#define VECSIZE 1024
  //#define VECSIZE 33
  //#define VECSIZE 32
  unsigned c = VECSIZE;
  Integer **items = GC_malloc(VECSIZE*(sizeof(Integer*)));

  unsigned i;
  for (i = 0; i < c; ++i)
    items[i] = Integer_new(i);

  PersistentVector *v = EMPTY;

  struct timeval before, after;
  gettimeofday(&before, NULL);
  unsigned j;
  //for (j = 0; j < 50; ++j)
  for (i = 0; i < c; i++)
    v = PersistentVector_cons(v, items[i]);
  gettimeofday(&after, NULL);
  long b = (long)before.tv_sec * 1000000,
       a = (long)after.tv_sec * 1000000;
  b += before.tv_usec;
  a += after.tv_usec;
  printf("%.4fms\n", (a-b)/1000.0);

  PersistentVector *v2 = v;
  v2 = PersistentVector_assocN(v2, v2->count, Integer_new(42));
  v2 = PersistentVector_assocN(v2, v2->count>>1, Integer_new(42));
  v2 = PersistentVector_assocN(v2, 0, Integer_new(42));

  PersistentVector *vs[] = { v, v2 };
  for (j = 0; j < 2; ++j) {
    PersistentVector *v = vs[j];
    printf("[ ");
    for (i = 0; i < 5; ++i)
      printf("%li ", Integer_value(PersistentVector_nth(v, i, NULL)));
    printf("... ");
    for (i = (v->count>>1)-2; i < (v->count>>1)+3; ++i)
      printf("%li ", Integer_value(PersistentVector_nth(v, i, NULL)));
    printf("... ");
    for (i = v->count-5; i < v->count; ++i)
      printf("%li ", Integer_value(PersistentVector_nth(v, i, NULL)));
    printf("]\n");
  }

  for (i=0; i < v->count; i++)
    assert(PersistentVector_nth(v, i, NULL) == items[i]);
  
  return 0;
}