/*
 * A batch is a cacheline-friendly way of storing a batch of pointers.
 * Each level of the batch contains 15 pointers plus a pointer to the
 * next batch_list.  Except the last list of batches which contains 16
 * user pointers.
 */

#include <linux/kernel.h>
#include <linux/slab.h>

/*
 * count encodes the number of elements in the batch in an unusual
 * manner.  If there are one or no elements in the batch, count is 0.
 * Once we alloc a subbatch, the element at head is placed in slot 15,
 * and count has value 31.  We then fill in slots 14-0 at which point
 * count has the value 16.  When we store the 17th pointer in the batch,
 * count soars to 46, then counts down to 32.  batch_count() decodes this
 * for you if it's important to know exactly how many elements are in the
 * batch.  If you only need to know approximately, testing it for
 * batch->count >= 256, for example, will be close enough.
 */
struct batch {
	unsigned int count;
	struct batch_sub *head;
};

struct batch_sub {
	void *slots[16];
};

unsigned int batch_size(struct batch *batch)
{
	if (!batch->count)
		return batch->head ? 1 : 0;
	return (batch->count >> 4) * 15 + (16 - batch->count);
}

bool batch_add(struct batch *batch, void *data, gfp_t gfp)
{
	struct batch_sub *sub;

	if (batch->count & 0xf) {
		batch->count--;
		batch->head->slots[batch->count & 0xf] = data;
		return true;
	}

	if (!batch->head) {
		batch->head = data;
		return true;
	}

	sub = kmalloc(sizeof(*sub), gfp);
	if (!sub)
		return false;

	sub->slots[15] = batch->head;
	sub->slots[14] = data;
	batch->head = sub;
	batch->count += 0x1e;
	return true;
}

void *batch_remove(struct batch *batch)
{
	void *data;
	if (batch->count > 0xf) {
		data = batch->head->slots[batch->count & 0xf];
		batch->count++;
		if ((batch->count & 0xf) == 0xf) {
			void *head = batch->head->slots[15];
			kfree(batch->head);
			batch->head = head;
			batch->count -= 0x1f;
		}
	} else {
		data = batch->head;
		if (data)
			batch->head = NULL;
	}
	return data;
}

#ifndef __KERNEL__
#include <assert.h>
#include <stdio.h>
int main(void)
{
	struct batch batch = { 0 };
	void *items[50];
	int i;

	assert(batch_remove(&batch) == NULL);

	for (i = 0; i < 50; i++) {
		items[i] = kmalloc(16, GFP_KERNEL);
		assert(items[i]);
	}

	for (i = 0; i < 50; i++) {
		assert(batch_add(&batch, items[i], GFP_KERNEL));
	}

	for (i = 49; i >= 0; i--) {
		assert(batch_remove(&batch) == items[i]);
	}

	assert(batch_remove(&batch) == NULL);
}
#endif