Hash trie iterator example (integer set)

intset.c

#include <stddef.h>
#include <stdint.h>
#include <stdlib.h>
#include <string.h>

#define assert(c)  while (!(c)) *(volatile int *)0 = 0
#define new(a, t)  (t *)alloc(a, sizeof(t))

typedef struct {
    char *beg, *end;
} arena;

static void *alloc(arena *a, ptrdiff_t size)
{
    ptrdiff_t available = a->end - a->beg;
    if (available < size) {
        assert(0);
    }
    return memset(a->end -= size, 0, size);
}

typedef struct intset intset;
struct intset {
    intset *child[4];
    int     element;
};

static _Bool insert(intset **s, int element, arena *perm)
{
    for (uint64_t h = element*1111111111111111111u; *s; h <<= 2) {
        if ((*s)->element == element) {
            return 0;
        }
        s = &(*s)->child[h>>62];
    }
    (*s = new(perm, intset))->element = element;
    return 1;
}

typedef struct iternode iternode;
struct iternode {
    iternode *next;
    intset   *node;
    int       index;
};

typedef struct {
    iternode *head;
    iternode *free;
} iter;

static iter *newiter(intset *s, arena *perm)
{
    iter *it = new(perm, iter);
    if (s) {
        it->head = new(perm, iternode);
        it->head->node = s;
    }
    return it;
}

typedef struct {
    int   element;
    _Bool ok;
} iterresult;

static iterresult iternext(iter *it, arena *perm)
{
    iterresult r = {0};
    while (it->head) {
        int index = it->head->index++;
        if (index == 0) {
            r.element = it->head->node->element;
            r.ok = 1;
            return r;
        } else if (index == 5) {
            iternode *dead = it->head;
            it->head = dead->next;
            dead->next = it->free;
            it->free = dead;
        } else if (it->head->node->child[index-1]) {
            iternode *new = it->free;
            if (new) {
                it->free = it->free->next;
                new->index = 0;
            } else {
                new = new(perm, iternode);
            }
            new->node = it->head->node->child[index-1];
            new->next = it->head;
            it->head = new;
        }
    }
    return r;
}

static int randint(uint64_t *rng, int lo, int hi)
{
    *rng = *rng*0x3243f6a8885a308d + 1;
    return (unsigned)((*rng>>32)*(hi - lo)>>32) + lo;
}

static arena newarena(ptrdiff_t cap)
{
    arena a = {0};
    a.end = (a.beg = malloc(cap)) + cap;
    return a;
}

int main(void)
{
    uint64_t rng = 1;
    arena perm = newarena(1<<20);

    {
        // Test empty iteration
        arena scratch = perm;
        iter *it = newiter(0, &scratch);
        iterresult r = iternext(it, &scratch);
        assert(!r.ok);
    }

    for (int n = 0; n < 10000; n++) {
        intset *set = 0;
        int64_t seen = 0;
        arena scratch = perm;

        // Populate the intset, checking against a bitset
        for (int i = 0; i < 1000; i++) {
            int element = randint(&rng, 1, 61);
            int64_t bit = (int64_t)1 << element;
            _Bool result = insert(&set, element, &scratch);
            assert(result == !(seen & bit));
            seen |= bit;
        }

        // Traverse the intset, checking against a bitset
        int64_t visited = 0;
        for (iter *it = newiter(set, &scratch);;) {
            iterresult r = iternext(it, &scratch);
            if (!r.ok) {
                break;
            }
            int64_t bit = (int64_t)1 << r.element;
            assert(!(visited & bit));
            visited |= bit;
        }
        assert(seen == visited);
    }
}