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

typedef uint32_t unset_t;

#define BITS_PER_WORD	(sizeof(unset_t) * 8)
#define bits_unset(x)	(BITS_PER_WORD - __builtin_popcount(x))

#define TREE_ARY	16
#define BITS_PER_LEAF	(TREE_ARY * BITS_PER_WORD)

struct bitmap_tree {
	struct {
		uint64_t	size;
		uint64_t	nr_unset;
		size_t		first_leaf;
	} __attribute__((aligned(64)));

	struct bitmap_tree_node {
		unset_t nr_unset[TREE_ARY];
	} tree[];
};

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) {
			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, bit = 0;

	if (i < tree->first_leaf) {
		i = TREE_ARY * i + 1;
		__builtin_prefetch(tree->tree + i);

		while (bit_to_set > n->nr_unset[j]) {
			bit_to_set -= n->nr_unset[j];
			j++;
		}

		n->nr_unset[j]--;

		return set_nth_unset_bit_recurse(tree, i + j, bit_to_set);
	} else {
		while (bit_to_set > bits_unset(n->nr_unset[j])) {
			bit_to_set -= bits_unset(n->nr_unset[j]);
			j++;
		}

		do {
			while (n->nr_unset[j] & (1ULL << bit))
				bit++;
		} while (bit_to_set--);

		n->nr_unset[j] &= ~(1ULL << bit);

		return ((i - tree->first_leaf) * TREE_ARY + j) * BITS_PER_WORD + bit;
	}
}

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)
		set_nth_unset_bit(tree, lrand48());

	return 0;
}