Dev (#1)

* Fixed size extend check in realloc

* Fixed size usage overlapping BlockHeaders

* Fixed usage of sizing and sbrk invocations

* Fixed condition for realloc resizing to contiguously available free mem

1 files changed, 30 insertions, 39 deletions

diff --git a/libpsn00b/libc/malloc.c b/libpsn00b/libc/malloc.c
index ed12e0b..4477834 100644
--- a/libpsn00b/libc/malloc.c
+++ b/libpsn00b/libc/malloc.c
@@ -17,11 +17,12 @@
 #include <stdint.h>
 #include <stdlib.h>
 
+#define ALIGN_SIZE 8
 #define _align(x, n) (((x) + ((n) - 1)) & ~((n) - 1))
 
 /* Private types */
 
-typedef struct _BlockHeader {
+typedef struct __attribute__((aligned(ALIGN_SIZE))) _BlockHeader {
 	struct _BlockHeader	*prev, *next;
 	void				*ptr;
 	size_t				size;
@@ -52,7 +53,7 @@ __attribute__((weak)) void InitHeap(void *addr, size_t size) {
 
 __attribute__((weak)) void *sbrk(ptrdiff_t incr) {
 	void *old_end = _heap_end;
-	void *new_end = (void *) _align((uintptr_t) old_end + incr, 8);
+	void *new_end = (void *) _align((uintptr_t) old_end + incr, ALIGN_SIZE);
 
 	if (new_end > _heap_limit)
 		return 0;
@@ -81,7 +82,6 @@ __attribute__((weak)) void GetHeapUsage(HeapUsage *usage) {
 
 static BlockHeader *_find_fit(BlockHeader *head, size_t size) {
 	BlockHeader *prev = head;
-
 	for (; prev; prev = prev->next) {
 		if (prev->next) {
 			uintptr_t next_bot = (uintptr_t) prev->next;
@@ -91,7 +91,6 @@ static BlockHeader *_find_fit(BlockHeader *head, size_t size) {
 				return prev;
 		}
 	}
-
 	return prev;
 }
 
@@ -99,28 +98,28 @@ __attribute__((weak)) void *malloc(size_t size) {
 	if (!size)
 		return 0;
 
-	size_t _size = _align(size + sizeof(BlockHeader), 8);
+	size_t _size    = _align(size + sizeof(BlockHeader), ALIGN_SIZE);
+	size_t _size_nh = _size - sizeof(BlockHeader);
 
 	// Nothing's initialized yet? Let's just initialize the bottom of our heap,
 	// flag it as allocated.
 	if (!_alloc_head) {
 		//if (!_alloc_start)
 			//_alloc_start = sbrk(0);
-
 		BlockHeader *new = (BlockHeader *) sbrk(_size);
 		if (!new)
 			return 0;
 
 		void *ptr = (void *) &new[1];
 		new->ptr  = ptr;
-		new->size = _size - sizeof(BlockHeader);
+		new->size = _size_nh;
 		new->prev = 0;
 		new->next = 0;
 
 		_alloc_head = new;
 		_alloc_tail = new;
 
-		TrackHeapUsage(size);
+		TrackHeapUsage(_size);
 		return ptr;
 	}
 
@@ -132,14 +131,14 @@ __attribute__((weak)) void *malloc(size_t size) {
 
 		void *ptr = (void *) &new[1];
 		new->ptr  = ptr;
-		new->size = _size - sizeof(BlockHeader);
+		new->size = _size_nh;
 		new->prev = 0;
 		new->next = _alloc_head;
 
 		_alloc_head->prev = new;
 		_alloc_head       = new;
 
-		TrackHeapUsage(size);
+		TrackHeapUsage(_size);
 		return ptr;
 	}
 
@@ -148,16 +147,16 @@ __attribute__((weak)) void *malloc(size_t size) {
 	if (prev) {
 		BlockHeader *new = (BlockHeader *) ((uintptr_t) prev->ptr + prev->size);
 
-		void *ptr = (void *)((uintptr_t) new + sizeof(BlockHeader));
+		void *ptr = (void *) &new[1];
 		new->ptr  = ptr;
-		new->size = _size - sizeof(BlockHeader);
+		new->size = _size_nh;
 		new->prev = prev;
 		new->next = prev->next;
 
 		(new->next)->prev = new;
 		prev->next        = new;
 
-		TrackHeapUsage(size);
+		TrackHeapUsage(_size);
 		return ptr;
 	}
 
@@ -168,14 +167,14 @@ __attribute__((weak)) void *malloc(size_t size) {
 
 	void *ptr = (void *) &new[1];
 	new->ptr  = ptr;
-	new->size = _size - sizeof(BlockHeader);
+	new->size = _size_nh;
 	new->prev = _alloc_tail;
 	new->next = 0;
 
 	_alloc_tail->next = new;
 	_alloc_tail       = new;
 
-	TrackHeapUsage(size);
+	TrackHeapUsage(_size);
 	return ptr;
 }
 
@@ -191,35 +190,35 @@ __attribute__((weak)) void *realloc(void *ptr, size_t size) {
 	if (!ptr)
 		return malloc(size);
 
-	size_t      _size = _align(size + sizeof(BlockHeader), 8);
+	size_t _size      = _align(size + sizeof(BlockHeader), ALIGN_SIZE);
+	size_t _size_nh   = _size - sizeof(BlockHeader);
 	BlockHeader *prev = (BlockHeader *) ((uintptr_t) ptr - sizeof(BlockHeader));
 
 	// New memory block shorter?
-	if (prev->size >= _size) {
-		TrackHeapUsage(size - prev->size);
-		prev->size = _size - sizeof(BlockHeader);
+	if (prev->size >= _size_nh) {
+		TrackHeapUsage(_size_nh - prev->size);
+		prev->size = _size_nh;
 
 		if (!prev->next)
-			sbrk((ptr - sbrk(0)) + _size);
+			sbrk((ptr - sbrk(0)) + _size_nh);
 
 		return ptr;
 	}
 
 	// New memory block larger; is it the last one?
 	if (!prev->next) {
-		void *new = sbrk(_size - prev->size);
+		void *new = sbrk(_size_nh - prev->size);
 		if (!new)
 			return 0;
-
-		TrackHeapUsage(size - prev->size);
-		prev->size = _size - sizeof(BlockHeader);
+		TrackHeapUsage(_size_nh - prev->size);
+		prev->size = _size_nh;
 		return ptr;
 	}
 
 	// Do we have free memory after it?
-	if (((prev->next)->ptr - ptr) > _size) {
-		TrackHeapUsage(size - prev->size);
-		prev->size = _size - sizeof(BlockHeader);
+	if ((uintptr_t) prev->next - (uintptr_t) ptr >= _size_nh) {
+		TrackHeapUsage(_size_nh - prev->size);
+		prev->size = _size_nh;
 		return ptr;
 	}
 
@@ -236,7 +235,6 @@ __attribute__((weak)) void *realloc(void *ptr, size_t size) {
 __attribute__((weak)) void free(void *ptr) {
 	if (!ptr || !_alloc_head)
 		return;
-
 	// First block; bumping head ahead.
 	if (ptr == _alloc_head->ptr) {
 		size_t size = _alloc_head->size;
@@ -250,30 +248,23 @@ __attribute__((weak)) void free(void *ptr) {
 			sbrk(-size);
 		}
 
-		TrackHeapUsage(-(_alloc_head->size));
+		TrackHeapUsage(-size);
 		return;
 	}
 
 	// Finding the proper block
-	BlockHeader *cur = _alloc_head;
-
-	for (cur = _alloc_head; ptr != cur->ptr; cur = cur->next) {
-		if (!cur->next)
-			return;
-	}
-
+	BlockHeader* cur = (BlockHeader*) (ptr - sizeof(BlockHeader));
 	if (cur->next) {
 		// In the middle, just unlink it
 		(cur->next)->prev = cur->prev;
 	} else {
 		// At the end, shrink heap
-		void  *top  = sbrk(0);
+		void   *top = sbrk(0);
 		size_t size = (top - (cur->prev)->ptr) - (cur->prev)->size;
 		_alloc_tail = cur->prev;
 
 		sbrk(-size);
 	}
-
-	TrackHeapUsage(-(cur->size));
+	TrackHeapUsage(-(cur->size + sizeof(BlockHeader)));
 	(cur->prev)->next = cur->next;
 }