From 2bd91c2d940301592b3e8e21220cb3741787137a Mon Sep 17 00:00:00 2001
From: "Justin C. Miller" <justin@devjustinian.com>
Date: Sun, 10 May 2020 16:43:18 -0700
Subject: [PATCH] [boot] Split get_uefi_mappings and module creation

The `get_mappings()` function was getting too large, and some of its
output is needed by more than just the building of the kernel map. Split
it out into two.

Tags: boot memory
---
 src/boot/main.cpp   |  4 +++-
 src/boot/memory.cpp | 24 ++++++++++++++++--------
 src/boot/memory.h   |  9 ++++++++-
 3 files changed, 27 insertions(+), 10 deletions(-)

diff --git a/src/boot/main.cpp b/src/boot/main.cpp
index 4842e55..b070e3e 100644
--- a/src/boot/main.cpp
+++ b/src/boot/main.cpp
@@ -171,7 +171,9 @@ bootloader_main_uefi(uefi::handle image, uefi::system_table *st, console &con)
 	loader::loaded_elf kernel_elf =
 		loader::load(kernel->location, kernel->size, bs);
 
-	memory::get_mappings(bs);
+	memory::efi_mem_map efi_map = memory::get_uefi_mappings(bs);
+	memory::build_kernel_mem_map(efi_map, args, bs);
+
 	return kernel_elf;
 }
 
diff --git a/src/boot/memory.cpp b/src/boot/memory.cpp
index 99c9ca1..63a54f6 100644
--- a/src/boot/memory.cpp
+++ b/src/boot/memory.cpp
@@ -1,6 +1,5 @@
 #include <stddef.h>
 #include <uefi/types.h>
-#include "kernel_args.h"
 
 #include "console.h"
 #include "error.h"
@@ -160,8 +159,10 @@ can_merge(mem_entry &prev, mem_type type, uefi::memory_descriptor *next)
 }
 
 efi_mem_map
-get_mappings(uefi::boot_services *bs)
+get_uefi_mappings(uefi::boot_services *bs)
 {
+	status_line(L"Getting UEFI memory map");
+
 	size_t needs_size = 0;
 	size_t map_key = 0;
 	size_t desc_size = 0;
@@ -173,8 +174,6 @@ get_mappings(uefi::boot_services *bs)
 	if (status != uefi::status::buffer_too_small)
 		error::raise(status, L"Error getting memory map size");
 
-	console::print(L"memory map needs %lu bytes\r\n", needs_size);
-
 	size_t buffer_size = needs_size + 10*desc_size;
 	uefi::memory_descriptor *buffer = nullptr;
 	try_or_raise(
@@ -193,9 +192,15 @@ get_mappings(uefi::boot_services *bs)
 	map.key = map_key;
 	map.version = desc_version;
 	map.entries = buffer;
+	return map;
+}
 
+void
+build_kernel_mem_map(efi_mem_map &efi_map, kernel::args::header *args, uefi::boot_services *bs)
+{
+	status_line(L"Creating kernel memory map");
 
-	size_t map_size = map.num_entries() * sizeof(mem_entry);
+	size_t map_size = efi_map.num_entries() * sizeof(mem_entry);
 	kernel::args::mem_entry *kernel_map = nullptr;
 	try_or_raise(
 		bs->allocate_pages(
@@ -209,7 +214,7 @@ get_mappings(uefi::boot_services *bs)
 
 	size_t i = 0;
 	bool first = true;
-	for (auto desc : map) {
+	for (auto desc : efi_map) {
 		/*
 		console::print(L"   Range %lx (%lx) %x(%s) [%lu]\r\n",
 			desc->physical_start, desc->attribute, desc->type, memory_type_name(desc->type), desc->number_of_pages);
@@ -293,6 +298,11 @@ get_mappings(uefi::boot_services *bs)
 		}
 	}
 
+	kernel::args::module &module = args->modules[args->num_modules++];
+	module.location = reinterpret_cast<void*>(kernel_map);
+	module.size = map_size;
+	module.type = kernel::args::mod_type::memory_map;
+
 	/*
 	for (size_t i = 0; i<map.num_entries(); ++i) {
 		mem_entry &ent = kernel_map[i];
@@ -300,8 +310,6 @@ get_mappings(uefi::boot_services *bs)
 			ent.start, ent.attr, ent.type, ent.pages);
 	}
 	*/
-
-	return map;
 }
 
 
diff --git a/src/boot/memory.h b/src/boot/memory.h
index d685ec5..230e565 100644
--- a/src/boot/memory.h
+++ b/src/boot/memory.h
@@ -4,6 +4,7 @@
 #include <uefi/boot_services.h>
 #include <uefi/runtime_services.h>
 #include <stdint.h>
+#include "kernel_args.h"
 #include "pointer_manipulation.h"
 
 namespace boot {
@@ -78,7 +79,13 @@ struct efi_mem_map
 };
 
 /// Get the memory map from UEFI.
-efi_mem_map get_mappings(uefi::boot_services *bs);
+efi_mem_map get_uefi_mappings(uefi::boot_services *bs);
+
+/// Add the kernel's memory map as a module to the kernel args.
+void build_kernel_mem_map(
+	efi_mem_map &efi_map,
+	kernel::args::header *args,
+	uefi::boot_services *bs);
 
 } // namespace boot
 } // namespace memory