Tree bitmap thingy

2 files changed, 156 insertions, 0 deletions

diff --git a/Makefile b/Makefile
new file mode 100644
index 0000000..00fd0ba
--- /dev/null
+++ b/Makefile
@@ -0,0 +1 @@
+CFLAGS=-Wall -g -O2
diff --git a/bitmap.c b/bitmap.c
new file mode 100644
index 0000000..5e7bcb5
--- /dev/null
+++ b/bitmap.c
@@ -0,0 +1,155 @@
+#include <assert.h>
+#include <stdint.h>
+#include <stdio.h>
+#include <stdlib.h>
+#include <string.h>
+
+#define TREE_ARY	8
+#define BITS_PER_LEAF	(TREE_ARY * 64)
+
+struct bitmap_tree {
+	uint64_t	size;
+	uint64_t	nr_unset;
+	size_t		first_leaf;
+
+	struct bitmap_tree_node {
+		uint64_t nr_unset[TREE_ARY];
+	} tree[];
+};
+
+unsigned popcount_64(uint64_t x)
+{
+	static const uint64_t m1  = 0x5555555555555555LLU;
+	static const uint64_t m2  = 0x3333333333333333LLU;
+	static const uint64_t m4  = 0x0f0f0f0f0f0f0f0fLLU;
+	static const uint64_t h01 = 0x0101010101010101LLU;
+
+	x -= (x >> 1) & m1;
+	x = (x & m2) + ((x >> 2) & m2);
+	x = (x + (x >> 4)) & m4;
+	return (x * h01) >> 56;
+}
+
+static void bitmap_tree_init(struct bitmap_tree *tree, size_t i,
+			     uint64_t nr_unset, uint64_t nr_this_level)
+{
+	struct bitmap_tree_node *n = tree->tree + i;
+	unsigned j = 0;
+
+	if (i < tree->first_leaf) {
+		i = TREE_ARY * i + 1;
+
+		nr_this_level /= TREE_ARY;
+
+		while (nr_unset) {
+			assert(j < TREE_ARY);
+
+			n->nr_unset[j] = nr_unset < nr_this_level
+				? nr_unset : nr_this_level;
+
+			bitmap_tree_init(tree, i + j, n->nr_unset[j],
+					 nr_this_level);
+
+			nr_unset -= n->nr_unset[j];
+			j++;
+		}
+	} else {
+		memset(n, 0, sizeof(*n));
+	}
+}
+
+struct bitmap_tree *bitmap_tree_alloc(uint64_t size)
+{
+	struct bitmap_tree *tree;
+	size_t nr_leaf_nodes = (size + BITS_PER_LEAF - 1) / BITS_PER_LEAF;
+	size_t nr_nodes = 0;
+	size_t nr_cur_level = 1;
+
+	while (1) {
+		if (nr_cur_level >= nr_leaf_nodes)
+			break;
+
+		nr_nodes += nr_cur_level;
+		nr_cur_level *= TREE_ARY;
+	}
+
+	nr_nodes += nr_leaf_nodes;
+
+	tree = malloc(sizeof(struct bitmap_tree) +
+		      sizeof(struct bitmap_tree_node) * nr_nodes);
+
+	tree->size = size;
+	tree->nr_unset = size;
+	tree->first_leaf = nr_nodes - nr_leaf_nodes;
+
+	bitmap_tree_init(tree, 0, tree->nr_unset,
+			 nr_cur_level * BITS_PER_LEAF);
+
+	return tree;
+}
+
+static uint64_t set_nth_unset_bit_recurse(struct bitmap_tree *tree, size_t i,
+					  uint64_t bit_to_set)
+{
+	struct bitmap_tree_node *n = tree->tree + i;
+	unsigned j = 0;
+
+	if (i < tree->first_leaf) {
+		i = TREE_ARY * i + 1;
+
+		while (bit_to_set > n->nr_unset[j]) {
+			bit_to_set -= n->nr_unset[j];
+			j++;
+			assert(j < TREE_ARY);
+		}
+
+		n->nr_unset[j]--;
+
+		return set_nth_unset_bit_recurse(tree, i + j, bit_to_set);
+	} else {
+		uint64_t ret = (i - tree->first_leaf) * BITS_PER_LEAF;
+		unsigned bit = 0;
+
+		while (bit_to_set > 64 - popcount_64(n->nr_unset[j])) {
+			bit_to_set -= 64 - popcount_64(n->nr_unset[j]);
+
+			ret += 64;
+			j++;
+			assert(j < TREE_ARY);
+		}
+
+		do {
+			while (n->nr_unset[j] & (1ULL << bit)) {
+				ret++;
+				bit++;
+				assert(bit < 64);
+			}
+		} while (bit_to_set--);
+
+		n->nr_unset[j] &= ~(1ULL << bit);
+
+		return ret;
+	}
+}
+
+uint64_t set_nth_unset_bit(struct bitmap_tree *tree, uint64_t bit)
+{
+	bit %= tree->nr_unset;
+	tree->nr_unset--;
+
+	return set_nth_unset_bit_recurse(tree, 0, bit);
+}
+
+int main()
+{
+	struct bitmap_tree *tree = bitmap_tree_alloc(123320987);
+
+	while (tree->nr_unset) {
+		uint64_t bit = set_nth_unset_bit(tree, lrand48());
+
+		if (!(tree->nr_unset % 10000))
+			printf("set bit %lu\n", bit);
+	}
+
+	return 0;
+}