ec3/src/b-tree.c

#include <blue/core/queue.h>
#include <stdlib.h>
#include <assert.h>
#include <string.h>

#include "b-tree.h"

#define SWAP(type, x, y)                                                       \
	do {                                                                   \
		type tmp = x;                                                  \
		x = y;                                                         \
		y = tmp;                                                       \
	} while (0)

struct cache_entry {
	int c_allocated;
	b_queue_entry c_entry;
	/* ... node data here ... */
};

#define GET_ENTRY(tree, entries, index) \
	(b_tree_node_entry *)((unsigned char *)(entries) + ((index) * (tree)->tree_entry_size));


static b_tree_node *cache_alloc_node(struct b_tree *tree)
{
	b_queue_iterator it = { 0 };
	b_queue_foreach (&it, &tree->tree_cache) {
		struct cache_entry *e = b_unbox(struct cache_entry, it.entry, c_entry);
		if (!e->c_allocated) {
			e->c_allocated = 1;
			return (b_tree_node *)(e + 1);
		}
	}

	struct cache_entry *e = malloc(sizeof *e + tree->tree_node_size);
	if (!e) {
		return NULL;
	}

	memset(e, 0x0, sizeof *e);
	e->c_allocated = 1;
	e->c_entry = B_QUEUE_ENTRY_INIT;

	b_queue_push_back(&tree->tree_cache, &e->c_entry);

	return (b_tree_node *)(e + 1);
}

static void cache_free_node(struct b_tree *tree, b_tree_node *node)
{
	struct cache_entry *e = (struct cache_entry *)((unsigned char *)node - sizeof *e);
	e->c_allocated = 0;
}

void b_tree_init(struct b_tree *tree, const struct b_tree_ops *ops, unsigned int node_size, unsigned int entry_size, unsigned int order)
{
	memset(tree, 0x0, sizeof *tree);
	assert((order % 2) == 0);
	tree->tree_ops = ops;
	tree->tree_node_size = node_size;
	tree->tree_entry_size = entry_size;
	tree->tree_order = order;
}

static long tree_alloc_id(struct b_tree *tree)
{
	return tree->tree_ops->tree_alloc_node(tree);
}

static int node_read(struct b_tree *tree, long id, b_tree_node *out)
{
	return tree->tree_ops->tree_get_node(tree, id, out);
}

static int node_write(struct b_tree *tree, long id, const b_tree_node *in)
{
	return tree->tree_ops->tree_put_node(tree, id, in);
}

static unsigned long node_get_nr_entries(struct b_tree *tree, b_tree_node *node)
{
	return tree->tree_ops->node_get_nr_entries(node);
}

static void node_set_nr_entries(struct b_tree *tree, b_tree_node *node, unsigned long val)
{
	tree->tree_ops->node_set_nr_entries(node, val);
}

static long node_get_child(struct b_tree *tree, b_tree_node *node, unsigned long index)
{
	uint16_t *children = tree->tree_ops->node_get_children(node);
	if (children[index] == 0xFFFF) {
		return B_TREE_INVALID_PTR;
	}

	return children[index];
}

static void node_set_child(struct b_tree *tree, b_tree_node *node, unsigned long index, long ptr)
{
	uint16_t *children = tree->tree_ops->node_get_children(node);
	if (ptr == B_TREE_INVALID_PTR) {
		children[index] = 0xFFFF;
	} else {
		children[index] = ptr;
	}
}

static b_tree_node_entry *node_get_entry(struct b_tree *tree, b_tree_node *node, unsigned long index)
{
	b_tree_node_entry *entries = tree->tree_ops->node_get_entries(node);
	return GET_ENTRY(tree, entries, index);
}

static void node_set_entry(struct b_tree *tree, b_tree_node *node, unsigned long index, const b_tree_node_entry *entry)
{
	b_tree_node_entry *entries = tree->tree_ops->node_get_entries(node);
	b_tree_node_entry *dest = GET_ENTRY(tree, entries, index);
	memcpy(dest, entry, tree->tree_entry_size);
}

static void node_kill_entry(struct b_tree *tree, b_tree_node *node, unsigned long index)
{
	b_tree_node_entry *entries = tree->tree_ops->node_get_entries(node);
	b_tree_node_entry *dest = GET_ENTRY(tree, entries, index);
	memset(dest, 0x0, tree->tree_entry_size);
}

static void node_shift_entries(struct b_tree *tree, b_tree_node *node, unsigned long shift_from, long shift_by)
{
	unsigned long nr_entries = node_get_nr_entries(tree, node);
	b_tree_node_entry *entries = tree->tree_ops->node_get_entries(node);

	if (shift_from >= nr_entries) {
		return;
	}

	b_tree_node_entry *src = GET_ENTRY(tree, entries, shift_from);
	b_tree_node_entry *dest = GET_ENTRY(tree, entries, shift_from + shift_by);
	unsigned int count = nr_entries - shift_from;

	memmove(dest, src, tree->tree_entry_size * count);
	if (shift_by < 0) {
		dest = GET_ENTRY(tree, entries, nr_entries + shift_by);
		count = abs(shift_by);
	}
	else {
		dest = src;
		count = shift_by;
	}

	memset(dest, 0x0, count * tree->tree_entry_size);
	return;
}

static void node_shift_children(struct b_tree *tree, b_tree_node *node, unsigned long shift_from, long shift_by)
{
	unsigned long nr_entries = node_get_nr_entries(tree, node);
	uint16_t *children = tree->tree_ops->node_get_children(node);

	if (shift_from >= nr_entries + 1) {
		return;
	}

	if (shift_from == nr_entries && shift_by > 0) {
		return;
	}

	uint16_t *src = &children[shift_from];
	uint16_t *dest = &children[shift_from + shift_by];
	unsigned int count = nr_entries + 1 - shift_from;
	if (shift_by > 0) {
		count -= shift_by;
	}

	memmove(dest, src, sizeof(uint16_t) * count);
	if (shift_by < 0) {
		dest = &children[nr_entries + 1 + shift_from];
		count = abs(shift_by);
	}
	else {
		dest = src;
		count = shift_by;
	}

	memset(dest, 0xFF, count * sizeof(uint16_t));
}

static void node_move_entries(struct b_tree *tree, b_tree_node *from, b_tree_node *to, unsigned long src_index, unsigned long dest_index, unsigned long count)
{
	unsigned long from_nr_entries = node_get_nr_entries(tree, from);
	unsigned long to_nr_entries = node_get_nr_entries(tree, to);

	b_tree_node_entry *from_entries = tree->tree_ops->node_get_entries(from);
	b_tree_node_entry *to_entries = tree->tree_ops->node_get_entries(to);

	b_tree_node_entry *src = GET_ENTRY(tree, from_entries, src_index);
	b_tree_node_entry *dest = GET_ENTRY(tree, to_entries, dest_index);

	memmove(dest, src, count * tree->tree_entry_size);
	if (src_index + count >= from_nr_entries) {
		memset(src, 0x0, count * tree->tree_entry_size);
	} else {
		node_shift_entries(tree, from, src_index + count, (int)count * -1);
	}

	from_nr_entries -= count;
	to_nr_entries += count;

	node_set_nr_entries(tree, from, from_nr_entries);
	node_set_nr_entries(tree, to, to_nr_entries);
}

static void node_move_children(struct b_tree *tree, b_tree_node *from, b_tree_node *to, unsigned long src_index, unsigned long dest_index, unsigned long count)
{
	unsigned long from_nr_entries = node_get_nr_entries(tree, from);

	uint16_t *from_children = tree->tree_ops->node_get_children(from);
	uint16_t *to_children = tree->tree_ops->node_get_children(to);
	
	uint16_t *src = &from_children[src_index];
	uint16_t *dest = &to_children[dest_index];

	memmove(dest, src, count * sizeof(uint16_t));
	if (src_index + count >= from_nr_entries + 1) {
		memset(src, 0xFF, count * sizeof(uint16_t));
	}
	else {
		node_shift_children(tree, from, src_index + count, (int)count * -1);
	}
}

static int node_put(struct b_tree *tree, b_tree_node *n, const b_tree_node_entry *e)
{
	unsigned long nr_entries = node_get_nr_entries(tree, n);
	if (nr_entries == 0) {
		node_set_entry(tree, n, 0, e);
		node_set_nr_entries(tree, n, 1);
		return 0;
	}

	int insert_at = -1;
	for (unsigned long i = 0; i < nr_entries; i++) {
		const b_tree_node_entry *cur = node_get_entry(tree, n, i);
		int cmp = entry_compare(tree, cur, e);
		if (cmp == 1) {
			insert_at = i;
			break;
		}
	}

	if (insert_at != -1) {
		node_shift_entries(tree, n, insert_at, 1);
	}
	else {
		insert_at = nr_entries;
	}

	node_set_entry(tree, n, insert_at, e);
	nr_entries++;
	node_set_nr_entries(tree, n, nr_entries);
	return insert_at;
}

static int node_is_leaf(struct b_tree *tree, b_tree_node *node)
{
	long first_child = node_get_child(tree, node, 0);
	return first_child == B_TREE_INVALID_PTR;
}

static void node_init(struct b_tree *tree, b_tree_node *out)
{
	memset(out, 0x0, tree->tree_node_size);
	for (unsigned int i = 0; i < tree->tree_order; i++) {
		node_set_child(tree, out, i, B_TREE_INVALID_PTR);
	}
}

static int entry_compare(struct b_tree *tree, const b_tree_node_entry *a, const b_tree_node_entry *b)
{
	return tree->tree_ops->entry_compare(a, b);
}

int b_tree_put(struct b_tree *tree, const b_tree_node_entry *to_put)
{
	int depth = 0;

	long current_id = tree->tree_root;
	long next_id = -1;
	
	b_tree_node *current = cache_alloc_node(tree);
	b_tree_node *next = cache_alloc_node(tree);

	int err = node_read(tree, current_id, current);

	if (err != 0) {
		return err;
	}

	unsigned long nr_entries = node_get_nr_entries(tree, current);

	if (nr_entries == tree->tree_order - 1) {
		/* root node is full. split it pre-emptively. */
		b_tree_node *a = next;
		b_tree_node *b = cache_alloc_node(tree);
		
		node_init(tree, a);
		node_init(tree, b);

		int median = nr_entries / 2;

		node_move_entries(tree, current, a, 0, 0, median);
		node_move_entries(tree, current, b, 1, 0, median);

		node_move_children(tree, current, a, 0, 0, median + 1);
		node_move_children(tree, current, b, median + 1, 0, median + 1);

		int a_index = tree_alloc_id(tree);
		node_write(tree, a_index, a);

		int b_index = tree_alloc_id(tree);
		node_write(tree, b_index, b);

		node_set_child(tree, current, 0, a_index);
		node_set_child(tree, current, 1, b_index);

		node_write(tree, current_id, current);

		cache_free_node(tree, b);
	}

	bool leaf = node_is_leaf(tree, current);

	while (!leaf) {
		unsigned int i;
		bool found_bigger_key = false;

		nr_entries = node_get_nr_entries(tree, current);
		next_id = node_get_child(tree, current, nr_entries);

		for (i = 0; i < nr_entries; i++) {
			b_tree_node_entry *entry = node_get_entry(tree, current, i);
			int cmp = entry_compare(tree, entry, to_put);

			if (cmp == 1) {
				next_id = node_get_child(tree, current, i);
				break;
			}
		}

		err = node_read(tree, next_id, next);
		if (err != 0) {
			return err;
		}

		nr_entries = node_get_nr_entries(tree, next);
		if (nr_entries < tree->tree_order - 1) {
			/* swap current and next pointers. */
			SWAP(b_tree_node *, current, next);
			current_id = next_id;

			leaf = node_is_leaf(tree, current);

			depth++;
			continue;
		}

		/* this node is full, split it. */

		/* half the nodes from `next` will be moved into `n3`. the other
		   half will stay put. the median entry will be moved to
		   `current`.

		   `n3` will become the right sibling of `next` */
		b_tree_node *n3 = cache_alloc_node(tree);
		node_init(tree, n3);

		bool n3_is_next = false;
		nr_entries = node_get_nr_entries(tree, next);
		int median = nr_entries / 2;
		b_tree_node_entry *median_node = node_get_entry(tree, next, median);
		int cmp = entry_compare(tree, median_node, to_put);
		if (cmp == -1) {
			// median++;
			n3_is_next = true;
		}

		node_move_entries(
			tree,
			next,
			n3,
			median + 1,
			0,
			nr_entries - median - 1);
		node_move_children(
			tree,
			next,
			n3,
			median + 1,
			0,
			nr_entries - median);
		
		int index = node_put(tree, current, median_node);
		node_kill_entry(tree, next, median);
		nr_entries = node_get_nr_entries(tree, next);
		nr_entries--;
		node_set_nr_entries(tree, next, nr_entries);
		node_shift_children(tree, current, index + 1, 1);

		int n3_id = tree_alloc_id(tree);

		node_set_child(tree, current, index + 1, n3_id);
		node_write(tree, current_id, current);
		node_write(tree, next_id, next);
		node_write(tree, n3_id, n3);

		if (n3_is_next) {
			cache_free_node(tree, current);
			current = n3;
			current_id = n3_id;
		}
		else {
			SWAP(b_tree_node *, current, next);
			current_id = next_id;
		}

		depth++;
		leaf = node_is_leaf(tree, current);
	}

	node_put(tree, current, to_put);
	node_write(tree, current_id, current);

	return 0;
}