1 files changed, 21 insertions, 19 deletions

diff --git a/src/bumpptr_alloc.c b/src/bumpptr_alloc.c
index 3985904..6cb82df 100644
--- a/src/bumpptr_alloc.c
+++ b/src/bumpptr_alloc.c
@@ -1,3 +1,4 @@
+#include <assert.h>
 #include <errno.h>
 #include <stddef.h> /* NULL, size_t */
 #include <stdint.h> /* uintptr_t */
@@ -19,11 +20,12 @@ extern "C" {
 #endif
 
 __thread void* mem_start = NULL;
-__thread void* mem_end = NULL;
-
+__thread uintptr_t mem_end = 0;
 __thread uintptr_t ptr = 0;
 
-void* malloc(size_t size) {
+inline void* bump_up(size_t size, size_t align) {
+	assert(align % 2 == 0);
+
 	if (unlikely(mem_start == NULL)) {
 		mem_start = mmap(NULL, MEMSIZE, PROT_READ|PROT_WRITE, MAP_PRIVATE|MAP_ANONYMOUS, -1, 0);
 		if(mem_start == MAP_FAILED) {
@@ -31,17 +33,23 @@ void* malloc(size_t size) {
 			return NULL;
 		}
 		ptr = (uintptr_t)mem_start;
+		mem_end = ptr + MEMSIZE;
 	}
 
 	// align ptr;
-	if (ptr % MIN_ALIGNMENT != 0) {
-		size_t mask = MIN_ALIGNMENT -1;
-		ptr = (ptr + mask) & ~mask;
+	uintptr_t aligned = (ptr + align - 1) & ~(align - 1);
+
+	uintptr_t new_ptr = aligned + size;
+	if (new_ptr > mem_end)
+		return NULL;
+	else {
+		ptr = new_ptr;
+		return (void*)aligned;
 	}
+}
 
-	void* ret = (void*)ptr;
-	ptr += size;
-	return ret;
+void* malloc(size_t size) {
+	return bump_up(size, MIN_ALIGNMENT);
 }
 
 void free(__attribute__ ((unused)) void* ptr) {
@@ -51,20 +59,14 @@ void* realloc(void* ptr, size_t size) {
 	if(ptr == NULL)
 		return malloc(size);
 
-	if(size == 0)
-		return NULL;
-
-	void* new_ptr = malloc(size);
+	void* new_ptr = bump_up(size, MIN_ALIGNMENT);
 	// this may copies to much
 	memcpy(new_ptr, ptr, size);
 	return new_ptr;
 }
 
 void* memalign(size_t alignment, size_t size) {
-	// align ptr to alignment and malloc
-	size_t mask = alignment - 1;
-	ptr = (ptr + mask) & ~mask;
-	return malloc(size);
+	return bump_up(size, alignment);
 }
 
 int posix_memalign(void **memptr, size_t alignment, size_t size)
@@ -89,7 +91,7 @@ int posix_memalign(void **memptr, size_t alignment, size_t size)
 		return 0;
 	}
 
-	out = memalign(alignment, size);
+	out = bump_up(size, alignment);
 	if(out == NULL) {
 		return 12;
 	}
@@ -107,7 +109,7 @@ void* calloc(size_t nmemb, size_t size)
 		return NULL;
 	}
 	
-	out = malloc(fullsize);
+	out = bump_up(fullsize, MIN_ALIGNMENT);
 	if(out == NULL) {
 		return NULL;
 	}