3 files changed, 235 insertions, 41 deletions

diff --git a/src/memory/memory.c b/src/memory/memory.c
new file mode 100644
index 0000000..1c69bae
--- /dev/null
+++ b/src/memory/memory.c
@@ -0,0 +1,214 @@
+#include <memory.h>
+#include <stdint.h>
+#include <lib.h>
+
+#ifdef MEMORY_PANIC
+#include <panic.h>
+#endif
+
+#define MAGIC 0xBEEFCAFE
+
+struct page_header {
+	struct page_header *next;
+	struct page_header *prev;
+	size_t node_number;			// all headers on the same page alloc have the same node number (so they can be merged)
+	size_t free;				// free space after the node (if its the last node in the alloc block)
+	size_t used;				// how much space this allocation is using
+	uint64_t magic;
+};
+
+static const size_t header_len = sizeof(struct page_header);
+struct page_header *start_header = NULL;
+struct page_header *end_header = NULL;
+
+struct page_header* get_header(void *ptr) {
+	struct page_header *header =
+		(struct page_header *) ((uintptr_t) ptr - header_len);
+
+	// PERF: do we want to make sure this pointer is paged
+	// before reading it???
+	if (header->magic != MAGIC) {
+		return NULL; // invalid pointer
+	}
+
+	return header;
+}
+
+void *kalloc_new(size_t size) {
+	size_t pages = ((size + header_len) / PAGE_SIZE) + 1;
+	void *addr = alloc_pages(pages);
+	void *mem = (char *)addr + header_len;
+
+	size_t total = pages * PAGE_SIZE;
+	size_t free = total - (size + header_len);
+
+	if (addr == NULL) {
+		return NULL;
+	}
+
+	size_t node;
+	if (end_header != NULL) {
+		node = end_header->node_number + 1;
+	} else {
+		node = 0;
+	}
+
+	struct page_header *header = addr;
+	header->magic = MAGIC;
+	header->used = size;
+	header->free = free;
+	header->prev = end_header;
+	header->next = NULL;
+	header->node_number = node;
+
+	if (end_header == NULL) {
+		start_header = header;
+	} else {
+		end_header->next = header;
+	}
+
+	end_header = header;
+
+	return mem;
+}
+
+void *kalloc_block(size_t size, struct page_header *block) {
+	struct page_header *header =
+		(struct page_header *) ((char *) block + block->used + header_len);
+
+	size_t free = block->free - (size + header_len);
+	block->free = 0;
+
+	header->magic = MAGIC;
+	header->used = size;
+	header->free = free;
+	header->prev = block;
+	header->next = block->next;
+	block->next = header;
+	header->node_number = block->node_number;
+
+	void *mem = (char *) header + header_len;
+
+	return mem;
+}
+
+void *kalloc(size_t size) {
+	struct page_header *header = start_header;
+
+	for (; header != NULL; header = header->next) {
+		size_t free = header->free;
+		if (size <= (free - header_len)) { // we must be able to fit data + header
+			break;
+		}
+	}
+
+	if (header != NULL) {
+		return kalloc_block(size, header);
+	} else {
+		return kalloc_new(size);
+	}
+}
+
+void *krealloc(void *src, size_t dst_len) {
+	struct page_header *header;
+	size_t src_len;
+	void *dst;
+
+	// realloc of 0 means free pointer
+	if (dst_len == 0) {
+		kfree(src);
+		return NULL;
+	}
+
+	// NULL src means allocate ptr
+	if (src == NULL) {
+		dst = kalloc(dst_len);
+		return dst;
+	}
+
+	header = get_header(src);
+
+	if (header == NULL) {
+#ifdef MEMORY_PANIC
+		panic("attempted to realloc on a invalid ptr");
+#else
+		return NULL; // invalid pointer passed
+#endif
+	}
+
+	src_len = header->used;
+
+	if (src_len == 0) {
+#ifdef MEMORY_PANIC
+		panic("attempted to realloc on an empty ptr");
+#else
+		return NULL; // likely double free :(
+#endif
+	}
+
+	dst = kalloc(dst_len);
+
+	if (dst == NULL) {
+		return NULL; // allocation failed
+	}
+
+	memcpy(dst, src, src_len);
+	return dst;
+}
+
+void kfree(void *ptr) {
+	struct page_header *header;
+
+	if (ptr == NULL) {
+		return;
+	}
+
+	header = get_header(ptr);
+
+	if (header == NULL || header->used == 0) {
+#ifdef MEMORY_PANIC
+		panic("attempted to kfree invalid pointer");
+#else
+		return;
+#endif
+	}
+
+	header->free += header->used;
+	header->used = 0;
+
+	struct page_header *neighbor;
+
+	// merge left
+	for (neighbor = header->prev; neighbor != NULL; neighbor = neighbor->prev) {
+		if (neighbor->node_number != header->node_number)
+			break;
+		if (neighbor->used && header->used)
+			break;
+		neighbor->free += header->free + header_len;
+		neighbor->next = header->next;
+		header = neighbor;
+	}
+
+	// merge right
+	for (neighbor = header->next; neighbor != NULL; neighbor = neighbor->next) {
+		if (neighbor->node_number != header->node_number)
+			break;
+		if (neighbor->used)
+			break;
+		header->free += neighbor->free + header_len;
+		header->next = neighbor->next;
+	}
+
+	if (
+			(header->next == NULL || header->next->node_number != header->node_number) &&
+			(header->prev == NULL || header->prev->node_number != header->node_number) &&
+			header->used == 0
+	) {
+		if (header->next)
+			header->next->prev = header->prev;
+		if (header->prev)
+			header->prev->next = header->next;
+		free_pages(header);
+	}
+
+}
diff --git a/src/memory/physalloc.c b/src/memory/physalloc.c
index dccd980..1140258 100644
--- a/src/memory/physalloc.c
+++ b/src/memory/physalloc.c
@@ -163,13 +163,11 @@ void memory_init(struct memory_map *map) {
 	page_free_start = 0;
 	page_start = NULL;
 
-	uintptr_t end = (uintptr_t) map;
-	end += map->size;
-
-	struct memory_segment *segment = &map->entries[0];
 	segment_count = 0;
-	
-	for(; (uintptr_t) segment < end; segment++) {
+
+	for(uint32_t i = 0; i < map->entry_count; i++) {
+		struct memory_segment *segment = &map->entries[i];
+
 		if (segment_invalid(segment))
 			continue;
 
@@ -189,16 +187,17 @@ void memory_init(struct memory_map *map) {
 	char *page_area_addr = (char *)bitmap + bitmap_size;
 	bitmap = mmap(bitmap, bitmap_size);
 	memset(bitmap, 0, bitmap_size);
-	
+
 	memory_start = page_align(kaddr(kernel_end) + bitmap_size + page_area_size);
 
 	page_area_addr = mmap(page_area_addr, page_area_size);
 	page_start = (struct memory_area *) page_area_addr;
-	
+
 	struct memory_area *area = page_start;
-	segment = &map->entries[0];
-	
-	for(; (uintptr_t) segment < end; segment++) {
+
+	for(uint32_t i = 0; i < map->entry_count; i++) {
+		struct memory_segment *segment = &map->entries[i];
+
 		if (segment_invalid(segment))
 			continue;
 
@@ -206,7 +205,7 @@ void memory_init(struct memory_map *map) {
 		*area = temp;
 		area++;
 	}
-	
+
 	page_count -= bitmap_pages;
 
 	memory_unlock();
@@ -229,24 +228,3 @@ uint64_t memory_used(void) {
 	return total_memory - free_memory;
 }
 
-// stubs
-// simon do these ik you want to
-// :3
-
-void *malloc(size_t size) {
-	//TODO: implement
-	(void)size;
-	return NULL;
-}
-
-void *realloc(void *ptr, size_t size) {
-	//TODO: implement
-	(void)ptr;
-	(void)size;
-	return NULL;
-}
-
-void free(void *ptr) {
-	//TODO: implement
-	(void)ptr;
-}
diff --git a/src/memory/virtalloc.c b/src/memory/virtalloc.c
index f8d8169..781fc9b 100644
--- a/src/memory/virtalloc.c
+++ b/src/memory/virtalloc.c
@@ -21,20 +21,23 @@ static struct addr_node nodes[BOOTSTRAP_BSS_NODES];
 static struct addr_node *start_node;
 
 static struct addr_node *alloc_node(void) {
+	struct addr_node *node = NULL;
 	if (bss_nodes >= BOOTSTRAP_BSS_NODES) {
-		//FIXME: alloc on heap
+		node = kalloc(sizeof(struct addr_node));
+		if (node == NULL)
+			return NULL;
+		node->is_bss = false;
 	} else {
-		struct addr_node *node = &nodes[bss_nodes];
+		node = &nodes[bss_nodes];
 		bss_nodes += 1;
 		node->is_bss = true;
-		return node;
 	}
-	return NULL;
+	return node;
 }
 
 static void free_node(struct addr_node *node) {
 	if (!node->is_bss)
-		free(node);
+		kfree(node);
 }
 
 void virtaddr_init(void) {
@@ -53,7 +56,6 @@ void virtaddr_init(void) {
 
 void *virtaddr_alloc(int n_pages) {
 
-
 	if (n_pages < 1)
 		return NULL;
 
@@ -84,7 +86,7 @@ void *virtaddr_alloc(int n_pages) {
 			return (void *) new->start;
 		}
 	}
-	
+
 	return NULL;
 }
 
@@ -121,7 +123,7 @@ long virtaddr_free(void *virtaddr) {
 	for (; node != NULL; node = node->next) {
 		if (node->start == virt) {
 			int length = node->end - node->start;
-			int pages = length / PAGE_SIZE; 
+			int pages = length / PAGE_SIZE;
 			merge_back(node);
 			merge_forward(node);
 			return pages;