#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;
}