From 6d195d82f54f9f445eb9438f957a0f9eb00a4ae6 Mon Sep 17 00:00:00 2001
From: Achilleas Anagnostopoulos <achilleas.imap.mbox@gmail.com>
Date: Sun, 18 Jun 2017 08:58:05 +0100
Subject: [PATCH] Make vmm use the bitmap allocator after it is initialized

---
 kernel/mem/pmm/allocator/bitmap_allocator.go      | 13 ++++++++++++-
 kernel/mem/pmm/allocator/bitmap_allocator_test.go |  9 +++++++--
 2 files changed, 19 insertions(+), 3 deletions(-)

diff --git a/kernel/mem/pmm/allocator/bitmap_allocator.go b/kernel/mem/pmm/allocator/bitmap_allocator.go
index 7fcd817..392f237 100644
--- a/kernel/mem/pmm/allocator/bitmap_allocator.go
+++ b/kernel/mem/pmm/allocator/bitmap_allocator.go
@@ -306,11 +306,22 @@ func earlyAllocFrame() (pmm.Frame, *kernel.Error) {
 	return earlyAllocator.AllocFrame()
 }
 
+// sysAllocFrame is a helper that delegates a frame allocation request to the
+// bitmap allocator instance.
+func sysAllocFrame() (pmm.Frame, *kernel.Error) {
+	return FrameAllocator.AllocFrame()
+}
+
 // Init sets up the kernel physical memory allocation sub-system.
 func Init(kernelStart, kernelEnd uintptr) *kernel.Error {
 	earlyAllocator.init(kernelStart, kernelEnd)
 	earlyAllocator.printMemoryMap()
 
 	vmm.SetFrameAllocator(earlyAllocFrame)
-	return FrameAllocator.init()
+	if err := FrameAllocator.init(); err != nil {
+		return err
+	}
+	vmm.SetFrameAllocator(sysAllocFrame)
+
+	return nil
 }
diff --git a/kernel/mem/pmm/allocator/bitmap_allocator_test.go b/kernel/mem/pmm/allocator/bitmap_allocator_test.go
index 702e329..0bb97e6 100644
--- a/kernel/mem/pmm/allocator/bitmap_allocator_test.go
+++ b/kernel/mem/pmm/allocator/bitmap_allocator_test.go
@@ -338,7 +338,7 @@ func TestBitmapAllocatorAllocAndFreeFrame(t *testing.T) {
 		for expFrame := pool.startFrame; expFrame <= pool.endFrame; expFrame++ {
 			got, err := alloc.AllocFrame()
 			if err != nil {
-				t.Fatalf("[pool %d] unexpected error: %v", err)
+				t.Fatalf("[pool %d] unexpected error: %v", poolIndex, err)
 			}
 
 			if got != expFrame {
@@ -364,7 +364,7 @@ func TestBitmapAllocatorAllocAndFreeFrame(t *testing.T) {
 	for poolIndex, pool := range alloc.pools {
 		for frame := pool.startFrame; frame <= pool.endFrame; frame++ {
 			if err := alloc.FreeFrame(frame); err != nil {
-				t.Fatalf("[pool %d] unexpected error: %v", err)
+				t.Fatalf("[pool %d] unexpected error: %v", poolIndex, err)
 			}
 		}
 
@@ -411,6 +411,11 @@ func TestAllocatorPackageInit(t *testing.T) {
 		if err := Init(0x100000, 0x1fa7c8); err != nil {
 			t.Fatal(err)
 		}
+
+		// At this point sysAllocFrame should work
+		if _, err := sysAllocFrame(); err != nil {
+			t.Fatal(err)
+		}
 	})
 
 	t.Run("error", func(t *testing.T) {