SquidgeOS/src/kernel/memory.c

#include <stdint.h>
#include <stddef.h>
#include "drivers/uart.h"
#include "memory.h"
#include "panic.h"


HeapHeader* heap_free_list;
extern uint8_t _heap_start[]; // named in the linker script

void page_init() {
	kprint("Initialising page allocator.\n");
	uintptr_t start = ((uintptr_t)_heap_start + PAGE_SIZE - 1) & ~(PAGE_SIZE - 1);
	uintptr_t end = 0x88000000; // Default QEMU RAM limit

	for (uintptr_t addr = start; addr + PAGE_SIZE <= end; addr += PAGE_SIZE) {
		page_free((void *)addr);
	}
}

void page_free(void *addr) {
	if (addr == NULL) return;

	struct Page *p = (struct Page *)addr;
	p->next = free_list;
	free_list = p;
}

void *page_alloc() {
	if (free_list == NULL) {
		kpanic("No free pages!");
	}
	struct Page *p = free_list;
	free_list = free_list->next;

	//zero out the page
	for (int i = 0; i < (PAGE_SIZE); i++) {
		((uint8_t *)p)[i] = 0;
	}

	return (void *)p;
}

void *kmalloc(size_t size) {
	if (size == 0) return NULL;

	// Align size to 8 bytes
	size = (size + 7) & ~7;

	if (heap_free_list == NULL) {
		HeapHeader* header = (HeapHeader *)page_alloc();
		header->size = PAGE_SIZE - sizeof(HeapHeader);
		header->is_free = 1;
		header->next = NULL;
		
		heap_free_list = header;
	}

	HeapHeader *current = heap_free_list;
	HeapHeader *prev = NULL;

	while (current != NULL) {
		if (current->is_free && current->size >= size) {
			break;
		}
		prev = current;
		current = current->next;
	}

	if(current == NULL) {
		// No suitable block found, allocate a new page
		HeapHeader* header = (HeapHeader *)page_alloc();
		header->size = PAGE_SIZE - sizeof(HeapHeader);
		header->is_free = 1;
		header->next = heap_free_list;
		heap_free_list = header;
		current = header;
	}

	// Now, current is a block that can be used
    if (!current) {
        return NULL;
    }

	current->is_free = 0;

	// If the block is larger than needed, split it
	size_t min_split_size = sizeof(HeapHeader) + 16;
	if (current->size >= size + min_split_size) {
		HeapHeader *new_header = (HeapHeader *)((uint8_t *)current + sizeof(HeapHeader) + size);
		// Safety check: if new_header address is wild, abort!
	    if ((uintptr_t)new_header < 0x80000000 || (uintptr_t)new_header > 0x88000000) {
        	kpanic("Splitting created invalid pointer!");
    	}
		new_header->size = current->size - size - sizeof(HeapHeader);
		new_header->is_free = 1;
		new_header->next = current->next;

		current->size = size;
		current->next = new_header;
	}

    return (void*)((char*)current + sizeof(HeapHeader));
}

void kfree(void *ptr) {
	if (ptr == NULL) return;

	HeapHeader *header = (HeapHeader *)((uint8_t *)ptr - sizeof(HeapHeader));
	if(header->is_free) {
		kpanic("Double free detected!");
	}
	header->is_free = 1;

    HeapHeader *temp = heap_free_list;
    while (temp != NULL) {
        if (temp->is_free) {
            kcoalesce(temp);
        }
        temp = temp->next;
    }
}

void kcoalesce(HeapHeader *header) {
	if(!header || !header->next || !header->is_free) return;
	uintptr_t current_end = (uintptr_t)header + sizeof(HeapHeader) + header->size;
	if (current_end == (uintptr_t)header->next && header->next->is_free) {
		header->size += sizeof(HeapHeader) + header->next->size;
		header->next = header->next->next;
		kcoalesce(header);
	}
}

void test_memory_integrity() {
    kprint("Running Integrity Test...\n");
    uint64_t *a = (uint64_t*)kmalloc(16);
    uint64_t *b = (uint64_t*)kmalloc(16);
    
    *a = 0x1122334455667788;
    *b = 0x99AABBCCDDEEFF00;

    if (*a == 0x1122334455667788) {
        kprint("Integrity Pass!\n");
    } else {
        kprint("CORRUPTION DETECTED!\n");
    }
    kfree(a);
    kfree(b);
}

void test_memory_alignment() {
	kprint("Running Alignment Test...\n");
	for (int i = 1; i <= 64; i++) {
		void *ptr = kmalloc(i);
		if (((uintptr_t)ptr % 8) != 0) {
			kprint("Misaligned allocation detected!\n");
			return;
		}
		kfree(ptr);
	}
	kprint("All allocations are properly aligned!\n");
}

void test_memory_stress() {
    kprintf("Starting Stress Test...\n");
	heap_stats();
    void *ptrs[50] = {0}; // Track allocated pointers
    uint32_t seed = 0xACE2026; // Example seed

    for (int i = 0; i < 100; i++) {
		kprintf("Iteration %d\n", i);
        // 1. Randomly allocate or free
        int idx = i % 50;
        if (ptrs[idx] == NULL) {
            size_t size = (seed % 256) + 1;
            ptrs[idx] = kmalloc(size);
            // Optional: fill with data to check integrity later
        } else {
            kfree(ptrs[idx]);
            ptrs[idx] = NULL;
        }
        // Simple LCG to "randomize" seed
        seed = (seed * 1103515245 + 12345) & 0x7fffffff;
    }
    kprintf("Stress Test Finished. Check heap_stats() for sanity.\n");
    heap_stats();
}

void heap_stats() {
    size_t free_size = 0;
    size_t used_size = 0;
    size_t free_blocks = 0;
    size_t used_blocks = 0;

    HeapHeader *current = heap_free_list;

    while (current != NULL) {
        if (current->is_free) {
            free_size += current->size;
            free_blocks++;
        } else {
            used_size += current->size;
            used_blocks++;
        }
        current = current->next;
    }

    kprint("--- Kernel Heap Stats ---\n");
    kprint("Used: "); kprint_hex(used_size); kprint(" bytes in "); kprint_int(used_blocks); kprint(" blocks\n");
    kprint("Free: "); kprint_hex(free_size); kprint(" bytes in "); kprint_int(free_blocks); kprint(" blocks\n");
    kprint("Total metadata overhead: "); kprint_int((used_blocks + free_blocks) * sizeof(HeapHeader)); kprint(" bytes\n");
    kprint("-------------------------\n");
}