src/memory.c: improve paging code

uses untracked memory for paging, avoiding wasting memory by not
creating headers for memory that won't ever be freed.

also simplify some functions

Signed-off-by: Anna (navi) Figueiredo Gomes <navi@vlhl.dev>

2 files changed, 152 insertions, 90 deletions

diff --git a/include/memory.h b/include/memory.h
index b349e62..0bf3715 100644
--- a/include/memory.h
+++ b/include/memory.h
@@ -2,12 +2,14 @@
 #define _PAGE_H_
 
 #include <stddef.h>
+#include <stdint.h>
 #include <stdbool.h>
 
 void *kmalloc(size_t size);
 void *kmalloc_a(size_t size, size_t alignment);
 void kfree(void *ptr);
-void print_blocks();
-bool paging_init();
+void print_blocks(void);
+void memory_map(uintptr_t addr, size_t size);
+bool paging_init(void);
 
 #endif
diff --git a/src/memory.c b/src/memory.c
index eac0718..a73b616 100644
--- a/src/memory.c
+++ b/src/memory.c
@@ -10,9 +10,10 @@
 extern uint64_t multiboot_magic;
 extern struct multiboot_info *multiboot_header;
 extern void *kernel_end;
-
 extern uint64_t pml4[512];
 
+#define PAGE_SIZE sizeof(uint64_t) * 512
+
 struct memory {
 	struct memory *prev, *next;
 	bool free;
@@ -20,11 +21,24 @@ struct memory {
 	uint8_t block[];
 } *blocklist;
 
-void print_blocks() {
+struct page_mem {
+	void *ptr;
+	size_t size;
+	struct page_mem *next;
+} pagelist = {
+	.ptr = &kernel_end,
+	.size = 0x100000,
+	.next = NULL
+};
+
+
+void print_blocks(void) {
 	struct memory *block = NULL;
 	for (block = blocklist; block; block = block->next) {
-		vga_puts("block at: ");
+		vga_puts("block header at: ");
 		vga_putx((size_t)block);
+		vga_puts(", block at: ");
+		vga_putx((size_t)block->block);
 		vga_puts(", size: ");
 		vga_putx(block->size);
 		if (block->free)
@@ -33,47 +47,45 @@ void print_blocks() {
 	}
 }
 
+static struct memory *split_block(struct memory *block, size_t size) {
+	struct memory *nextblock = (void *)((uintptr_t)block->block + size);
+	*nextblock = (struct memory) {
+		.size = block->size - size - sizeof(*nextblock),
+		.free = true,
+		.prev = block,
+		.next = block->next
+	};
+	block->next = nextblock;
+	block->size = size;
+	block->free = false;
+	return nextblock;
+}
+
 void *kmalloc_a(size_t size, size_t alignment) {
 	struct memory *block = NULL;
 	for (block = blocklist; block; block = block->next) {
 		if (!block->free || block->size < size + sizeof(struct memory))
 			continue;
-		size_t offset = (uintptr_t)block->block % alignment;
-		if (offset) {
-			if (block->size < (2 * sizeof(struct memory)) + size + (alignment - offset)) {
-				continue;
-			}
-
-			struct memory *new_block = (void *)block + (alignment - offset);
-			*new_block = (struct memory) {
-				.size = block->size - (alignment - offset),
-				.free = true,
-				.prev = block,
-				.next = block->next,
-			};
-			block->next = new_block;
-			block->size = alignment - offset;
 
-			block = new_block;
+		size_t offset = (uintptr_t)block->block % alignment;
+		if (!offset)
 			break;
-		}
+
+		if (block->size < (2 * sizeof(struct memory)) + size + (alignment - offset))
+			continue;
+
+		struct memory *new_block = split_block(block, (alignment - offset - sizeof(*block)));
+		block = new_block;
+		break;
 	}
+
 	if (!block) {
-		print_blocks();
+		//print_blocks();
 		vga_puts("\n\noom");
 		while(1);
 	}
 	
-	struct memory *nextblock = (void *)block->block + size;
-	*nextblock = (struct memory) {
-		.size = block->size - size,
-		.free = true,
-		.prev = block,
-		.next = block->next
-	};
-	block->next = nextblock;
-	block->size = size;
-	block->free = false;
+	split_block(block, size);
 	return block->block;
 }
 
@@ -81,34 +93,24 @@ void *kmalloc(size_t size) {
 	struct memory *block = NULL;
 	for (block = blocklist; block && (!block->free || block->size < size + sizeof(struct memory)); block = block->next);
 	if (!block) {
-		print_blocks();
+		//print_blocks();
 		vga_puts("\n\noom");
 		while(1);
 	}
 	
-	struct memory *nextblock = (void *)block->block + size;
-	*nextblock = (struct memory) {
-		.size = block->size - size,
-		.free = true,
-		.prev = block,
-		.next = block->next
-	};
-	block->next = nextblock;
-	block->size = size;
-	block->free = false;
-
+	split_block(block, size);
 	return block->block;
 }
 
 void kfree(void *ptr) {
-	struct memory *header = ptr - sizeof(struct memory);
-	if (header->next && header->next == (void*)header->block + header->size && header->next->free) {
+	struct memory *header = (void *)((uintptr_t)ptr - sizeof(struct memory));
+	if (header->next && header->next == (void*)((uintptr_t)header->block + header->size) && header->next->free) {
 		header->size += header->next->size + sizeof(struct memory);
 		header->next = header->next->next;
 		if (header->next)
 			header->next->prev = header;
 	}
-	if (header->prev && header == (void*)header->prev->block + header->prev->size && header->prev->free) {
+	if (header->prev && header == (void*)((uintptr_t)header->prev->block + header->prev->size) && header->prev->free) {
 		header->prev->size += header->size + sizeof(struct memory);
 		header->prev->next = header->next;
 		if (header->next)
@@ -119,67 +121,125 @@ void kfree(void *ptr) {
 	header->free = true;
 }
 
-uint64_t alloc_page() {
-	uint64_t *page = kmalloc_a(sizeof(uint64_t) * 512, 0x1000);
-	memset(page, 0, sizeof(uint64_t) * 512);
-	return (uint64_t)page | 0x3;
-}
+static uint64_t alloc_page(void) {
+	uint64_t *page;
 
-uint64_t *addr_to_page(uint64_t addr) {
-	size_t index = (addr >> 39) & 0x1ff;
-	if (!(pml4[index] & 0x1))
-		pml4[index] = alloc_page();
+	if (blocklist) {
+		page = kmalloc_a(PAGE_SIZE, 0x1000);
+		memset(page, 0, PAGE_SIZE);
+		return (uint64_t)page | 0x3;
+	}
 
-	uint64_t *pdpr = (uint64_t*)(pml4[index] & ~0xfff);
-	index = (addr >> 30) & 0x1ff;
-	if (!(pdpr[index] & 0x1))
-		pdpr[index] = alloc_page();
+	if (pagelist.size < PAGE_SIZE) {
+		if (!pagelist.next) {
+			vga_puts("oom");
+			while(1);
+		}
+		pagelist = *pagelist.next;
+	}
 
-	uint64_t *pd = (uint64_t*)(pdpr[index] & ~0xfff);
-	index = (addr >> 21) & 0x1ff;
-	if (!(pd[index] & 0x1))
-		pd[index] = alloc_page();
+	if ((uintptr_t) pagelist.ptr & 0xfff) {
+		pagelist.ptr = (void *) (((uintptr_t) pagelist.ptr & ~0xfff) + 0x1000);
+		pagelist.size -= (0x1000 - (uintptr_t) pagelist.ptr % 0x1000);
+	}
 
-	uint64_t *pt = (uint64_t*)(pd[index] & ~0xfff);
-	index = (addr >> 12) & 0x1ff;
-	return &pt[index];
-};
+	page = pagelist.ptr;
+	pagelist.ptr = (void *) ((uintptr_t) pagelist.ptr + PAGE_SIZE);
+	pagelist.size -= PAGE_SIZE;
 
-bool paging_init() {
-	blocklist = (void *)0x7e00;
-	*blocklist = (struct memory) {
-		.size = 0x7ffff - 0x7e00,
-		.free = true
-	};
+	memset(page, 0, PAGE_SIZE);
+	return (uint64_t)page | 0x3;
+}
 
+static uint64_t *addr_to_page(uint64_t addr) {
+	uint64_t *page = pml4;
+	size_t offset = 39;
+	size_t index = (addr >> offset) & 0x1ff;
+
+	for (size_t i = 0; i < 3; i++) {
+		if (!(page[index] & 0x1))
+			page[index] = alloc_page();
+		page = (uint64_t*)(page[index] & ~0xfff);
+		index = (addr >> (offset -= 9)) & 0x1ff;
+	}
+	return &page[index];
+}
+
+void memory_map(uintptr_t addr, size_t size) {
+	for (uintptr_t ptr = addr & ~0xfff; ptr < addr + size; ptr += 0x1000) {
+		*addr_to_page(ptr) = ptr | 0x3 | (1 << 6);
+	}
+}
+
+bool paging_init(void) {
 	if (multiboot_magic != 0x2badb002)
 		return false;
 
 	for (size_t i = 0; i < multiboot_header->mmap_length / sizeof(struct multiboot_mmap_entry); i++) {
 		struct multiboot_mmap_entry *mmap = &((struct multiboot_mmap_entry *) (uintptr_t) multiboot_header->mmap_addr)[i];
 
-		if (mmap->type == MULTIBOOT_MEMORY_AVAILABLE && mmap->addr > 0x1000) {
-			struct memory *header = mmap->addr > (uint64_t)&kernel_end ? (void*) mmap->addr : (void *)&kernel_end;
+		vga_puts("region addr ");
+		vga_putx(mmap->addr);
+		vga_puts(" len ");
+		vga_putx(mmap->len);
+		switch (mmap->type) {
+			case MULTIBOOT_MEMORY_AVAILABLE:
+				vga_puts(" avail");
+				break;
+			case MULTIBOOT_MEMORY_NVS:
+				vga_puts(" nvs");
+				break;
+			case MULTIBOOT_MEMORY_BADRAM:
+				vga_puts(" badram");
+				break;
+			case MULTIBOOT_MEMORY_RESERVED:
+				vga_puts(" reserv");
+				break;
+			case MULTIBOOT_MEMORY_ACPI_RECLAIMABLE:
+				vga_puts(" acpi");
+				break;
+		}
+		vga_putc('\n');
 
-			// map the first header, in case we need to use this region for paging itself
-			uint64_t *page = addr_to_page((uintptr_t) header);
-			*page = ((uintptr_t)header & ~0xfff) | 0x3;
+		if (mmap->type == MULTIBOOT_MEMORY_AVAILABLE && mmap->addr >= 0x100000
+				&& mmap->addr + mmap->len > (uintptr_t)&kernel_end + 0x100000) {
+			struct page_mem *next_page = mmap->addr > (uintptr_t)&kernel_end + 0x100000 ?
+				(void*) mmap->addr : (void *) ((uintptr_t)&kernel_end + 0x100000);
 
-			*header = (struct memory) {
+			*next_page = (struct page_mem) {
+				.ptr = next_page,
 				.size = mmap->len,
-				.free = true
+				.next = NULL
 			};
-			struct memory *last;
-			for (last = blocklist; last->next; last = last->next);
-			last->next = header;
-			header->prev = last;
+
+			struct page_mem *last = &pagelist;
+			for (; last->next; last = last->next);
+			last->next = next_page;
 		}
 
-		size_t end = mmap->addr + mmap->len;
-		for (size_t addr = mmap->addr & ~0xfff; addr < end; addr += 0x1000) {
-			uint64_t *page = addr_to_page(addr);
-			*page = addr | 0x3;
+		for (size_t addr = mmap->addr & ~0xfff; addr <= mmap->addr + mmap->len; addr += 0x1000)
+			*addr_to_page(addr) = addr | 0x3;
+	}
+
+	for (struct page_mem *region = &pagelist; region; region = pagelist.next) {
+		pagelist = *region;
+		struct memory *mem_block = pagelist.ptr;
+		(*mem_block) = (struct memory) {
+			.size = pagelist.size - sizeof(struct memory),
+			.free = true
+		};
+
+		if (!blocklist) {
+			blocklist = mem_block;
+			continue;
 		}
+
+		struct memory *free_block = blocklist;
+		while (free_block->next)
+			free_block = free_block->next;
+		free_block->next = mem_block;
+		mem_block->prev = free_block;
 	}
+
 	return true;
 }