Tries are really neat

While working with some sonar data, I found myself wanting to index a data file. I needed a key-value map where both the keys and the values are 64 bit integers. The values are essentially offsets into the original data file, but the keys are the concatenation of three integer fields from the data record itself. The first byte is the sonar subsystem that the data originated from, the next six bytes are a millisecond-resolution timestamp, and the final byte distinguishes the port from the starboard channel. I organize the keys this way because I have found that a lot of data access patterns require accessing the two channels of data from each timestamp for a particular subsystem. To scan the index in this way, we need to be able to iterate over keys that share a common prefix. A trie is a natural fit for this problem.

Tries are trees where the nodes represent common prefixes of the keys, and the edges represent the next symbol in the key. With integer keys, we can interpret the key as a string of bits, so that a very simple trie node consists of two pointers, one of which points to the subtrie where the prefix of the present node is followed by a 0 bit and one of which points to the subtrie where the prefix is followed by a 1 bit. In C, this might look something like

typedef struct trie trie;
struct trie {
    trie *children[2];
};

A 64 bit integer key can be found in the trie by scanning from the root node, which represents the empty prefix, and following 64 pointers by consuming bits from the key starting from the most significant bit. If we ever hit a NULL pointer, the key is not present, and we break out from the loop early. We don't need to store the keys in the trie, because they are implicitly represented by the path from the root node.

trie *trie_lookup(trie *t, uint64_t key)
{
    for (int32_t level = 63; level >=0; level--) {
        uint8_t child = (key >> level) & 1;
        if (!t->children[child]) {
            return NULL;
        }
        t = t->children[(key >> level) & 1];
    }
    return t;
}

trie_lookup returns a pointer to a trie, and a NULL pointer indicates that the key was not found. Note that this requires that the child pointers at level 0 are non-zero. They don't necessarily need to point to a valid trie, if we can ensure that we never dereference the pointer returned from trie_lookup. However, it is helpful when we start storing values in the trie to instead allocate a leaf node at a fictitious level equivalent to -1, so that the level 0 child pointers are either NULL or valid pointers to leaf nodes. Insertion is very similar to lookup.

trie *new_node();
trie *trie_insert(trie *t, uint64_t key)
{
    for (int32_t level = 64; level>=0; level--) {
        uint8_t child = (key >> level) & 1;
        if (!t->children[child]) {
            t->children[child] = new_node();
        }
        t = t->children[child];
    }
}

The function new_node depends on your chosen strategy for memory allocation. One could simply use malloc(sizeof(trie)), but I would usually allocate from a pool of trie nodes or a linear allocator. A simple pool allocator is just a buffer of trie nodes that tracks the next available node.

typedef struct {
    trie* buffer;
    ptrdiff_t capacity;
    ptrdiff_t count;
} trie_pool;

trie *pool_allocate(trie_pool *pool)
{
    assert(pool->count < pool->capacity);
    return pool->buffer[pool->count++];
}

Before considering how to iterate over a trie, we can store values in the trie by changing the trie to a union.

typedef union trie trie;
union trie {
    trie *children[2];
    struct {
        uint64_t key;
        uint64_t value;
    };
};

The nodes on the fictitious -1 level will be leaf nodes with key-value pairs stored instead of child pointers. Rather than tagging the union, we will distinguish the internal from the leaf nodes by their level, which we will always keep track of outside of the data structure. While we don't need to store the keys, there is a spare 8 bytes in the leaf nodes that we might as well fill up, and storing the keys prevents us from having to keep track of the keys while iterating.

Our trie is a binary tree and iterating over it amounts to a depth-first tree traversal, which means we need a stack.

typedef struct {
    trie *stack[65];
    int32_t levels[65];
    int32_t top;
} trie_iterator;

void trie_push(trie_iterator *iter, trie *node, int32_t level)
{
    assert(iter->top < 65);
    iter->levels[iter->top] = level;
    iter->stack[iter->top++] = node;
}

int32_t trie_pop(trie *node, trie_iterator *iter)
{
    assert(iter->top > 0);
    node = iter->stack[--iter->top];
    return iter->levels[iter->top];
}

The trie has a fixed depth, so the stack is bounded by the number of levels, which is 64 plus the leaf level, so 65 nodes will suffice. We maintain a stack of trie pointers and their corresponding levels. If we didn't store the keys in the leaves, we would also want to have a stack of keys, which we would reconstruct as we progress through the trie.

We initialize the iterator by pushing the root node at level 63:

trie_iterator trie_begin(trie *root)
{
    trie_iterator iter = {0}; 
    trie_push(&iter,root,63);
    return iter;
}

To iterate, we pop nodes from the stack and push their children.

trie *trie_iterate(trie_iterator *iter, int32_t end_level)
{
    trie *node;
    int32_t level;
    while (iter->top > 0) {
        level = trie_pop(node,iter);
        if (level < end_level) {
            return node;
        }
        if (node->children[1]) {
            trie_push(iter,node->children[1],level - 1);
        }
        if (node->children[0]) {
            trie_push(iter,node->children[0],level - 1);
        }
    }
    return NULL;
}

We use the end_level parameter to control which nodes we iterate over. If we want to iterate over the leaf nodes, we pass 0 for end_level. If we want to iterate over prefixes, we pass end_level corresponding to the level of the prefix. For our sonar data example, we might iterate first over the subsystems, then over the timestamps, then the channels.

trie_iterator subsystem_iterator = trie_begin(root);
trie *subsystem;
while ((subsystem = trie_iterate(&subsystem_iterator,56))) {
    trie_iterator timestamp_iterator;
    trie_push(&timestamp_iterator,subsystem,55);
    trie *timestamp;
    while ((timestamp = trie_iterate(&subsystem_iterator,8))) {
        trie_iterator channel_iterator;
        trie_push(&channel_iterator,timestamp,7);
        trie *channel;
        while ((channel = trie_iterate(&channel_iterator,0))) {
            process_data(channel->key,channel->value);
        }
    }
}

And that is all we need to start building and scanning indices for the sonar data. There are a lot of improvements that we can make on this basic design. For example, we can consume more than a single bit of the key at a time, which speeds up lookups while requiring more space to store the child pointers. We can also construct a variety of hybrid tree variants that use other search data structures at each level of the trie. It is also relatively straightforward to implement lock-free concurrent inserts to this simple array-based trie.