Nucleus

diff --git a/memory/frame.c b/memory/frame.c
index 40075b0..512c47d 100644
--- a/memory/frame.c
+++ b/memory/frame.c
@@ -7,9 +7,14 @@
 
 #include <stdint.h>
 #include <stddef.h>
+#include <string.h>
+#include <sys/types.h>
 #include <nucleus/memory.h>
 #include <nucleus/panic.h>
 
+#define INDEX(a) ((a)/32)
+#define OFFSET(a) ((a)%32)
+
 /* Types of an E820 Memory Map Entry */
 enum E820Type {
 	E820_NONE,
@@ -28,24 +33,99 @@ struct E820Entry {
 	uint32_t type;
 } __attribute__((packed));
 
-static uintptr_t nextFrame;
+/* Descriptor of a Frame Region */
+struct FrameRegion {
+	uintptr_t base;
+	size_t numFrames, usedFrames;
+	struct FrameRegion *next;
+	uint32_t bitmap[];
+};
+
+size_t numFrames, usedFrames;
+struct FrameRegion *regions;
+
+/* Count the number of pages a number of bytes occupies */
+static size_t
+page_count(size_t bytes)
+{
+	size_t pages = bytes / PAGE_SIZE;
+	if (bytes & (PAGE_SIZE - 1))
+		return pages + 1;
+	return pages;
+}
+
+/* Set a bit in a frame's bitset */
+static void
+set_frame(struct FrameRegion *region, off_t index)
+{
+	region->bitmap[INDEX(index)] |= (1 << OFFSET(index));
+	region->usedFrames++;
+	usedFrames++;
+}
+
+/* Clear a bit in a frame's bitset */
+static void
+clear_frame(struct FrameRegion *region, off_t index)
+{
+	region->bitmap[INDEX(index)] &= ~(1 << OFFSET(index));
+	region->usedFrames--;
+	usedFrames--;
+}
+
+/* Test a bit in a frame's bitset */
+static uint32_t
+test_frame(struct FrameRegion *region, off_t index)
+{
+	return (region->bitmap[INDEX(index)] & (1 << OFFSET(index)));
+}
+
+/* Find first free frame */
+static off_t
+find_frame(struct FrameRegion *region)
+{
+	off_t i, j;
+	for (i = 0; i < INDEX(region->numFrames); i++) {
+		if (!~region->bitmap[i])
+			continue;
+		for (j = 0; j < 32; j++) {
+			if ((region->bitmap[i] >> j) & 1)
+				continue;
+			return (i * 32) + j;
+		}
+	}
+	return (off_t) -1;
+}
 
 /* Allocate a page frame */
 uintptr_t
 alloc_frame(void)
 {
-	uintptr_t frame = nextFrame;
-	nextFrame = *((uintptr_t *) frame);
-	*((uintptr_t *) frame) = 0;
-	return frame;
+	off_t idx;
+	/* Walk the regions, first fit */
+	struct FrameRegion *region;
+	for (region = regions; region; region = region->next) {
+		idx = find_frame(region);
+		if (idx != -1)
+			break;
+	}
+	if (idx == -1)
+		return 0x00000000;
+
+	set_frame(region, idx);
+	return region->base + (idx << 12);
 }
 
 /* Free a page frame */
 void
 free_frame(uintptr_t frame)
 {
-	*((uintptr_t *) frame) = nextFrame;
-	nextFrame = frame;
+	/* Walk the regions */
+	struct FrameRegion *region;
+	for (region = regions; region->next; region = region->next) {
+		if (region->next->base >= frame)
+			break;
+	}
+	clear_frame(region, (frame - region->base) >> 12);
 }
 
 /* Setup the page frame allocator */
@@ -57,20 +137,54 @@ init_frames(uint32_t size, void *addr)
 	 * memory.  This map must be read so the memory manager can avoid bad
 	 * areas of memory and areas that are mapped to hardware, ACPI, etc.
 	 */
-
-	size_t i, j;
+	int i, j;
 	struct E820Entry *memMap = addr;
+	struct FrameRegion *prev = NULL, *head = regions;
+	uintptr_t bumpAlloc = PAGE_SIZE;
 	for (i = 0; i < size / sizeof(struct E820Entry); i++) {
 		if (memMap[i].baseHigh > 0
 		 || memMap[i].base > 0xFFFFFFFF
 		 || memMap[i].type != E820_USABLE)
 			continue;
 
-		for (j = 0; j < memMap[i].length; j += PAGE_SIZE) {
-			if ((memMap[i].base + j) >= 0x100000
-			 && (memMap[i].base + j) < 0x180000)
-				continue;
-			free_frame(memMap[i].baseHigh + memMap[i].base + j);
-		}
+		/* Usable region - create bitmap */
+		size_t frameSize = memMap[i].length / PAGE_SIZE;
+		size_t bitmapSize = (frameSize / 8);
+		head = (struct FrameRegion *) bumpAlloc;
+		bumpAlloc += bitmapSize + sizeof(struct FrameRegion);
+		head->base = memMap[i].base;
+		if (sizeof(head->base) > sizeof(memMap[i].base))
+			head->base += memMap[i].baseHigh;
+		head->numFrames = frameSize;
+		head->usedFrames = 0;
+
+		memset(head->bitmap, 0, bitmapSize);
+		/* Bits that don't refer to real frames must be set */
+		for (j = OFFSET(head->numFrames); j < 32; j++)
+			set_frame(head, head->numFrames + j);
+
+		if (prev) prev->next = head;
+		else regions = head;
+		prev = head;
+	}
+
+	/* Regions to be remapped */
+	struct {uint32_t start, end;} remaps[] = {
+		{.start = 0x0000,   .end = bumpAlloc}, /* PMM bitmaps */
+		{.start = 0x100000, .end = 0x180000 }, /* Kernel */
+	};
+
+	/* Check bitmaps */
+	numFrames = usedFrames = 0;
+	struct FrameRegion *region;
+	for (region = regions; region; region = region->next) {
+		numFrames += region->numFrames;
+		uintptr_t end = region->base + (region->numFrames * PAGE_SIZE);
+		/* Iterate the remaps[] to find overlapping regions */
+		for (i = 0; i < sizeof(remaps)/sizeof(remaps[0]); i++)
+			for (j = remaps[i].start;
+			     j < remaps[i].end && j >= region->base && j < end;
+			     j += PAGE_SIZE)
+				set_frame(region, (j - region->base) >> 12);
 	}
 }