Add allocator
diff --git a/src/kernel/Makefile b/src/kernel/Makefile
index f8f2117..4137f84 100644
--- a/src/kernel/Makefile
+++ b/src/kernel/Makefile
@@ -11,7 +11,9 @@
 			timer.o \
 			paging.o \
 			switch_table.o \
-			scan_codes.o
+			scan_codes.o \
+			kheap.o \
+			alloc.o
 CFLAGS = -nostdlib -nostdinc -fno-builtin -fno-stack-protector -ffreestanding \
 			-m32 -O2 -g -Wall -Wno-unused-function -Wno-unused-variable
 LDFLAGS = -Tlink.ld -melf_i386
diff --git a/src/kernel/alloc.c b/src/kernel/alloc.c
new file mode 100644
index 0000000..fee5edb
--- /dev/null
+++ b/src/kernel/alloc.c
@@ -0,0 +1,164 @@
+#include "alloc.h"
+#include "kheap.h"
+#include "kint.h"
+#include "log.h"
+#include "paging.h"
+#include "io.h"
+
+extern uint end;
+static size_t palloc_base = (size_t)&end;
+static size_t malloc_base = (size_t)&end + 0x2000;
+
+#define HEADER_SIZE sizeof(struct heap_alloc_header)
+#define FOOTER_SIZE sizeof(struct heap_alloc_footer)
+
+static struct min_heap heap = {0};
+
+void *_kmalloc(size_t size, bool align, void **phys)
+{
+	if (align && (palloc_base & 0xfff)) // if not yet aligned
+	{
+		palloc_base &= ~0xfff;
+		palloc_base += 0x1000;
+	}
+
+	if (phys)
+	{
+		*phys = (void *)palloc_base;
+	}
+
+	size_t addr = palloc_base;
+	palloc_base += size;
+	return (void *)addr;
+}
+
+void *kmalloc(size_t size)
+{
+	return _kmalloc(size, false, NULL);
+}
+
+void *kmalloc_a(size_t size)
+{
+	return _kmalloc(size, true, NULL);
+}
+
+void *kmalloc_ap(size_t size, void **p)
+{
+	return _kmalloc(size, true, p);
+}
+
+// Proper allocators
+
+void init_allocator()
+{
+	struct heap_alloc_header *h = (struct heap_alloc_header *)malloc_base;
+	h->allocated = false;
+	h->magic = HEAP_MAGIC;
+	h->size = 0xC0400000 - malloc_base;
+
+	heap.size = 1;
+	heap.elements[0] = (struct heap_entry){
+		.key = 0xC0200000 - malloc_base,
+		.address = malloc_base,
+	};
+	memset((void *)malloc_base, 0, heap.elements[0].key);
+}
+
+void *malloc(size_t size)
+{
+	bool ok;
+	size_t full_size = size + HEADER_SIZE + FOOTER_SIZE;
+	int i;
+
+	struct heap_entry e = heap_lookup_min(&heap, full_size, &ok, false, &i);
+
+	if (ok)
+	{
+		// Found smallest hole
+		struct heap_alloc_header *h = (struct heap_alloc_header *)e.address;
+
+		kassert(!h->allocated,
+				"Gap already allocated (this should never happen)");
+
+		size_t old_size = h->size;
+
+		if (full_size == old_size)
+		{
+			// Completely used, no need to change anything!
+			heap_delete(&heap, i);
+		}
+		else
+		{
+			// If there isn't very much space left
+			size_t new_size = old_size - full_size;
+			if (new_size <= HEADER_SIZE + FOOTER_SIZE + 8)
+			{
+				full_size = old_size;
+				heap_delete(&heap, i);
+			}
+			else
+			{
+				struct heap_alloc_footer *old_f =
+					(struct heap_alloc_footer *)(e.address + old_size -
+												 FOOTER_SIZE);
+
+				// Else create a new header
+				size_t new_header_addr = e.address + full_size;
+				struct heap_alloc_header *h =
+					(struct heap_alloc_header *)new_header_addr;
+
+				h->allocated = true;
+				h->magic = HEAP_MAGIC;
+				h->size = new_size;
+				old_f->size = new_size;
+
+				heap_decrease(&heap, i,
+							  (struct heap_entry){
+								  .key = new_size,
+								  .address = new_header_addr,
+							  });
+			}
+
+			struct heap_alloc_footer *f =
+				(struct heap_alloc_footer *)(e.address + full_size -
+											 FOOTER_SIZE);
+
+			h->size = full_size;
+			f->size = h->size;
+		}
+
+		return (void *)(e.address + HEADER_SIZE);
+	}
+	else
+	{
+		// We need more memory :L
+		kpanic("Whoops, malloc ran out of memory");
+	}
+}
+
+void free(void *mem)
+{
+	// TODO: expand memory
+	
+	if (!mem)
+		return; // freeing NULL ptr
+
+	struct heap_alloc_header *base =
+		(struct heap_alloc_header *)((size_t)mem - HEADER_SIZE);
+
+	if (base->magic != HEAP_MAGIC)
+	{
+		kpanic("Freeing memory not allocated with malloc()");
+	}
+
+	base->allocated = false;
+	// Add entry into heap
+
+	struct heap_entry entry =
+	{
+		.key = base->size,
+		.address = (size_t)base,
+	};
+	
+	heap_insert(&heap, entry);
+}
diff --git a/src/kernel/alloc.h b/src/kernel/alloc.h
new file mode 100644
index 0000000..4cb046d
--- /dev/null
+++ b/src/kernel/alloc.h
@@ -0,0 +1,13 @@
+#pragma once
+
+#include "kint.h"
+
+void *_kmalloc(size_t size, bool align, void **phys);
+void *kmalloc(size_t size);
+void *kmalloc_a(size_t size);
+void *kmalloc_ap(size_t size, void **p);
+
+void *malloc(size_t size);
+void free(void *mem);
+
+void init_allocator();
diff --git a/src/kernel/kheap.c b/src/kernel/kheap.c
new file mode 100644
index 0000000..31de626
--- /dev/null
+++ b/src/kernel/kheap.c
@@ -0,0 +1,138 @@
+#include "kheap.h"
+#include "log.h"
+
+static void swap_entries(struct heap_entry *a, struct heap_entry *b)
+{
+	// Static to not waste stack space, NOT thread safe
+	static struct heap_entry tmp;
+	tmp = *a;
+	*a = *b;
+	*b = tmp;
+}
+
+int heap_parent(int i)
+{
+	return (i - 1) / 2;
+}
+
+int heap_left(int i)
+{
+	return 2 * i + 1;
+}
+
+int heap_right(int i)
+{
+	return 2 * i + 2;
+}
+
+static void heap_sort(struct min_heap *heap, int i)
+{
+	int l = heap_left(i), r = heap_right(i);
+	int smallest = i;
+	int size = heap->size;
+	struct heap_entry *e = heap->elements;
+
+	if (l < size && e[l].key < e[smallest].key)
+		smallest = l;
+	if (r < size && e[r].key < e[smallest].key)
+		smallest = r;
+
+	if (smallest != i)
+	{
+		swap_entries(&e[i], &e[smallest]);
+		heap_sort(heap, smallest);
+	}
+}
+
+static void heap_delete_root(struct min_heap *heap)
+{
+	if (heap->size <= 0)
+		return;
+	if (heap->size == 1)
+		heap->size--;
+
+	heap->elements[0] = heap->elements[--heap->size];
+	heap_sort(heap, 0);
+}
+
+void heap_decrease(struct min_heap *heap, int i, struct heap_entry to)
+{
+	heap->elements[i] = to;
+	int k = to.key;
+
+	while (i && heap->elements[heap_parent(i)].key > k)
+	{
+		swap_entries(&heap->elements[i], &heap->elements[heap_parent(i)]);
+		i = heap_parent(i);
+	}
+}
+
+void heap_delete(struct min_heap *heap, int i)
+{
+	heap_decrease(heap, i,
+				  (struct heap_entry){
+					  .key = -1,
+					  .address = 0,
+				  });
+	heap_delete_root(heap);
+}
+
+void heap_insert(struct min_heap *heap, struct heap_entry e)
+{
+	kassert(heap->size < HEAP_SIZE, "Heap overflow!");
+
+	int i = heap->size++;
+	uint k = e.key;
+	heap->elements[i] = e;
+
+	// While the parent element is GREATER than the child
+	while (i && heap->elements[heap_parent(i)].key > k)
+	{
+		swap_entries(&heap->elements[i], &heap->elements[heap_parent(i)]);
+		i = heap_parent(i);
+	}
+}
+
+struct heap_entry heap_lookup_min(struct min_heap *heap, int min, bool *ok,
+								  bool delete, int *i)
+{
+	int j = 0;
+	struct heap_entry *e = heap->elements;
+
+	while (j < heap->size && e[j].key < min)
+	{
+		int left = heap_left(j), right = heap_right(j);
+
+		// If left is invalid, use right, if right is invalid, use left.
+		// If both are valid, chose the smallest one.
+		// In the first case, right may too be invalid, but it will be
+		// sorted out in the while loop (see else clause at end of function).
+
+		if (left >= heap->size)
+			j = right;
+		else if (right >= heap->size)
+			j = left;
+		else
+			j = e[left].key < e[right].key ? left : right;
+	}
+
+	if (j < heap->size)
+	{
+		struct heap_entry val = e[j];
+
+		*ok = true;
+
+		if (i)
+			*i = j;
+
+		if (delete)
+			heap_delete(heap, j);
+
+		return val;
+	}
+	else
+	{
+		*ok = false;
+		return (struct heap_entry){0};
+	}
+}
diff --git a/src/kernel/kheap.h b/src/kernel/kheap.h
new file mode 100644
index 0000000..bc42620
--- /dev/null
+++ b/src/kernel/kheap.h
@@ -0,0 +1,42 @@
+#pragma once
+
+#include "kint.h"
+
+#define HEAP_SIZE 4096
+#define HEAP_MAGIC 0xCAFEBABE
+
+struct heap_entry
+{
+	uint key;
+	size_t address;
+};
+
+struct min_heap
+{
+	struct heap_entry elements[HEAP_SIZE];
+	uint size;
+};
+
+struct heap_alloc_header
+{
+	uint magic;
+	bool allocated;
+	// size = size from beginning of header to end of footer
+	size_t size;
+};
+
+struct heap_alloc_footer
+{
+	// size = size from beginning of header to end of footer
+	size_t size;
+};
+
+int heap_parent(int i);
+int heap_left(int i);
+int heap_right(int i);
+void heap_delete(struct min_heap *heap, int k);
+void heap_insert(struct min_heap *heap, struct heap_entry e);
+void heap_decrease(struct min_heap *heap, int k, struct heap_entry to);
+
+struct heap_entry heap_lookup_min(struct min_heap *heap, int min, bool *ok,
+								  bool delete, int *i);
diff --git a/src/kernel/kint.h b/src/kernel/kint.h
index 44b9bfd..8f934fa 100644
--- a/src/kernel/kint.h
+++ b/src/kernel/kint.h
@@ -7,7 +7,7 @@
 
 typedef unsigned long size_t;
 
-typedef uchar bool;
+typedef _Bool bool;
 
 enum
 {
diff --git a/src/kernel/main.c b/src/kernel/main.c
index a866852..9485def 100644
--- a/src/kernel/main.c
+++ b/src/kernel/main.c
@@ -1,9 +1,10 @@
+#include "alloc.h"
 #include "descriptor_tables.h"
 #include "io.h"
 #include "log.h"
+#include "paging.h"
 #include "timer.h"
 #include "vga.h"
-#include "paging.h"
 
 int kmain(void *mboot)
 {
@@ -22,20 +23,34 @@
 	vga_set_color(WHITE, BLACK);
 
 	init_timer(20);
-
+	init_allocator();
 	init_kbd();
+
+	// Test allocator
+
+	int *one = malloc(sizeof(int));
+	int *two = malloc(sizeof(int));
+
+	*one = 1;
+	*two = 2;
+
+	int *array = malloc(sizeof(int[12]));
+
+	for (int i = 0; i < 12; i++)
+		array[i] = i;
+
+	kprintf("Allocated one, two, array[3] = %d, %d, %d\n", *one, *two,
+			array[3]);
+	kprintf("[%x, %x, %x]\n", one, two, array);
+
+	kprintf("Freeing two\n");
+	free(two);
+	int *four = malloc(sizeof(int));
+	*four = 4;
+	kprintf("Allocated four = %d (%x)\n", *four, four);
+
 	asm volatile("sti");
 
-
-#ifdef TEST_PAGE_FAULT
-	kprintf("Causing page fault:\n");
-
-	volatile uint *ptr = (uint *) 0xa0000000;
-	volatile uint cause_page_fault = *ptr;
-
-	kprintf("Should have caused page fault: %d...\n", cause_page_fault);
-#endif
-
 	while (true)
 		asm volatile("hlt");
 
diff --git a/src/kernel/paging.c b/src/kernel/paging.c
index 9da17dc..62636b3 100644
--- a/src/kernel/paging.c
+++ b/src/kernel/paging.c
@@ -1,51 +1,16 @@
 #include "paging.h"
+#include "alloc.h"
 #include "io.h"
 #include "kint.h"
 #include "log.h"
 #include "pic.h"
 
-extern uint end;
-static size_t alloc_base = (size_t)&end;
-
 /* frames bitset, 0 = free, 1 = used */
-static uint *frames;
+static uint frames[0xffffffff / 0x1000 / 32];
 static ulong num_frames;
 
 static uint first_page_table[1024] __attribute__((aligned(4096)));
-static uint page_directory[1024] __attribute__((aligned(4096)));
-
-void *_kmalloc(size_t size, bool align, void **phys)
-{
-	if (align && (alloc_base & 0xfff)) // if not yet aligned
-	{
-		alloc_base &= ~0xfff;
-		alloc_base += 0x1000;
-	}
-
-	if (phys)
-	{
-		*phys = (void *)alloc_base;
-	}
-
-	size_t addr = alloc_base;
-	alloc_base += size;
-	return (void *)addr;
-}
-
-void *kmalloc(size_t size)
-{
-	return _kmalloc(size, false, NULL);
-}
-
-void *kmalloc_a(size_t size)
-{
-	return _kmalloc(size, true, NULL);
-}
-
-void *kmalloc_ap(size_t size, void **p)
-{
-	return _kmalloc(size, true, p);
-}
+static uint kernel_page_directory[1024] __attribute__((aligned(4096)));
 
 /* frame utils */
 
@@ -94,15 +59,15 @@
 	kpanic("first_free_frame failed! no free frames");
 }
 
-void alloc_frame(uint *page, bool user, bool writable)
+void alloc_frame(uint *page_table_entry, bool user, bool writable)
 {
-	if (*page >> 12)
+	if (*page_table_entry >> 12)
 		return; /* frame already allocated */
 
 	uint frame = first_free_frame();
 	//	kprintf("first_free_frame found %d\n", frame);
 	set_frame(frame * 0x1000); /* mark as mapped */
-	*page = frame | 1 | writable << 1 | user << 2;
+	*page_table_entry = frame | 1 | writable << 1 | user << 2;
 }
 
 void free_frame(uint page)
@@ -136,9 +101,32 @@
 	return page_table;
 }
 
+void alloc_kernel_page(uint *virt)
+{
+	// Page number % pages per table
+	uint page = ((size_t)virt / 0x1000) % 1024;
+	uint *table = get_or_create_table(kernel_page_directory, (size_t)virt >> 22,
+									  false, false);
+
+	alloc_frame(&table[page], false, false);
+}
+
+void alloc_kernel_page_range(uint *from, uint *to)
+{
+	uint f = (size_t)from / 0x1000,
+		t = (size_t)to / 0x1000;
+
+	do
+	{
+		alloc_kernel_page((uint *)(size_t)t);
+		t += 0x1000; // next page
+	} while (f < t);
+}
+
 void map_page(uint *dir, size_t virt_start, bool user, bool rw)
 {
-	uint page = virt_start / 0x1000;
+	// Page number % pages per table
+	uint page = (virt_start / 0x1000) % 1024;
 	uint table = virt_start >> 22;
 	uint frame = first_free_frame();
 
@@ -158,11 +146,11 @@
 
 void init_paging()
 {
-	frames = kmalloc_a(0x1000);
-	memset(page_directory, 0, 1024 * 4);
-	map_4mb(page_directory, (size_t)KERNEL_VIRTUAL_BASE, 0, false, false);
+	memset(kernel_page_directory, 0, 1024 * 4);
+	map_4mb(kernel_page_directory, (size_t)KERNEL_VIRTUAL_BASE, 0, false,
+			false);
 
-	load_page_directory((uint)page_directory - 0xC0000000);
+	load_page_directory((uint)kernel_page_directory - 0xC0000000);
 	add_interrupt_handler(14, page_fault);
 }
 
diff --git a/src/kernel/paging.h b/src/kernel/paging.h
index 921524f..8128d43 100644
--- a/src/kernel/paging.h
+++ b/src/kernel/paging.h
@@ -12,11 +12,8 @@
 extern uint load_page_directory(uint table_address);
 extern void enable_paging();
 
-void *_kmalloc(size_t size, bool align, void **phys);
-void *kmalloc(size_t size);
-void *kmalloc_a(size_t size);
-void *kmalloc_ap(size_t size, void **p);
-
 void init_paging();
 
+void alloc_frame(uint *page_table_entry, bool user, bool writable);
+void alloc_kernel_page(uint *page);
 void page_fault(struct registers *regs);