A compat memory alloc for lua

compat_alloc.c

// 	gcc -g -Wall -o alloc alloc.c

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

#if defined(_MSC_VER) || defined(__MINGW32__) || defined(__MINGW64__)

#include <windows.h>

static void *
memory_map(size_t total) {
	return VirtualAlloc(NULL, total, MEM_RESERVE, PAGE_NOACCESS);
}

static void
memory_unmap(void * ptr, size_t total) {
	VirtualFree(ptr, total, MEM_RELEASE);
}

static void *
memory_alloc(void *ptr, size_t sz) {
	return VirtualAlloc(ptr, sz, MEM_COMMIT, PAGE_READWRITE);
}

static void
memory_free(void *ptr, size_t sz) {
	VirtualAlloc(ptr, sz, MEM_RESET, PAGE_NOACCESS);
}

#else

#include <sys/mman.h>

static void *
memory_map(size_t total) {
	void *ptr = mmap(NULL, total, PROT_NONE, MAP_ANONYMOUS | MAP_PRIVATE, -1, 0);
	if (ptr == MAP_FAILED) {
		return NULL;
	}
	return ptr;
}

static void
memory_unmap(void * ptr, size_t sz) {
	munmap(ptr, sz);
}

static void *
memory_alloc(void *ptr, size_t sz) {
	void *p = mmap(ptr, sz, PROT_READ | PROT_WRITE, MAP_ANONYMOUS | MAP_PRIVATE | MAP_FIXED, -1, 0);
	if (p == MAP_FAILED) {
		return NULL;
	}
	return p;
}

static void
memory_free(void *ptr, size_t sz) {
	mmap(ptr, sz, PROT_NONE, MAP_ANONYMOUS | MAP_PRIVATE | MAP_FIXED, -1, 0);
}

#endif

typedef unsigned int hoffset_t;
typedef unsigned int hsize_t;

// at least sizeof(hoffset_t) + sizeof(hsize_t)
#define ALIGN 8
#define NULL_OFFSET 0

struct heap {
	hsize_t total_size;
	hsize_t chunk_size;
	hsize_t size;
	hoffset_t freelist;
	hoffset_t freehead;
	hsize_t used_size;
	int alive_n;
	int free_n;
};

struct block {
	hsize_t sz;
	hoffset_t next;
};

struct free_block {
	hsize_t sz;
	hoffset_t offset;
};

static inline struct block *
heap_offset(struct heap *h, hoffset_t off) {
	return (struct block *)((char *)h + off * ALIGN);
}

static inline struct block *
heap_freelist(struct heap *h) {
	return heap_offset(h, h->freelist);
}

static inline void *
heap_alloc_(struct heap *h, hsize_t sz) {
	if (h->freelist == NULL_OFFSET)
		return NULL;
	struct block * b = heap_freelist(h);
	hsize_t hsize = b->sz;
	if (hsize <= sz) {
		// try next free block
		h->freelist = b->next;
		if (h->freelist == NULL_OFFSET)
			return NULL;
		b = heap_freelist(h);
		hsize = b->sz;
		if (hsize <= sz) {
			return NULL;
		}
	}
	++h->alive_n;
	h->used_size += sz;
	b->sz = hsize - sz;	// split block
	return heap_offset(h, h->freelist + b->sz);	// return last part
}

static inline void
heap_free_(struct heap *h, void *ptr, hsize_t sz) {
	assert(ptr >= (void *)h && ptr <= (void *)heap_offset(h, h->size - sz));
	hoffset_t off = ((char *)ptr - (char *)h) / ALIGN;
	struct block * b = ptr;
	--h->alive_n;
	++h->free_n;
	h->used_size -= sz;
	b->sz = sz;
	// add to head of free list
	b->next = h->freehead;
	h->freehead = off;
}

static void
heap_shrink(struct heap *h, void *ptr, size_t osize, size_t nsize) {
	hsize_t os = (osize + ALIGN -1) / ALIGN;
	hsize_t ns = (nsize + ALIGN -1) / ALIGN;
	if (ns >= os)
		return;
	heap_free_(h, (char *)ptr + ns * ALIGN, os-ns);
}

static int
comp_offset(const void *a, const void *b) {
	const struct free_block *aa = a;
	const struct free_block *bb = b;
	return (signed int)(aa->offset - bb->offset);
}

static int
comp_size(const void *a, const void *b) {
	const struct free_block *aa = a;
	const struct free_block *bb = b;
	return (signed int)(bb->sz - aa->sz);
}

static void
heap_arrange(struct heap *h) {
	int i;
	int n = h->free_n;
	if (n == 0)
		return;

	size_t tmp_size = n * sizeof(struct free_block);
	if (tmp_size > (h->total_size - h->size) * ALIGN)
		return;	// not enough tmp size

	struct free_block *dead = (struct free_block *)memory_alloc(heap_offset(h, h->size), tmp_size);
	struct block * b = heap_offset(h, h->freehead);
	for (i=0;i<n;i++) {
		assert(b);
		dead[i].sz = b->sz;
		dead[i].offset = ((char *)b - (char *)h) / ALIGN;	// offset
		if (b->next == NULL_OFFSET) {
			b = NULL;
		} else {
			b = heap_offset(h, b->next);
		}
	}

	assert(b == NULL);

	qsort(dead, n, sizeof(struct free_block), comp_offset);

	// combine continues free block
	int p=0;
	hoffset_t next_offset = dead[p].offset + dead[p].sz;
	for (i=1;i<n;i++) {
		if (next_offset == dead[i].offset) {
			// combine
			dead[p].sz += dead[i].sz;
			next_offset += dead[i].sz;
		} else {
			++p;
			dead[p] = dead[i];
			next_offset = dead[p].offset + dead[p].sz;
		}
	}

	// sort by size
	qsort(dead, p+1, sizeof(struct free_block), comp_size);

	for (i=0;i<p;i++) {
		struct block *b = heap_offset(h, dead[i].offset);
		b->sz = dead[i].sz;
		b->next = dead[i+1].offset;
	}
	b = heap_offset(h, dead[p].offset);
	b->sz = dead[p].sz;
	b->next = NULL_OFFSET;
	
	h->freehead = h->freelist = dead[0].offset;
	h->free_n = p + 1;

	memory_free(dead, tmp_size);
}

void
heap_dump_info(struct heap *h) {
	printf("HEAP size = %fM, use = %fM, alive block = %d, freed block = %d\n", 
		(float)h->size / 1024 * ALIGN / 1024,
		(float)h->used_size * ALIGN / 1024.0f / 1024.0f,
		h->alive_n,
		h->free_n);
	struct block *b = heap_freelist(h);
	printf("Free block = %fK\n", (float)b->sz * ALIGN / 1024.0f);
	int n = 0;
	while (b->sz >= 1024 * 1024 / ALIGN) {	// big than 1M
		++n;
		if (b->next == NULL_OFFSET) {
			b = NULL;
			break;
		}
		b = heap_offset(h, b->next);
	}
	printf("Block >1M %d\n", n);
	n = 0;
	while (b && b->sz >= 1024) {	// big than 8K
		++n;
		if (b->next == NULL_OFFSET) {
			break;
		}
		b = heap_offset(h, b->next);
	}
	printf("Block >8K %d\n", n);
}

void
heap_dump_freelist(struct heap *h) {
	int i;
	struct block * b = heap_offset(h, h->freehead);
	for (i=0;i<h->free_n;i++) {
		if (b == NULL) {
			printf("(Corrupt block)");
			break;
		}
		printf("(%d %d %d) ", 
			(int)(((char *)b - (char *)h) / ALIGN),  // offset
			(int)(b->sz),	// size
			(int)(b->next));
		if (b->next == NULL_OFFSET) {
			b = NULL;
		} else {
			b = heap_offset(h, b->next);
		}
	}
	printf("\n");
}

struct heap *
heap_new(size_t total_size, size_t chunk_size) {
	assert(total_size / ALIGN <= 0xffffffff && chunk_size > sizeof(struct heap));	// max 32G
	total_size = total_size / ALIGN * ALIGN;
	void * mem = memory_map(total_size);
	if (mem == NULL)
		return NULL;
	hsize_t csz = (hsize_t)(chunk_size / ALIGN);
	void * ptr = memory_alloc(mem, csz * ALIGN);
	if (ptr == NULL) {
		memory_unmap(mem, total_size);
		return NULL;
	}
	hsize_t meta_size = (sizeof(struct heap) + ALIGN -1) / ALIGN;
	struct heap *h = mem;
	h->total_size = (hsize_t)total_size / ALIGN;
	h->chunk_size = csz;
	h->size = csz;
	h->freelist = meta_size;
	h->freehead = meta_size;
	h->used_size = 0;
	h->alive_n = 1;
	h->free_n = 1;
	struct block *b = heap_freelist(h);
	b->sz = h->size - meta_size;
	b->next = NULL_OFFSET;
	return h;
}

void
heap_delete(struct heap *h) {
	memory_unmap((void *)h, h->total_size * ALIGN);
}

void *
heap_alloc(struct heap *h, size_t sz) {
	hsize_t s = (sz + ALIGN -1) / ALIGN;
	if (s <= h->size - h->used_size) {
		void * ret = heap_alloc_(h, s);
		if (ret) {
			return ret;
		}
		heap_arrange(h);
		ret = heap_alloc_(h, s);
		if (ret) {
			return ret;
		}
	}
	// expand
	hsize_t new_size = h->size + h->chunk_size;
	if (new_size > h->total_size)
		return NULL;
	void * ptr = memory_alloc(heap_offset(h, h->size), h->chunk_size * ALIGN);
	if (ptr == NULL)
		return NULL;	// never failed
	++h->free_n;
	struct block *b = (void *)ptr;
	b->sz = h->chunk_size;
	b->next = h->freehead;
	h->freelist = h->freehead = h->size;
	h->size = new_size;
	return heap_alloc_(h, s);
}

void
heap_free(struct heap *h, void *ptr, size_t sz) {
	if (ptr == NULL)
		return;
	hsize_t s = (sz + ALIGN -1) / ALIGN;
	heap_free_(h, ptr, s);
}

void *
heap_lua_alloc(void *ud, void *ptr, size_t osize, size_t nsize) {
	struct heap * h = ud;
	if (nsize == 0) {
		heap_free(h, ptr, osize);
		return NULL;
	}
	if (ptr == NULL) {
		// new block
		return heap_alloc(h, nsize);
	}
	if (nsize <= osize) {
		heap_shrink(h, ptr, osize, nsize);
		return ptr;
	}
	void *newptr = heap_alloc(h, nsize);
	if (newptr == NULL) {
		return NULL;
	}
	memcpy(newptr, ptr, osize);
	heap_free(h, ptr, osize);
	return newptr;
}

int
main() {
	struct heap *h = heap_new(1024 * 1024 * 1024, 1024 * ALIGN);
	void *p1 = heap_alloc(h, 3);
	void *p2 = heap_alloc(h, 13);
	void *p3 = heap_alloc(h, 23);

	printf("alloc %p %p\n", p1, p2);
	heap_free(h, p1, 3);
	heap_dump_info(h);
	heap_free(h, p2, 13);
	heap_dump_info(h);
	heap_arrange(h);
	heap_dump_info(h);
	heap_free(h, p3, 23);
	heap_dump_info(h);
	heap_arrange(h);
	heap_dump_info(h);

	int i;
	for (i=0;i<10;i++) {
		heap_alloc(h, 1024);
	}

	heap_dump_info(h);

	heap_delete(h);
	return 0;
}