Fixed runtime errors.

include/CHM.h

@@ -6,12 +6,6 @@
 #include <stdlib.h>
 #include <sys/syscall.h>
 
-#define ALIGN4(s) (((((s)-1)>>2)<<2)+4)
-#define ALIGN16(s) (((s) + 15) & ~0x0F)
-#define ALIGNMENT 4
-#define BLOCK_SIZE sizeof(struct Block)
-#define MINIMUM_BLOCK_SIZE (sizeof(struct Block) + 16)
-
 enum MemMode : unsigned char
 {
 	MEM_MODE_SPEED, // Memory will be allocated more quickly but loses memory efficiency.
@@ -34,9 +28,10 @@ void *realloc(void *ptr, size_t newSize);
 /// @param [in] ptr The memory to flag as free.
 void free(void *ptr);
 
+size_t get_header_size();
+
 /// Will retrieve the current real byte size of the heap.
 /// @returns The size in bytes.
 size_t get_heap_size();
 
-/// Will check memory blocks on program exit if there's dangling memory.
+void cleanup_memory();
-void check_memory();

src/CHM.c

@@ -6,6 +6,11 @@
 #include <stdlib.h>
 #include <sys/syscall.h>
 
+#define ALIGN4(s) (((((s)-1)>>2)<<2)+4)
+#define ALIGN16(s) (((s) + 15) & ~0x0F)
+#define ALIGNMENT 4
+#define BLOCK_SIZE sizeof(struct Block)
+
 /// The memory block's header.
 struct Block
 {
@@ -25,7 +30,7 @@ void set_memory_mode(const enum MemMode new_mode)
 {
 	if (new_mode > MEM_MODE_BALANCED)
 	{
-		printf("The given memory mode, \"%i\", is invalid.", new_mode);
+		fprintf(stderr, "The given memory mode, \"%i\", is invalid.", new_mode);
 		return;
 	}
 
@@ -62,61 +67,6 @@ struct Block *extend_heap(const size_t s)
 	return b;
 }
 
-struct Block *extend_block(struct Block *in, const size_t newSize)
-{
-	if (!in->next)
-	{
-		void *addr = (struct Block *)syscall(SYS_brk, NULL);
-
-		if ((void *)syscall(SYS_brk, (char *)addr + (newSize - in->size)) == (void *)-1)
-			return NULL;
-
-		in->size = newSize;
-
-		return in;
-	}
-
-	size_t totalSize = in->size;
-	struct Block *start = in->prev;
-	struct Block *end = in->next;
-
-	while (start && start->free && totalSize < newSize)
-	{
-		totalSize += BLOCK_SIZE + in->prev->size;
-		start = in->prev;
-	}
-
-	while (end && end->free && totalSize < newSize)
-	{
-		totalSize += BLOCK_SIZE + end->size;
-		end = end->next;
-	}
-
-	if (totalSize < newSize || (mode == MEM_MODE_EFFICIENCY && totalSize > newSize))
-		return NULL;
-
-	if (!start)
-		start = in;
-
-	if (!end)
-		end = in;
-
-	struct Block *b = start->next;
-	while (b)
-	{
-		start->size += BLOCK_SIZE + b->size;
-		b = b->next;
-	}
-
-	start->next = end->next;
-	start->free = 0;
-
-	if (!start->next)
-		last = start;
-
-	return start;
-}
-
 /// Will defragment adjacent memory blocks of the given memory block.
 /// @param [in] in The memory block to defragment for.
 void defragment_adjacent_blocks(struct Block *in)
@@ -151,6 +101,76 @@ void fragment_block(struct Block *in, size_t s)
 	in->free = 0;
 }
 
+struct Block *extend_block(struct Block *in, const size_t newSize)
+{
+	if (!in->next)
+	{
+		void *addr = (struct Block *)syscall(SYS_brk, NULL);
+
+		if ((void *)syscall(SYS_brk, (char *)addr + (newSize - in->size)) == (void *)-1)
+			return NULL;
+
+		in->size = newSize;
+
+		return in;
+	}
+
+	size_t totalSize = in->size;
+	struct Block *start = in->prev;
+	struct Block *end = in->next;
+
+	while (start && start->free && totalSize < newSize)
+	{
+		totalSize += BLOCK_SIZE + start->size;
+		if (totalSize >= newSize)
+			break;
+
+		start = in->prev;
+	}
+
+	while (end && end->free && totalSize < newSize)
+	{
+		totalSize += BLOCK_SIZE + end->size;
+		if (totalSize >= newSize)
+			break;
+
+		end = end->next;
+	}
+
+	if (totalSize < newSize || (mode == MEM_MODE_EFFICIENCY && totalSize > newSize))
+		return NULL;
+
+	if (!start)
+		start = in;
+
+	if (!end)
+		end = in;
+
+	struct Block *b = start->next;
+	while (b)
+	{
+		start->size += BLOCK_SIZE + b->size;
+
+		if (b == end)
+			break;
+
+		b = b->next;
+	}
+
+	start->next = end->next;
+	start->free = 0;
+	start->next->prev = start;
+
+	if (!start->next)
+		last = start;
+
+	const size_t remainder = start->size - (BLOCK_SIZE + newSize);
+	if (remainder && !(remainder % ALIGNMENT))
+		fragment_block(start, newSize);
+
+	return start;
+}
+
 /// Finds the first block that will fit the given size.
 /// @param [in] s The memory aligned size to look for.
 /// @returns The matching available memory block.
@@ -194,7 +214,7 @@ struct Block *find_balanced(const size_t s)
 			break;
 
 		const size_t remainder = current->size - totalSize;
-		if (!(remainder % ALIGNMENT))
+		if (remainder && !(remainder % ALIGNMENT))
 		{
 			fragment_block(current, s);
 			break;
@@ -231,9 +251,10 @@ struct Block *find_best_fit(const size_t s)
 			break;
 		}
 
+		const size_t remainder = current->size - totalSize;
 		if (best)
 		{
-			if (current->size < best->size && !((current->size - totalSize) % ALIGNMENT))
+			if (current->size < best->size && remainder && !(remainder % ALIGNMENT))
 			{
 				fragment = 1;
 				best = current;
@@ -241,7 +262,7 @@ struct Block *find_best_fit(const size_t s)
 		}
 		else
 		{
-			if (!((current->size - totalSize) % ALIGNMENT))
+			if (remainder && !(remainder % ALIGNMENT))
 			{
 				fragment = 1;
 				best = current;
@@ -311,21 +332,29 @@ void *realloc(void *ptr, size_t newSize)
 
 	newSize = ALIGN16(newSize);
 
-	struct Block *b = ((struct Block *)ptr) - 1;
+	struct Block *b = (struct Block *)ptr - 1;
 	if (b->free)
 		return NULL;
 
-	size_t totalSize = BLOCK_SIZE + newSize;
+	if (b->size == newSize)
-	size_t remainder = b->size - totalSize;
+		return b + 1;
 
-	if (b->size < newSize && !extend_block(b, newSize))
+	const size_t totalSize = BLOCK_SIZE + newSize;
-		return NULL;
+	const size_t remainder = b->size - totalSize;
-	else if (b->size > newSize && !(remainder % ALIGNMENT))
+	if (b->size < newSize) // Expand memory block if possible.
+	{
+		if (!extend_block(b, newSize))
+			return NULL;
+
+		return b + 1;
+	}
+	else if (b->size > newSize && remainder && !(remainder % ALIGNMENT)) // Shrink memory block if possible.
+	{
 		fragment_block(b, newSize);
-	else
+		return b + 1;
-		return ptr;
+	}
 
-	return b + 1;
+	return NULL;
 }
 
 void free(void *ptr)
@@ -359,27 +388,39 @@ void free(void *ptr)
 	}
 }
 
+size_t get_header_size()
+{
+	return BLOCK_SIZE;
+}
+
 size_t get_heap_size()
 {
 	if (!first || !last)
 		return 0;
 
-	return (char*)first - (char*)last - last->size;
+	return (char*)last + last->size - (char*)first;
 }
 
-void check_memory()
+void cleanup_memory()
 {
 	if (!first || !last)
 		return;
 
+	const struct Block *l = last;
 	struct Block *current = first;
 	while (current)
 	{
-		if (current->free)
+		if (!current->free)
-			printf("Memory at the address %p has been marked free, but still persists.", current + 1);
+			fprintf(stderr, "Memory at the address %p has not been marked free. Possible memory leak?\n", current + 1);
-		else
+
-			printf("Memory at the address %p has not been marked free. Possible memory leak?", current + 1);
+		if (current == l)
+			break;
 
 		current = current->next;
 	}
+
+	syscall(SYS_brk, first);
+
+	first = NULL;
+	last = NULL;
 }

src/main.c

@@ -1,23 +1,25 @@
 #include "CHM.h"
 
+#include <unistd.h>
+
 int main()
 {
+	set_memory_mode(MEM_MODE_BALANCED);
+
 	int *a = (int *)malloc(sizeof(int));
+	*a = 1;
 	int *b = (int *)malloc(sizeof(int));
+	*b = 2;
+	int *c = (int *)malloc(sizeof(int));
+	*c = 3;
 
-	*a = 5;
-	*b = 12;
-
-	printf("Test 1: %i\n", *a);
-	printf("Test 2: %i\n", *b);
 	printf("Heap Size: %zu Bytes\n", get_heap_size());
 
 	free(a);
 	free(b);
+	free(c);
 
-	int *c = (int *)malloc(sizeof(int));
+	cleanup_memory();
-
-	check_memory();
 
 	return 0;
 }