[boot] Improve bootloader allocation accounting

The bootloader relied on the kernel to know which parts of memory to not
allocate over. For the future shift of having the init process load
other processes instead of the kernel, the bootloader needs a mechanism
to just hand the kernel a list of allocations. This is now done through
the new bootloader allocator, which all allocation goes through. Pool
memory will not be tracked, and so can be overwritten - this means the
args structure and its other structures like programs need to be handled
right away, or copied by the kernel.

- Add bootloader allocator
- Implement a new linked-list based set of pages that act as allocation
  registers
- Allow for operator new in the bootloader, which goes through the
  global allocator for pool memory
- Split memory map and frame accouting code in the bootloader into
  separate memory_map.* files
- Remove many includes that could be replaced by forward declaration in
  the bootloader
- Add a new global template type, `counted`, which replaces the
  bootloader's `buffer` type, and updated kernel args structure to use it.
- Move bootloader's pointer_manipulation.h to the global include dir
- Make offset_iterator try to return references instead of pointers to
  make it more consistent with static array iteration
- Implement a stub atexit() in the bootloader to satisfy clang

src/boot/allocator.cpp

@@ -0,0 +1,131 @@
+#include <uefi/boot_services.h>
+#include <uefi/types.h>
+
+#include "kutil/no_construct.h"
+#include "allocator.h"
+#include "error.h"
+#include "kernel_args.h"
+#include "memory.h"
+
+namespace boot {
+
+kutil::no_construct<memory::allocator> __g_alloc_storage;
+memory::allocator &g_alloc = __g_alloc_storage.value;
+
+namespace memory {
+
+using kernel::init::allocation_register;
+using kernel::init::page_allocation;
+
+static_assert(sizeof(allocation_register) == page_size);
+
+
+void
+init_allocator(uefi::boot_services *bs)
+{
+	new (&g_alloc) allocator(*bs);
+}
+
+
+allocator::allocator(uefi::boot_services &bs) :
+	m_bs(bs),
+	m_register(nullptr),
+	m_current(nullptr)
+{}
+
+void
+allocator::add_register()
+{
+	allocation_register *reg = nullptr;
+
+	try_or_raise(
+		m_bs.allocate_pages(uefi::allocate_type::any_pages,
+			uefi::memory_type::loader_data, 1, reinterpret_cast<void**>(&reg)),
+		L"Failed allocating allocation register page");
+
+	m_bs.set_mem(reg, sizeof(allocation_register), 0);
+
+	if (!m_register) {
+		m_register = m_current = reg;
+		return;
+	}
+
+	m_current->next = reg;
+	m_current = reg;
+	return;
+}
+
+void *
+allocator::allocate_pages(size_t count, alloc_type type, bool zero)
+{
+	if (count & ~0xffffffffull) {
+		error::raise(uefi::status::unsupported,
+			L"Cannot allocate more than 16TiB in pages at once.",
+			__LINE__);
+	}
+
+	if (!m_current || m_current->count == 0xff)
+		add_register();
+
+	void *pages = nullptr;
+
+	try_or_raise(
+		m_bs.allocate_pages(uefi::allocate_type::any_pages,
+			uefi::memory_type::loader_data, count, &pages),
+		L"Failed allocating usable pages");
+
+	page_allocation &ent = m_current->entries[m_current->count++];
+	ent.address = reinterpret_cast<uintptr_t>(pages);
+	ent.count = count;
+	ent.type = type;
+
+	if (zero)
+		m_bs.set_mem(pages, count * page_size, 0);
+
+	return pages;
+}
+
+void *
+allocator::allocate(size_t size, bool zero)
+{
+	void *p = nullptr;
+	try_or_raise(
+		m_bs.allocate_pool(uefi::memory_type::loader_data, size, &p),
+		L"Could not allocate pool memory");
+
+	if (zero)
+		m_bs.set_mem(p, size, 0);
+
+	return p;
+}
+
+void
+allocator::free(void *p)
+{
+	try_or_raise(
+		m_bs.free_pool(p),
+		L"Freeing pool memory");
+}
+
+void
+allocator::memset(void *start, size_t size, uint8_t value)
+{
+	m_bs.set_mem(start, size, value);
+}
+
+void
+allocator::copy(void *to, void *from, size_t size)
+{
+	m_bs.copy_mem(to, from, size);
+}
+
+} // namespace memory
+} // namespace boot
+
+
+void * operator new (size_t size, void *p) { return p; }
+void * operator new(size_t size)           { return boot::g_alloc.allocate(size); }
+void * operator new [] (size_t size)       { return boot::g_alloc.allocate(size); }
+void operator delete (void *p) noexcept    { return boot::g_alloc.free(p); }
+void operator delete [] (void *p) noexcept { return boot::g_alloc.free(p); }
+

src/boot/allocator.h

@@ -0,0 +1,57 @@
+#pragma once
+/// \file allocator.h
+/// Page allocator class definition
+
+#include "kernel_args.h"
+
+namespace uefi {
+	class boot_services;
+}
+
+namespace boot {
+namespace memory {
+
+using alloc_type = kernel::init::allocation_type;
+
+class allocator
+{
+public:
+	using allocation_register = kernel::init::allocation_register;
+
+	allocator(uefi::boot_services &bs);
+
+	void * allocate(size_t size, bool zero = false);
+	void free(void *p);
+
+	void * allocate_pages(size_t count, alloc_type type, bool zero = false);
+
+	void memset(void *start, size_t size, uint8_t value);
+	void copy(void *to, void *from, size_t size);
+
+	inline void zero(void *start, size_t size) { memset(start, size, 0); }
+
+	allocation_register * get_register() { return m_register; }
+
+private:
+	void add_register();
+
+	uefi::boot_services &m_bs;
+
+	allocation_register *m_register;
+	allocation_register *m_current;
+};
+
+/// Initialize the global allocator
+void init_allocator(uefi::boot_services *bs);
+
+} // namespace memory
+
+extern memory::allocator &g_alloc;
+
+} // namespace boot
+
+void * operator new (size_t size, void *p);
+void * operator new(size_t size);
+void * operator new [] (size_t size);
+void operator delete (void *p) noexcept;
+void operator delete [] (void *p) noexcept;

src/boot/console.cpp

@@ -1,9 +1,11 @@
 #include <stddef.h>
 #include <stdint.h>
 
-#include <uefi/types.h>
+#include <uefi/boot_services.h>
 #include <uefi/graphics.h>
 #include <uefi/protos/graphics_output.h>
+#include <uefi/protos/simple_text_output.h>
+#include <uefi/types.h>
 
 #include "console.h"
 #include "error.h"

src/boot/console.h

@@ -3,10 +3,13 @@
 #pragma once
 #include <stdarg.h>
 #include <stddef.h>
-#include <uefi/boot_services.h>
-#include <uefi/protos/simple_text_output.h>
 #include "kernel_args.h"
-#include "types.h"
+
+namespace uefi {
+	struct boot_services;
+namespace protos {
+	struct simple_text_output;
+}}
 
 namespace boot {
 

src/boot/fs.cpp

@@ -1,34 +1,35 @@
+#include <uefi/boot_services.h>
 #include <uefi/types.h>
 #include <uefi/protos/file.h>
 #include <uefi/protos/file_info.h>
 #include <uefi/protos/loaded_image.h>
 #include <uefi/protos/simple_file_system.h>
 
-#include "fs.h"
+#include "allocator.h"
 #include "console.h"
 #include "error.h"
+#include "fs.h"
 #include "memory.h"
 #include "status.h"
 
 namespace boot {
 namespace fs {
 
-file::file(uefi::protos::file *f, uefi::boot_services *bs) :
+using memory::alloc_type;
-	m_file(f),
+
-	m_bs(bs)
+file::file(uefi::protos::file *f) :
+	m_file(f)
 {
 }
 
 file::file(file &o) :
-	m_file(o.m_file),
+	m_file(o.m_file)
-	m_bs(o.m_bs)
 {
 	o.m_file = nullptr;
 }
 
 file::file(file &&o) :
-	m_file(o.m_file),
+	m_file(o.m_file)
-	m_bs(o.m_bs)
 {
 	o.m_file = nullptr;
 }
@@ -48,11 +49,11 @@ file::open(const wchar_t *path)
 		m_file->open(&fh, path, uefi::file_mode::read, uefi::file_attr::none),
 		L"Could not open relative path to file");
 
-	return file(fh, m_bs);
+	return file(fh);
 }
 
 buffer
-file::load(uefi::memory_type mem_type)
+file::load()
 {
 	uint8_t info_buf[sizeof(uefi::protos::file_info) + 100];
 	size_t size = sizeof(info_buf);
@@ -66,19 +67,14 @@ file::load(uefi::memory_type mem_type)
 		reinterpret_cast<uefi::protos::file_info*>(&info_buf);
 
 	size_t pages = memory::bytes_to_pages(info->file_size);
-	void *data = nullptr;
+	void *data = g_alloc.allocate_pages(pages, alloc_type::file);
-	try_or_raise(
-		m_bs->allocate_pages(
-			uefi::allocate_type::any_pages,
-			mem_type, pages, &data),
-		L"Could not allocate pages to load file");
 
 	size = info->file_size;
 	try_or_raise(
 		m_file->read(&size, data),
 		L"Could not read from file");
 
-	return { .size = size,  .data = data };
+	return { .pointer = data, .count = size };
 }
 
 file
@@ -106,7 +102,7 @@ get_boot_volume(uefi::handle image, uefi::boot_services *bs)
 		fs->open_volume(&f),
 		L"Could not open the boot volume");
 
-	return file(f, bs);
+	return file(f);
 }
 
 } // namespace fs

src/boot/fs.h

@@ -3,9 +3,13 @@
 #pragma once
 
 #include <uefi/types.h>
-#include <uefi/boot_services.h>
+#include "counted.h"
-#include <uefi/protos/file.h>
+
-#include "types.h"
+namespace uefi {
+	struct boot_services;
+namespace protos {
+	struct file;
+}}
 
 namespace boot {
 namespace fs {
@@ -23,15 +27,14 @@ public:
 	file open(const wchar_t *path);
 
 	/// Load the contents of this file into memory.
-	/// \arg mem_type  The UEFI memory type to use for allocation
 	/// \returns       A buffer describing the loaded memory. The
 	///                memory will be page-aligned.
-	buffer load(uefi::memory_type mem_type = uefi::memory_type::loader_data);
+	buffer load();
 
 private:
 	friend file get_boot_volume(uefi::handle, uefi::boot_services*);
 
-	file(uefi::protos::file *f, uefi::boot_services *bs);
+	file(uefi::protos::file *f);
 
 	uefi::protos::file *m_file;
 	uefi::boot_services *m_bs;

src/boot/loader.cpp

@@ -1,13 +1,15 @@
 #include <uefi/boot_services.h>
 #include <uefi/types.h>
 
-#include "loader.h"
+#include "allocator.h"
 #include "console.h"
 #include "elf.h"
 #include "error.h"
 #include "fs.h"
+#include "loader.h"
 #include "memory.h"
 #include "paging.h"
+#include "pointer_manipulation.h"
 #include "status.h"
 
 namespace init = kernel::init;
@@ -15,17 +17,18 @@ namespace init = kernel::init;
 namespace boot {
 namespace loader {
 
+using memory::alloc_type;
+
 buffer
 load_file(
 	fs::file &disk,
 	const wchar_t *name,
-	const wchar_t *path,
+	const wchar_t *path)
-	uefi::memory_type type)
 {
 	status_line status(L"Loading file", name);
 
 	fs::file file = disk.open(path);
-	buffer b = file.load(type);
+	buffer b = file.load();
 
 	//console::print(L"    Loaded at: 0x%lx, %d bytes\r\n", b.data, b.size);
 	return b;
@@ -51,59 +54,55 @@ void
 load_program(
 	init::program &program,
 	const wchar_t *name,
-	buffer data,
+	buffer data)
-	uefi::boot_services *bs)
 {
 	status_line status(L"Loading program:", name);
-	const elf::header *header = reinterpret_cast<const elf::header*>(data.data);
+	const elf::header *header = reinterpret_cast<const elf::header*>(data.pointer);
 
-	if (data.size < sizeof(elf::header) || !is_elfheader_valid(header))
+	if (data.count < sizeof(elf::header) || !is_elfheader_valid(header))
 		error::raise(uefi::status::load_error, L"ELF file not valid");
 
+	size_t num_sections = 0;
 	uintptr_t prog_base = uintptr_t(-1);
 	uintptr_t prog_end = 0;
 
 	for (int i = 0; i < header->ph_num; ++i) {
 		ptrdiff_t offset = header->ph_offset + i * header->ph_entsize;
 		const elf::program_header *pheader =
-			offset_ptr<elf::program_header>(data.data, offset);
+			offset_ptr<elf::program_header>(data.pointer, offset);
 
 		if (pheader->type != elf::PT_LOAD)
 			continue;
 
+		++num_sections;
 		uintptr_t end = pheader->vaddr + pheader->mem_size;
 		if (pheader->vaddr < prog_base) prog_base = pheader->vaddr;
 		if (end > prog_end) prog_end = end;
 	}
 
+	init::program_section *sections = new init::program_section [num_sections];
+	program.sections = { .pointer = sections, .count = num_sections };
+
 	size_t total_size = prog_end - prog_base;
 	size_t num_pages = memory::bytes_to_pages(total_size);
-	void *pages = nullptr;
+	void *pages = g_alloc.allocate_pages(num_pages, alloc_type::program, true);
+	program.phys_base = reinterpret_cast<uintptr_t>(pages);
 
-	try_or_raise(
+	size_t next_section = 0;
-		bs->allocate_pages(uefi::allocate_type::any_pages,
-			uefi::memory_type::loader_data, num_pages, &pages),
-		L"Failed allocating space for program");
-
-	bs->set_mem(pages, total_size, 0);
-
-	program.base = prog_base;
-	program.total_size = total_size;
-	program.num_sections = 0;
 	for (int i = 0; i < header->ph_num; ++i) {
 		ptrdiff_t offset = header->ph_offset + i * header->ph_entsize;
 		const elf::program_header *pheader =
-			offset_ptr<elf::program_header>(data.data, offset);
+			offset_ptr<elf::program_header>(data.pointer, offset);
 
 		if (pheader->type != elf::PT_LOAD)
 			continue;
 
-		init::program_section &section = program.sections[program.num_sections++];
+		init::program_section &section = program.sections[next_section++];
 
-		void *src_start = offset_ptr<void>(data.data, pheader->offset);
+		void *src_start = offset_ptr<void>(data.pointer, pheader->offset);
 		void *dest_start = offset_ptr<void>(pages, pheader->vaddr - prog_base);
 
-		bs->copy_mem(dest_start, src_start, pheader->file_size);
+		g_alloc.copy(dest_start, src_start, pheader->file_size);
 		section.phys_addr = reinterpret_cast<uintptr_t>(dest_start);
 		section.virt_addr = pheader->vaddr;
 		section.size = pheader->mem_size;
@@ -114,9 +113,7 @@ load_program(
 }
 
 void
-verify_kernel_header(
+verify_kernel_header(init::program &program)
-	init::program &program,
-	uefi::boot_services *bs)
 {
 	status_line status(L"Verifying kernel header");
 
@@ -132,7 +129,7 @@ verify_kernel_header(
     if (header->version < init::min_header_version)
         error::raise(uefi::status::unsupported, L"Kernel header version not supported");
 
-	console::print(L"    Loaded kernel vserion: %d.%d.%d %lx\r\n",
+	console::print(L"    Loaded kernel vserion: %d.%d.%d %x\r\n",
             header->version_major, header->version_minor, header->version_patch,
             header->version_gitsha);
 }

src/boot/loader.h

@@ -2,15 +2,18 @@
 /// Definitions for loading the kernel into memory
 #pragma once
 
-#include <uefi/boot_services.h>
+#include "counted.h"
 
-#include "kernel_args.h"
+namespace kernel {
-#include "memory.h"
+namespace init {
-#include "types.h"
+	struct program;
+}}
 
 namespace boot {
 
-namespace fs { class file; }
+namespace fs {
+	class file;
+}
 
 namespace loader {
 
@@ -18,33 +21,26 @@ namespace loader {
 /// \arg disk  The opened UEFI filesystem to load from
 /// \arg name  Name of the module (informational only)
 /// \arg path  Path on `disk` of the file to load
-/// \arg type  Memory type to use for allocation
 buffer
 load_file(
 	fs::file &disk,
 	const wchar_t *name,
-	const wchar_t *path,
+	const wchar_t *path);
-	uefi::memory_type type = uefi::memory_type::loader_data);
 
 /// Parse and load an ELF file in memory into a loaded image.
 /// \arg program  The program structure to fill
 /// \arg name     The name of the program being loaded
 /// \arg data     Buffer of the ELF file in memory
-/// \arg bs       Boot services
 void
 load_program(
 	kernel::init::program &program,
 	const wchar_t *name,
-	buffer data,
+	buffer data);
-	uefi::boot_services *bs);
 
 /// Verify that a loaded ELF has the j6 kernel header
 /// \arg program  The program to check for a header
-/// \arg bs       Boot services
 void
-verify_kernel_header(
+verify_kernel_header(kernel::init::program &program);
-	kernel::init::program &program,
-	uefi::boot_services *bs);
 
 } // namespace loader
 } // namespace boot

src/boot/main.cpp

@@ -7,6 +7,7 @@
 #include <stddef.h>
 #include <stdint.h>
 
+#include "allocator.h"
 #include "console.h"
 #include "cpu/cpu_id.h"
 #include "error.h"
@@ -14,6 +15,7 @@
 #include "hardware.h"
 #include "loader.h"
 #include "memory.h"
+#include "memory_map.h"
 #include "paging.h"
 #include "status.h"
 
@@ -50,49 +52,14 @@ void change_pointer(T *&pointer)
 	pointer = offset_ptr<T>(pointer, kernel::memory::page_offset);
 }
 
-/// Allocate space for kernel args. Allocates enough space from pool
-/// memory for the args header and the module and program headers.
-init::args *
-allocate_args_structure(
-	uefi::boot_services *bs,
-	size_t max_modules,
-	size_t max_programs)
-{
-	status_line status {L"Setting up kernel args memory"};
-
-	init::args *args = nullptr;
-
-	size_t args_size =
-		sizeof(init::args) + // The header itself
-		max_modules * sizeof(init::module) +  // The module structures
-		max_programs * sizeof(init::program); // The program structures
-
-	try_or_raise(
-		bs->allocate_pool(uefi::memory_type::loader_data, args_size,
-			reinterpret_cast<void**>(&args)),
-		L"Could not allocate argument memory");
-
-	bs->set_mem(args, args_size, 0);
-
-	args->modules =
-		reinterpret_cast<init::module*>(args + 1);
-	args->num_modules = 0;
-
-	args->programs =
-		reinterpret_cast<init::program*>(args->modules + max_modules);
-	args->num_programs = 0;
-
-	return args;
-}
-
 /// Add a module to the kernel args list
 inline void
 add_module(init::args *args, init::mod_type type, buffer &data)
 {
-	init::module &m = args->modules[args->num_modules++];
+	init::module &m = args->modules[args->modules.count++];
 	m.type = type;
-	m.location = data.data;
+	m.location = data.pointer;
-	m.size = data.size;
+	m.size = data.count;
 
 	change_pointer(m.location);
 }
@@ -121,7 +88,7 @@ check_cpu_supported()
 
 /// The main procedure for the portion of the loader that runs while
 /// UEFI is still in control of the machine. (ie, while the loader still
-/// has access to boot services.
+/// has access to boot services.)
 init::args *
 uefi_preboot(uefi::handle image, uefi::system_table *st)
 {
@@ -129,33 +96,36 @@ uefi_preboot(uefi::handle image, uefi::system_table *st)
 
 	uefi::boot_services *bs = st->boot_services;
 	uefi::runtime_services *rs = st->runtime_services;
+
+	memory::init_allocator(bs);
 	memory::init_pointer_fixup(bs, rs);
 
-	init::args *args =
+	init::args *args = new init::args;
-		allocate_args_structure(bs, max_modules, max_programs);
+	g_alloc.zero(args, sizeof(init::args));
+	args->programs.pointer = new init::program[5];
+	args->modules.pointer = new init::module[5];
 
 	args->magic = init::args_magic;
 	args->version = init::args_version;
 	args->runtime_services = rs;
 	args->acpi_table = hw::find_acpi_table(st);
-	paging::allocate_tables(args, bs);
-
 	memory::mark_pointer_fixup(&args->runtime_services);
 
+	paging::allocate_tables(args);
+
 	fs::file disk = fs::get_boot_volume(image, bs);
 
-	buffer symbols = loader::load_file(disk, L"symbol table", L"symbol_table.dat",
+	buffer symbols = loader::load_file(disk, L"symbol table", L"symbol_table.dat");
-				uefi::memory_type::loader_data);
 	add_module(args, init::mod_type::symbol_table, symbols);
 
 	for (auto &desc : program_list) {
 		buffer buf = loader::load_file(disk, desc.name, desc.path);
-		init::program &program = args->programs[args->num_programs++];
+		init::program &program = args->programs[args->programs.count++];
-		loader::load_program(program, desc.name, buf, bs);
+		loader::load_program(program, desc.name, buf);
 	}
 
     // First program *must* be the kernel
-    loader::verify_kernel_header(args->programs[0], bs);
+    loader::verify_kernel_header(args->programs[0]);
 
 	return args;
 }
@@ -165,9 +135,14 @@ uefi_exit(init::args *args, uefi::handle image, uefi::boot_services *bs)
 {
 	status_line status {L"Exiting UEFI", nullptr, false};
 
-	memory::efi_mem_map map =
+	memory::efi_mem_map map;
-		memory::build_kernel_mem_map(args, bs);
+	map.update(*bs);
 
+	args->mem_map = memory::build_kernel_map(map);
+	args->frame_blocks = memory::build_frame_blocks(args->mem_map);
+	args->allocations = g_alloc.get_register();
+
+	map.update(*bs);
 	try_or_raise(
 		bs->exit_boot_services(image, map.key),
 		L"Failed to exit boot services");
@@ -194,7 +169,6 @@ efi_main(uefi::handle image, uefi::system_table *st)
 	// Map the kernel to the appropriate address
 	init::program &kernel = args->programs[0];
 	for (auto &section : kernel.sections)
-		if (section.size)
 		paging::map_section(args, section);
 
 	memory::fix_frame_blocks(args);
@@ -209,8 +183,10 @@ efi_main(uefi::handle image, uefi::system_table *st)
 
 	change_pointer(args);
 	change_pointer(args->pml4);
-	change_pointer(args->modules);
+	change_pointer(args->modules.pointer);
-	change_pointer(args->programs);
+	change_pointer(args->programs.pointer);
+	for (auto &program : args->programs)
+		change_pointer(program.sections.pointer);
 
 	status.next();
 

src/boot/memory.cpp

@@ -1,67 +1,24 @@
 #include <stddef.h>
-#include <uefi/types.h>
 
-#include "kernel_memory.h"
+#include <uefi/boot_services.h>
+#include <uefi/runtime_services.h>
+#include <uefi/types.h>
+#include <kutil/no_construct.h>
 
 #include "console.h"
 #include "error.h"
+#include "kernel_memory.h"
 #include "memory.h"
+#include "memory_map.h"
 #include "paging.h"
 #include "status.h"
 
 namespace boot {
 namespace memory {
 
-using mem_entry = kernel::init::mem_entry;
-using mem_type = kernel::init::mem_type;
-using frame_block = kernel::init::frame_block;
-using kernel::init::frames_per_block;
-
 size_t fixup_pointer_index = 0;
 void **fixup_pointers[64];
 
-static const wchar_t *memory_type_names[] = {
-	L"reserved memory type",
-	L"loader code",
-	L"loader data",
-	L"boot services code",
-	L"boot services data",
-	L"runtime services code",
-	L"runtime services data",
-	L"conventional memory",
-	L"unusable memory",
-	L"acpi reclaim memory",
-	L"acpi memory nvs",
-	L"memory mapped io",
-	L"memory mapped io port space",
-	L"pal code",
-	L"persistent memory"
-};
-
-static const wchar_t *kernel_memory_type_names[] = {
-	L"free",
-	L"pending",
-	L"acpi",
-	L"uefi_runtime",
-	L"mmio",
-	L"persistent"
-};
-
-static const wchar_t *
-memory_type_name(uefi::memory_type t)
-{
-	if (t < uefi::memory_type::max_memory_type)
-		return memory_type_names[static_cast<uint32_t>(t)];
-
-	return L"Bad Type Value";
-}
-
-static const wchar_t *
-kernel_memory_type_name(kernel::init::mem_type t)
-{
-	return kernel_memory_type_names[static_cast<uint32_t>(t)];
-}
-
 void
 update_marked_addresses(uefi::event, void *context)
 {
@@ -99,275 +56,13 @@ mark_pointer_fixup(void **p)
 	fixup_pointers[fixup_pointer_index++] = p;
 }
 
-bool
-can_merge(mem_entry &prev, mem_type type, uefi::memory_descriptor *next)
-{
-	return
-		prev.type == type &&
-		prev.start + (page_size * prev.pages) == next->physical_start &&
-		prev.attr == (next->attribute & 0xffffffff);
-}
-
-void
-get_uefi_mappings(efi_mem_map &map, uefi::boot_services *bs)
-{
-	size_t length = map.total;
-	uefi::status status = bs->get_memory_map(
-		&length, map.entries, &map.key, &map.size, &map.version);
-	map.length = length;
-
-	if (status == uefi::status::success)
-		return;
-
-	if (status != uefi::status::buffer_too_small)
-		error::raise(status, L"Error getting memory map size");
-
-	if (map.entries) {
-		try_or_raise(
-			bs->free_pool(reinterpret_cast<void*>(map.entries)),
-			L"Freeing previous memory map space");
-	}
-
-	map.total = length + 10*map.size;
-
-	try_or_raise(
-		bs->allocate_pool(
-			uefi::memory_type::loader_data, map.total,
-			reinterpret_cast<void**>(&map.entries)),
-		L"Allocating space for memory map");
-
-	map.length = map.total;
-	try_or_raise(
-		bs->get_memory_map(&map.length, map.entries, &map.key, &map.size, &map.version),
-		L"Getting UEFI memory map");
-}
-
-inline size_t bitmap_size(size_t frames) { return (frames + 63) / 64; }
-inline size_t num_blocks(size_t frames) { return (frames + (frames_per_block-1)) / frames_per_block; }
-
-void
-build_kernel_frame_blocks(const mem_entry *map, size_t nent, kernel::init::args *args, uefi::boot_services *bs)
-{
-	status_line status {L"Creating kernel frame accounting map"};
-
-	size_t block_count = 0;
-	size_t total_bitmap_size = 0;
-	for (size_t i = 0; i < nent; ++i) {
-		const mem_entry &ent = map[i];
-		if (ent.type != mem_type::free)
-			continue;
-
-		block_count += num_blocks(ent.pages);
-		total_bitmap_size += bitmap_size(ent.pages) * sizeof(uint64_t);
-	}
-
-	size_t total_size = block_count * sizeof(frame_block) + total_bitmap_size;
-
-	frame_block *blocks = nullptr;
-	try_or_raise(
-		bs->allocate_pages(
-			uefi::allocate_type::any_pages,
-			uefi::memory_type::loader_data,
-			bytes_to_pages(total_size),
-			reinterpret_cast<void**>(&blocks)),
-		L"Error allocating kernel frame block space");
-
-	frame_block *next_block = blocks;
-	for (size_t i = 0; i < nent; ++i) {
-		const mem_entry &ent = map[i];
-		if (ent.type != mem_type::free)
-			continue;
-
-		size_t page_count = ent.pages;
-		uintptr_t base_addr = ent.start;
-		while (page_count) {
-			frame_block *blk = next_block++;
-			bs->set_mem(blk, sizeof(frame_block), 0);
-
-			blk->attrs = ent.attr;
-			blk->base = base_addr;
-			base_addr += frames_per_block * page_size;
-
-			if (page_count >= frames_per_block) {
-				page_count -= frames_per_block;
-				blk->count = frames_per_block;
-				blk->map1 = ~0ull;
-				bs->set_mem(blk->map2, sizeof(blk->map2), 0xff);
-			} else {
-				blk->count = page_count;
-				unsigned i = 0;
-
-				uint64_t b1 = (page_count + 4095) / 4096;
-				blk->map1 = (1 << b1) - 1;
-
-				uint64_t b2 = (page_count + 63) / 64;
-				uint64_t b2q = b2 / 64;
-				uint64_t b2r = b2 % 64;
-				bs->set_mem(blk->map2, b2q, 0xff);
-				blk->map2[b2q] = (1 << b2r) - 1;
-				break;
-			}
-		}
-	}
-
-	uint64_t *bitmap = reinterpret_cast<uint64_t*>(next_block);
-	bs->set_mem(bitmap, total_bitmap_size, 0);
-	for (unsigned i = 0; i < block_count; ++i) {
-		frame_block &blk = blocks[i];
-		blk.bitmap = bitmap;
-
-		size_t b = blk.count / 64;
-		size_t r = blk.count % 64;
-		bs->set_mem(blk.bitmap, b*8, 0xff);
-		blk.bitmap[b] = (1 << r) - 1;
-
-		bitmap += bitmap_size(blk.count);
-	}
-
-	args->frame_block_count = block_count;
-	args->frame_block_pages = bytes_to_pages(total_size);
-	args->frame_blocks = blocks;
-}
-
-void
-fix_frame_blocks(kernel::init::args *args)
-{
-	// Map the frame blocks to the appropriate address
-	paging::map_pages(args,
-		reinterpret_cast<uintptr_t>(args->frame_blocks),
-		::memory::bitmap_start,
-		args->frame_block_pages,
-		true, false);
-
-	uintptr_t offset = ::memory::bitmap_start -
-		reinterpret_cast<uintptr_t>(args->frame_blocks);
-
-	for (unsigned i = 0; i < args->frame_block_count; ++i) {
-		frame_block &blk = args->frame_blocks[i];
-		blk.bitmap = reinterpret_cast<uint64_t*>(
-			reinterpret_cast<uintptr_t>(blk.bitmap) + offset);
-	}
-}
-
-efi_mem_map
-build_kernel_mem_map(kernel::init::args *args, uefi::boot_services *bs)
-{
-	status_line status {L"Creating kernel memory map"};
-
-	efi_mem_map map;
-	get_uefi_mappings(map, bs);
-
-	size_t map_size = map.num_entries() * sizeof(mem_entry);
-
-	mem_entry *kernel_map = nullptr;
-	try_or_raise(
-		bs->allocate_pages(
-			uefi::allocate_type::any_pages,
-			uefi::memory_type::loader_data,
-			bytes_to_pages(map_size),
-			reinterpret_cast<void**>(&kernel_map)),
-		L"Error allocating kernel memory map module space");
-
-	bs->set_mem(kernel_map, map_size, 0);
-	get_uefi_mappings(map, bs);
-
-	size_t nent = 0;
-	bool first = true;
-	for (auto desc : map) {
-		/*
-		// EFI map dump
-		console::print(L"   eRange %lx (%lx) %x(%s) [%lu]\r\n",
-			desc->physical_start, desc->attribute, desc->type, memory_type_name(desc->type), desc->number_of_pages);
-		*/
-
-		mem_type type;
-		switch (desc->type) {
-			case uefi::memory_type::reserved:
-			case uefi::memory_type::unusable_memory:
-			case uefi::memory_type::acpi_memory_nvs:
-			case uefi::memory_type::pal_code:
-				continue;
-
-			case uefi::memory_type::loader_code:
-			case uefi::memory_type::boot_services_code:
-			case uefi::memory_type::boot_services_data:
-			case uefi::memory_type::conventional_memory:
-			case uefi::memory_type::loader_data:
-				type = mem_type::free;
-				break;
-
-			case uefi::memory_type::runtime_services_code:
-			case uefi::memory_type::runtime_services_data:
-				type = mem_type::uefi_runtime;
-				break;
-
-			case uefi::memory_type::acpi_reclaim_memory:
-				type = mem_type::acpi;
-				break;
-
-			case uefi::memory_type::memory_mapped_io:
-			case uefi::memory_type::memory_mapped_io_port_space:
-				type = mem_type::mmio;
-				break;
-
-			case uefi::memory_type::persistent_memory:
-				type = mem_type::persistent;
-				break;
-
-			default:
-				error::raise(
-					uefi::status::invalid_parameter,
-					L"Got an unexpected memory type from UEFI memory map");
-		}
-
-		// TODO: validate uefi's map is sorted
-		if (first) {
-			first = false;
-			mem_entry &ent = kernel_map[nent++];
-			ent.start = desc->physical_start;
-			ent.pages = desc->number_of_pages;
-			ent.type = type;
-			ent.attr = (desc->attribute & 0xffffffff);
-			continue;
-		}
-
-		mem_entry &prev = kernel_map[nent - 1];
-		if (can_merge(prev, type, desc)) {
-			prev.pages += desc->number_of_pages;
-		} else {
-			mem_entry &next = kernel_map[nent++];
-			next.start = desc->physical_start;
-			next.pages = desc->number_of_pages;
-			next.type = type;
-			next.attr = (desc->attribute & 0xffffffff);
-		}
-	}
-
-	// Give just the actually-set entries in the header
-	args->mem_map = kernel_map;
-	args->map_count = nent;
-
-	/*
-	// kernel map dump
-	for (unsigned i = 0; i < nent; ++i) {
-		const mem_entry &e = kernel_map[i];
-		console::print(L"   kRange %lx (%lx) %x(%s) [%lu]\r\n",
-			e.start, e.attr, e.type, kernel_memory_type_name(e.type), e.pages);
-	}
-	*/
-
-	build_kernel_frame_blocks(kernel_map, nent, args, bs);
-	get_uefi_mappings(map, bs);
-	return map;
-}
-
 void
 virtualize(void *pml4, efi_mem_map &map, uefi::runtime_services *rs)
 {
 	paging::add_current_mappings(reinterpret_cast<paging::page_table*>(pml4));
 
-	for (auto desc : map)
+	for (auto &desc : map)
-		desc->virtual_start = desc->physical_start + ::memory::page_offset;
+		desc.virtual_start = desc.physical_start + ::memory::page_offset;
 
 	// Write our new PML4 pointer to CR3
 	asm volatile ( "mov %0, %%cr3" :: "r" (pml4) );

src/boot/memory.h

@@ -1,15 +1,19 @@
+#pragma once
 /// \file memory.h
 /// Memory-related constants and functions.
-#pragma once
+
-#include <uefi/boot_services.h>
-#include <uefi/runtime_services.h>
 #include <stdint.h>
-#include "kernel_args.h"
+
-#include "pointer_manipulation.h"
+namespace uefi {
+	struct boot_services;
+	struct runtime_services;
+}
 
 namespace boot {
 namespace memory {
 
+class efi_mem_map;
+
 /// UEFI specifies that pages are always 4 KiB.
 constexpr size_t page_size = 0x1000;
 
@@ -33,44 +37,6 @@ void mark_pointer_fixup(void **p);
 
 /// @}
 
-/// Struct that represents UEFI's memory map. Contains a pointer to the map data
-/// as well as the data on how to read it.
-struct efi_mem_map
-{
-	using desc = uefi::memory_descriptor;
-	using iterator = offset_iterator<desc>;
-
-	size_t length;     ///< Total length of the map data
-	size_t total;      ///< Total allocated space for map data
-	size_t size;       ///< Size of an entry in the array
-	size_t key;        ///< Key for detecting changes
-	uint32_t version;  ///< Version of the `memory_descriptor` struct
-	desc *entries;     ///< The array of UEFI descriptors
-
-	efi_mem_map() : length(0), total(0), size(0), key(0), version(0), entries(nullptr) {}
-
-	/// Get the count of entries in the array
-	inline size_t num_entries() const { return length / size; }
-
-	/// Return an iterator to the beginning of the array
-	iterator begin() { return iterator(entries, size); }
-
-	/// Return an iterator to the end of the array
-	iterator end() { return offset_ptr<desc>(entries, length); }
-};
-
-/// Add the kernel's memory map as a module to the kernel args.
-/// \returns  The uefi memory map used to build the kernel map
-efi_mem_map build_kernel_mem_map(kernel::init::args *args, uefi::boot_services *bs);
-
-/// Create the kernel frame allocation maps
-void build_kernel_frame_blocks(
-		const kernel::init::mem_entry *map, size_t nent,
-		kernel::init::args *args, uefi::boot_services *bs);
-
-/// Map the frame allocation maps to the right spot and fix up pointers
-void fix_frame_blocks(kernel::init::args *args);
-
 /// Activate the given memory mappings. Sets the given page tables live as well
 /// as informs UEFI runtime services of the new mappings.
 /// \arg pml4  The root page table for the new mappings

src/boot/memory_map.cpp

@@ -0,0 +1,303 @@
+#include <uefi/boot_services.h>
+#include <uefi/types.h>
+
+#include "allocator.h"
+#include "error.h"
+#include "kernel_memory.h"
+#include "memory.h"
+#include "memory_map.h"
+#include "paging.h"
+#include "status.h"
+
+namespace boot {
+namespace memory {
+
+using kernel::init::frame_block;
+using kernel::init::frames_per_block;
+using kernel::init::mem_entry;
+using kernel::init::mem_type;
+
+
+void
+efi_mem_map::update(uefi::boot_services &bs)
+{
+	size_t l = total;
+	uefi::status status = bs.get_memory_map(
+		&l, entries, &key, &size, &version);
+	length = l;
+
+	if (status == uefi::status::success)
+		return;
+
+	if (status != uefi::status::buffer_too_small)
+		error::raise(status, L"Error getting memory map size");
+
+	if (entries) {
+		try_or_raise(
+			bs.free_pool(reinterpret_cast<void*>(entries)),
+			L"Freeing previous memory map space");
+	}
+
+	total = length + 10 * size;
+
+	try_or_raise(
+		bs.allocate_pool(
+			uefi::memory_type::loader_data, total,
+			reinterpret_cast<void**>(&entries)),
+		L"Allocating space for memory map");
+
+	length = total;
+	try_or_raise(
+		bs.get_memory_map(&length, entries, &key, &size, &version),
+		L"Getting UEFI memory map");
+}
+
+static const wchar_t *memory_type_names[] = {
+	L"reserved memory type",
+	L"loader code",
+	L"loader data",
+	L"boot services code",
+	L"boot services data",
+	L"runtime services code",
+	L"runtime services data",
+	L"conventional memory",
+	L"unusable memory",
+	L"acpi reclaim memory",
+	L"acpi memory nvs",
+	L"memory mapped io",
+	L"memory mapped io port space",
+	L"pal code",
+	L"persistent memory"
+};
+
+static const wchar_t *kernel_memory_type_names[] = {
+	L"free",
+	L"pending",
+	L"acpi",
+	L"uefi_runtime",
+	L"mmio",
+	L"persistent"
+};
+
+static const wchar_t *
+memory_type_name(uefi::memory_type t)
+{
+	if (t < uefi::memory_type::max_memory_type)
+		return memory_type_names[static_cast<uint32_t>(t)];
+
+	return L"Bad Type Value";
+}
+
+static const wchar_t *
+kernel_memory_type_name(kernel::init::mem_type t)
+{
+	return kernel_memory_type_names[static_cast<uint32_t>(t)];
+}
+
+inline bool
+can_merge(mem_entry &prev, mem_type type, uefi::memory_descriptor &next)
+{
+	return
+		prev.type == type &&
+		prev.start + (page_size * prev.pages) == next.physical_start &&
+		prev.attr == (next.attribute & 0xffffffff);
+}
+
+counted<mem_entry>
+build_kernel_map(efi_mem_map &map)
+{
+	status_line status {L"Creating kernel memory map"};
+
+	size_t map_size = map.num_entries() * sizeof(mem_entry);
+	size_t num_pages = bytes_to_pages(map_size);
+	mem_entry *kernel_map = reinterpret_cast<mem_entry*>(
+		g_alloc.allocate_pages(num_pages, alloc_type::mem_map, true));
+
+	size_t nent = 0;
+	bool first = true;
+	for (auto &desc : map) {
+		/*
+		// EFI map dump
+		console::print(L"   eRange %lx (%lx) %x(%s) [%lu]\r\n",
+			desc.physical_start, desc.attribute, desc.type, memory_type_name(desc.type), desc.number_of_pages);
+		*/
+
+		mem_type type;
+		switch (desc.type) {
+			case uefi::memory_type::reserved:
+			case uefi::memory_type::unusable_memory:
+			case uefi::memory_type::acpi_memory_nvs:
+			case uefi::memory_type::pal_code:
+				continue;
+
+			case uefi::memory_type::loader_code:
+			case uefi::memory_type::boot_services_code:
+			case uefi::memory_type::boot_services_data:
+			case uefi::memory_type::conventional_memory:
+			case uefi::memory_type::loader_data:
+				type = mem_type::free;
+				break;
+
+			case uefi::memory_type::runtime_services_code:
+			case uefi::memory_type::runtime_services_data:
+				type = mem_type::uefi_runtime;
+				break;
+
+			case uefi::memory_type::acpi_reclaim_memory:
+				type = mem_type::acpi;
+				break;
+
+			case uefi::memory_type::memory_mapped_io:
+			case uefi::memory_type::memory_mapped_io_port_space:
+				type = mem_type::mmio;
+				break;
+
+			case uefi::memory_type::persistent_memory:
+				type = mem_type::persistent;
+				break;
+
+			default:
+				error::raise(
+					uefi::status::invalid_parameter,
+					L"Got an unexpected memory type from UEFI memory map");
+		}
+
+		// TODO: validate uefi's map is sorted
+		if (first) {
+			first = false;
+			mem_entry &ent = kernel_map[nent++];
+			ent.start = desc.physical_start;
+			ent.pages = desc.number_of_pages;
+			ent.type = type;
+			ent.attr = (desc.attribute & 0xffffffff);
+			continue;
+		}
+
+		mem_entry &prev = kernel_map[nent - 1];
+		if (can_merge(prev, type, desc)) {
+			prev.pages += desc.number_of_pages;
+		} else {
+			mem_entry &next = kernel_map[nent++];
+			next.start = desc.physical_start;
+			next.pages = desc.number_of_pages;
+			next.type = type;
+			next.attr = (desc.attribute & 0xffffffff);
+		}
+	}
+
+	/*
+	// kernel map dump
+	for (unsigned i = 0; i < nent; ++i) {
+		const mem_entry &e = kernel_map[i];
+		console::print(L"   kRange %lx (%lx) %x(%s) [%lu]\r\n",
+			e.start, e.attr, e.type, kernel_memory_type_name(e.type), e.pages);
+	}
+	*/
+
+	return { .pointer = kernel_map, .count = nent };
+}
+
+inline size_t bitmap_size(size_t frames) { return (frames + 63) / 64; }
+inline size_t num_blocks(size_t frames) { return (frames + (frames_per_block-1)) / frames_per_block; }
+
+counted<kernel::init::frame_block>
+build_frame_blocks(const counted<kernel::init::mem_entry> &kmap)
+{
+	status_line status {L"Creating kernel frame accounting map"};
+
+	size_t block_count = 0;
+	size_t total_bitmap_size = 0;
+	for (size_t i = 0; i < kmap.count; ++i) {
+		const mem_entry &ent = kmap[i];
+		if (ent.type != mem_type::free)
+			continue;
+
+		block_count += num_blocks(ent.pages);
+		total_bitmap_size += bitmap_size(ent.pages) * sizeof(uint64_t);
+	}
+
+	size_t total_size = block_count * sizeof(frame_block) + total_bitmap_size;
+
+	frame_block *blocks = reinterpret_cast<frame_block*>(
+		g_alloc.allocate_pages(bytes_to_pages(total_size), alloc_type::frame_map, true));
+
+	frame_block *next_block = blocks;
+	for (size_t i = 0; i < kmap.count; ++i) {
+		const mem_entry &ent = kmap[i];
+		if (ent.type != mem_type::free)
+			continue;
+
+		size_t page_count = ent.pages;
+		uintptr_t base_addr = ent.start;
+		while (page_count) {
+			frame_block *blk = next_block++;
+
+			blk->flags = static_cast<kernel::init::frame_flags>(ent.attr);
+			blk->base = base_addr;
+			base_addr += frames_per_block * page_size;
+
+			if (page_count >= frames_per_block) {
+				page_count -= frames_per_block;
+				blk->count = frames_per_block;
+				blk->map1 = ~0ull;
+				g_alloc.memset(blk->map2, sizeof(blk->map2), 0xff);
+			} else {
+				blk->count = page_count;
+				unsigned i = 0;
+
+				uint64_t b1 = (page_count + 4095) / 4096;
+				blk->map1 = (1 << b1) - 1;
+
+				uint64_t b2 = (page_count + 63) / 64;
+				uint64_t b2q = b2 / 64;
+				uint64_t b2r = b2 % 64;
+				g_alloc.memset(blk->map2, b2q, 0xff);
+				blk->map2[b2q] = (1 << b2r) - 1;
+				break;
+			}
+		}
+	}
+
+	uint64_t *bitmap = reinterpret_cast<uint64_t*>(next_block);
+	for (unsigned i = 0; i < block_count; ++i) {
+		frame_block &blk = blocks[i];
+		blk.bitmap = bitmap;
+
+		size_t b = blk.count / 64;
+		size_t r = blk.count % 64;
+		g_alloc.memset(blk.bitmap, b*8, 0xff);
+		blk.bitmap[b] = (1 << r) - 1;
+
+		bitmap += bitmap_size(blk.count);
+	}
+
+	return { .pointer = blocks, .count = block_count };
+}
+
+void
+fix_frame_blocks(kernel::init::args *args)
+{
+	counted<frame_block> &blocks = args->frame_blocks;
+
+	size_t size = blocks.count * sizeof(frame_block);
+	for (unsigned i = 0; i < blocks.count; ++i)
+		size += bitmap_size(blocks[i].count) * sizeof(uint64_t);
+
+	size_t pages = bytes_to_pages(size);
+	uintptr_t addr = reinterpret_cast<uintptr_t>(blocks.pointer);
+
+	// Map the frame blocks to the appropriate address
+	paging::map_pages(args, addr,
+		::memory::bitmap_start, pages, true, false);
+
+	uintptr_t offset = ::memory::bitmap_start - addr;
+
+	for (unsigned i = 0; i < blocks.count; ++i) {
+		frame_block &blk = blocks[i];
+		blk.bitmap = offset_ptr<uint64_t>(blk.bitmap, offset);
+	}
+}
+
+
+} // namespace memory
+} // namespace boot

src/boot/memory_map.h

@@ -0,0 +1,63 @@
+#pragma once
+/// \file memory_map.h
+/// Memory-map related types and functions
+
+#include "counted.h"
+#include "pointer_manipulation.h"
+
+namespace uefi {
+	struct boot_services;
+	struct memory_descriptor;
+}
+
+namespace kernel {
+namespace init {
+	struct args;
+	struct frame_block;
+	struct mem_entry;
+}}
+
+namespace boot {
+namespace memory {
+
+/// Struct that represents UEFI's memory map. Contains a pointer to the map data
+/// as well as the data on how to read it.
+struct efi_mem_map
+{
+	using desc = uefi::memory_descriptor;
+	using iterator = offset_iterator<desc>;
+
+	size_t length;     ///< Total length of the map data
+	size_t total;      ///< Total allocated space for map data
+	size_t size;       ///< Size of an entry in the array
+	size_t key;        ///< Key for detecting changes
+	uint32_t version;  ///< Version of the `memory_descriptor` struct
+	desc *entries;     ///< The array of UEFI descriptors
+
+	efi_mem_map() : length(0), total(0), size(0), key(0), version(0), entries(nullptr) {}
+
+	/// Update the map from UEFI
+	void update(uefi::boot_services &bs);
+
+	/// Get the count of entries in the array
+	inline size_t num_entries() const { return length / size; }
+
+	/// Return an iterator to the beginning of the array
+	inline iterator begin() { return iterator(entries, size); }
+
+	/// Return an iterator to the end of the array
+	inline iterator end() { return offset_ptr<desc>(entries, length); }
+};
+
+/// Add the kernel's memory map as a module to the kernel args.
+/// \returns  The uefi memory map used to build the kernel map
+counted<kernel::init::mem_entry> build_kernel_map(efi_mem_map &map);
+
+/// Create the kernel frame allocation maps
+counted<kernel::init::frame_block> build_frame_blocks(const counted<kernel::init::mem_entry> &kmap);
+
+/// Map the frame allocation maps to the right spot and fix up pointers
+void fix_frame_blocks(kernel::init::args *args);
+
+} // namespace boot
+} // namespace memory

src/boot/module.toml

@@ -2,15 +2,17 @@ name = "boot"
 kind = "exe"
 output = "boot.efi"
 targets = ["boot"]
-deps = ["cpu"]
+deps = ["cpu", "kutil"]
 sources = [
-	"main.cpp",
+	"allocator.cpp",
 	"console.cpp",
 	"error.cpp",
 	"fs.cpp",
 	"hardware.cpp",
 	"loader.cpp",
+	"main.cpp",
 	"memory.cpp",
+	"memory_map.cpp",
 	"paging.cpp",
 	"status.cpp",
 	"support.cpp",

src/boot/paging.cpp

@@ -1,5 +1,6 @@
 #include "kernel_memory.h"
 
+#include "allocator.h"
 #include "console.h"
 #include "error.h"
 #include "loader.h"
@@ -11,6 +12,7 @@
 namespace boot {
 namespace paging {
 
+using memory::alloc_type;
 using memory::page_size;
 using ::memory::pml4e_kernel;
 using ::memory::table_entries;
@@ -54,6 +56,19 @@ constexpr uint64_t huge_page_flags = 0x18b;
 /// Page table entry flags for entries pointing at another table
 constexpr uint64_t table_flags = 0x003;
 
+
+inline void *
+pop_pages(counted<void> &pages, size_t count)
+{
+	if (count > pages.count)
+		error::raise(uefi::status::out_of_resources, L"Page table cache empty", 0x7ab1e5);
+
+	void *next = pages.pointer;
+	pages.pointer = offset_ptr<void>(pages.pointer, count*page_size);
+	pages.count -= count;
+	return next;
+}
+
 /// Iterator over page table entries.
 template <unsigned D = 4>
 class page_entry_iterator
@@ -62,15 +77,12 @@ public:
 	/// Constructor.
 	/// \arg virt    Virtual address this iterator is starting at
 	/// \arg pml4    Root of the page tables to iterate
-	/// \arg page_cache Pointer to pages that can be used for page tables
+	/// \arg pages   Cache of usable table pages
-	/// \arg page_count Number of pages pointed to by `page_cache`
 	page_entry_iterator(
 			uintptr_t virt,
 			page_table *pml4,
-			void *&page_cache,
+			counted<void> &pages) :
-			size_t &cache_count) :
+		m_pages(pages)
-		m_page_cache(page_cache),
-		m_cache_count(cache_count)
 	{
 		m_table[0] = pml4;
 		for (unsigned i = 0; i < D; ++i) {
@@ -117,12 +129,7 @@ private:
 		uint64_t & parent_ent = entry(level - 1);
 
 		if (!(parent_ent & 1)) {
-			if (!m_cache_count--)
+			page_table *table = reinterpret_cast<page_table*>(pop_pages(m_pages, 1));
-				error::raise(uefi::status::out_of_resources, L"Page table cache empty", 0x7ab1e5);
-
-			page_table *table = reinterpret_cast<page_table*>(m_page_cache);
-			m_page_cache = offset_ptr<void>(m_page_cache, page_size);
-
 			parent_ent = (reinterpret_cast<uintptr_t>(table) & ~0xfffull) | table_flags;
 			m_table[level] = table;
 		} else {
@@ -130,29 +137,25 @@ private:
 		}
 	}
 
-	void *&m_page_cache;
+	counted<void> &m_pages;
-	size_t &m_cache_count;
 	page_table *m_table[D];
 	uint16_t m_index[D];
 };
 
 
 static void
-add_offset_mappings(page_table *pml4, void *&page_cache, size_t &num_pages)
+add_offset_mappings(page_table *pml4, counted<void> &pages)
 {
 	uintptr_t phys = 0;
 	uintptr_t virt = ::memory::page_offset; // Start of offset-mapped area
-	size_t pages = 64 * 1024; // 64 TiB of 1 GiB pages
+	size_t page_count = 64 * 1024; // 64 TiB of 1 GiB pages
 	constexpr size_t GiB = 0x40000000ull;
 
-	page_entry_iterator<2> iterator{
+	page_entry_iterator<2> iterator{virt, pml4, pages};
-		virt, pml4,
-		page_cache,
-		num_pages};
 
 	while (true) {
 		*iterator = phys | huge_page_flags;
-		if (--pages == 0)
+		if (--page_count == 0)
 			break;
 
 		iterator.increment();
@@ -161,13 +164,10 @@ add_offset_mappings(page_table *pml4, void *&page_cache, size_t &num_pages)
 }
 
 static void
-add_kernel_pds(page_table *pml4, void *&page_cache, size_t &num_pages)
+add_kernel_pds(page_table *pml4, counted<void> &pages)
 {
-	for (unsigned i = pml4e_kernel; i < table_entries; ++i) {
+	for (unsigned i = pml4e_kernel; i < table_entries; ++i)
-		pml4->set(i, page_cache, table_flags);
+		pml4->set(i, pop_pages(pages, 1), table_flags);
-		page_cache = offset_ptr<void>(page_cache, page_size);
-		num_pages--;
-	}
 }
 
 void
@@ -186,7 +186,7 @@ add_current_mappings(page_table *new_pml4)
 }
 
 void
-allocate_tables(kernel::init::args *args, uefi::boot_services *bs)
+allocate_tables(kernel::init::args *args)
 {
 	status_line status(L"Allocating initial page tables");
 
@@ -198,28 +198,16 @@ allocate_tables(kernel::init::args *args, uefi::boot_services *bs)
 
 	static constexpr size_t tables_needed = kernel_tables + extra_tables;
 
-	void *addr = nullptr;
+	void *addr = g_alloc.allocate_pages(tables_needed, alloc_type::page_table, true);
-	try_or_raise(
-		bs->allocate_pages(
-			uefi::allocate_type::any_pages,
-			uefi::memory_type::loader_data,
-			tables_needed,
-			&addr),
-		L"Error allocating page table pages.");
-
-	bs->set_mem(addr, tables_needed*page_size, 0);
-
 	page_table *pml4 = reinterpret_cast<page_table*>(addr);
 
 	args->pml4 = pml4;
-	args->table_pages = tables_needed;
+	args->page_tables = { .pointer = pml4 + 1, .count = tables_needed - 1 };
-	args->table_count = tables_needed - 1;
-	args->page_tables = offset_ptr<void>(addr, page_size);
 
 	console::print(L"    First page (pml4) at: 0x%lx\r\n", pml4);
 
-	add_kernel_pds(pml4, args->page_tables, args->table_count);
+	add_kernel_pds(pml4, args->page_tables);
-	add_offset_mappings(pml4, args->page_tables, args->table_count);
+	add_offset_mappings(pml4, args->page_tables);
 
 	//console::print(L"    Set up initial mappings, %d spare tables.\r\n", args->table_count);
 }
@@ -243,10 +231,7 @@ map_pages(
 	paging::page_table *pml4 =
 		reinterpret_cast<paging::page_table*>(args->pml4);
 
-	page_entry_iterator<4> iterator{
+	page_entry_iterator<4> iterator{virt, pml4, args->page_tables};
-		virt, pml4,
-		args->page_tables,
-		args->table_count};
 
 	uint64_t flags = page_flags;
 	if (!exe_flag)

src/boot/paging.h

@@ -30,9 +30,7 @@ struct page_table
 /// and kernel args' `page_table_cache` and `num_free_tables` are updated with
 /// the leftover space.
 /// \arg args    The kernel args struct, used for the page table cache and pml4
-void allocate_tables(
+void allocate_tables(kernel::init::args *args);
-	kernel::init::args *args,
-	uefi::boot_services *bs);
 
 /// Copy existing page table entries to a new page table. Does not do a deep
 /// copy - the new PML4 is updated to point to the existing next-level page

src/boot/status.cpp

@@ -1,5 +1,6 @@
 #include <uefi/types.h>
 #include <uefi/graphics.h>
+#include <uefi/protos/simple_text_output.h>
 
 #include "console.h"
 #include "error.h"

src/boot/support.cpp

@@ -34,6 +34,9 @@ int _purecall()
 	::boot::error::raise(uefi::status::unsupported, L"Pure virtual call");
 }
 
-} // extern "C"
+/// Clang can emit calls to atexit() in constructors or destructors, but
+/// those calls don't make sense for a bootloader. Implement an empty stub
+/// to satisfy the linker.
+int atexit(void (*function)(void)) { return 0; }
 
-void operator delete (void *) {}
+} // extern "C"

src/boot/types.h

@@ -1,13 +0,0 @@
-/// \file types.h
-/// Definitions of shared types used throughout the bootloader
-#pragma once
-
-namespace boot {
-
-struct buffer
-{
-	size_t size;
-	void *data;
-};
-
-} // namespace boot

src/include/counted.h

@@ -0,0 +1,39 @@
+#pragma once
+/// \file counted.h
+/// Definition of the `counted` template class
+
+#include "pointer_manipulation.h"
+
+/// A pointer and an associated count. Memory pointed to is not managed.
+/// Depending on the usage, the count may be size or number of elements.
+/// Helper methods provide the ability to treat the pointer like an array.
+template <typename T>
+struct counted
+{
+	T *pointer;
+	size_t count;
+
+	/// Index this object as an array of type T
+	inline T & operator [] (int i) { return pointer[i]; }
+
+	/// Index this object as a const array of type T
+	inline const T & operator [] (int i) const { return pointer[i]; }
+
+	using iterator = offset_iterator<T>;
+
+	/// Return an iterator to the beginning of the array
+	inline iterator begin() { return iterator(pointer, sizeof(T)); }
+
+	/// Return an iterator to the end of the array
+	inline iterator end() { return offset_ptr<T>(pointer, sizeof(T)*count); }
+};
+
+/// Specialize for `void` which cannot be indexed or iterated
+template <>
+struct counted<void>
+{
+	void *pointer;
+	size_t count;
+};
+
+using buffer = counted<void>;

src/include/kernel_args.h

@@ -3,6 +3,7 @@
 #include <stdalign.h>
 #include <stddef.h>
 #include <stdint.h>
+#include "counted.h"
 
 namespace kernel {
 namespace init {
@@ -55,10 +56,8 @@ struct program_section {
 
 struct program {
 	uintptr_t entrypoint;
-	uintptr_t base;
+	uintptr_t phys_base;
-	size_t total_size;
+	counted<program_section> sections;
-	size_t num_sections;
-	program_section sections[3];
 };
 
 enum class mem_type : uint32_t {
@@ -79,13 +78,56 @@ struct mem_entry
 	uint32_t attr;
 };
 
+enum class allocation_type : uint8_t {
+	none, page_table, mem_map, frame_map, file, program,
+};
+
+/// A single contiguous allocation of pages
+struct page_allocation
+{
+	uintptr_t address;
+	uint32_t count;
+	allocation_type type;
+};
+
+/// A page-sized register of page_allocation entries
+struct allocation_register
+{
+	allocation_register *next;
+	uint8_t count;
+
+	uint8_t reserved0;
+	uint16_t reserved1;
+	uint32_t reserved2;
+
+	page_allocation entries[255];
+};
+
+enum class frame_flags : uint32_t {
+	uncacheable         = 0x00000001,
+	write_combining     = 0x00000002,
+	write_through       = 0x00000004,
+	write_back          = 0x00000008,
+	uncache_exported    = 0x00000010,
+
+	write_protect       = 0x00001000,
+	read_protect        = 0x00002000,
+	exec_protect        = 0x00004000,
+	non_volatile        = 0x00008000,
+
+	read_only           = 0x00020000,
+	earmarked           = 0x00040000,
+	hw_crypto           = 0x00080000,
+};
+
+
 constexpr size_t frames_per_block = 64 * 64 * 64;
 
 struct frame_block
 {
 	uintptr_t base;
 	uint32_t count;
-	uint32_t attrs;
+	frame_flags flags;
 	uint64_t map1;
 	uint64_t map2[64];
 	uint64_t *bitmap;
@@ -103,22 +145,14 @@ struct args
 	boot_flags flags;
 
 	void *pml4;
-	void *page_tables;
+	counted<void> page_tables;
-	size_t table_count;
-	size_t table_pages;
 
-	program *programs;
+	counted<program> programs;
-	size_t num_programs;
+	counted<module> modules;
+	counted<mem_entry> mem_map;
+	counted<frame_block> frame_blocks;
 
-	module *modules;
+	allocation_register *allocations;
-	size_t num_modules;
-
-	mem_entry *mem_map;
-	size_t map_count;
-
-	frame_block *frame_blocks;
-	size_t frame_block_count;
-	size_t frame_block_pages;
 
 	void *runtime_services;
 	void *acpi_table;

src/include/pointer_manipulation.h

@@ -2,8 +2,6 @@
 /// Helper functions and types for doing type-safe byte-wise pointer math.
 #pragma once
 
-namespace boot {
-
 /// Return a pointer offset from `input` by `offset` bytes.
 /// \tparam T  Cast the return value to a pointer to `T`
 /// \tparam S  The type pointed to by the `input` pointer
@@ -27,15 +25,17 @@ public:
 
 	T* operator++() { m_t = offset_ptr<T>(m_t, m_off); return m_t; }
 	T* operator++(int) { T* tmp = m_t; operator++(); return tmp; }
-	bool operator==(T* p) { return p == m_t; }
 
-	T* operator*() const { return m_t; }
+	bool operator==(T* p) { return p == m_t; }
-	operator T*() const { return m_t; }
+	bool operator!=(T* p) { return p != m_t; }
+	bool operator==(offset_iterator<T> &i) { return i.m_t == m_t; }
+	bool operator!=(offset_iterator<T> &i) { return i.m_t != m_t; }
+
+	T& operator*() const { return *m_t; }
+	operator T& () const { return *m_t; }
 	T* operator->() const { return m_t; }
 
 private:
 	T* m_t;
 	size_t m_off;
 };
-
-} // namespace boot

src/kernel/main.cpp

@@ -147,7 +147,7 @@ kernel_main(init::args *args)
 
 	cpu->tss->create_ist_stacks(cpu->idt->used_ist_entries());
 
-	for (size_t i = 0; i < args->num_modules; ++i) {
+	for (size_t i = 0; i < args->modules.count; ++i) {
 		init::module &mod = args->modules[i];
 
 		switch (mod.type) {
@@ -213,7 +213,7 @@ kernel_main(init::args *args)
 	scheduler_ready = true;
 
 	// Skip program 0, which is the kernel itself
-	for (unsigned i = 1; i < args->num_programs; ++i)
+	for (unsigned i = 1; i < args->programs.count; ++i)
 		load_simple_process(args->programs[i]);
 
 	if (!has_video)

src/kernel/memory_bootstrap.cpp

@@ -28,6 +28,7 @@ namespace init {
 is_bitfield(section_flags);
 }}
 
+using kernel::init::allocation_register;
 using kernel::init::section_flags;
 
 using namespace kernel;
@@ -82,36 +83,15 @@ memory_initialize_pre_ctors(init::args &kargs)
 	new (&g_kernel_heap) kutil::heap_allocator {heap_start, kernel_max_heap};
 
 	frame_block *blocks = reinterpret_cast<frame_block*>(memory::bitmap_start);
-	new (&g_frame_allocator) frame_allocator {blocks, kargs.frame_block_count};
+	new (&g_frame_allocator) frame_allocator {blocks, kargs.frame_blocks.count};
 
 	// Mark all the things the bootloader allocated for us as used
-	g_frame_allocator.used(
+	allocation_register *reg = kargs.allocations;
-		get_physical_page(&kargs),
+	while (reg) {
-		memory::page_count(sizeof(kargs)));
+		for (auto &alloc : reg->entries)
-
+			if (alloc.type != init::allocation_type::none)
-	g_frame_allocator.used(
+				g_frame_allocator.used(alloc.address, alloc.count);
-		get_physical_page(kargs.frame_blocks),
+		reg = reg->next;
-		kargs.frame_block_pages);
-
-	g_frame_allocator.used(
-		get_physical_page(kargs.pml4),
-		kargs.table_pages);
-
-	for (unsigned i = 0; i < kargs.num_modules; ++i) {
-		const kernel::init::module &mod = kargs.modules[i];
-		g_frame_allocator.used(
-			get_physical_page(mod.location),
-			memory::page_count(mod.size));
-	}
-
-	for (unsigned i = 0; i < kargs.num_programs; ++i) {
-		const kernel::init::program &prog = kargs.programs[i];
-		for (auto &sect : prog.sections) {
-			if (!sect.size) continue;
-			g_frame_allocator.used(
-				sect.phys_addr,
-				memory::page_count(sect.size));
-		}
 	}
 
 	process *kp = process::create_kernel_process(kpml4);
@@ -141,8 +121,8 @@ memory_initialize_post_ctors(init::args &kargs)
 	vm.add(memory::buffers_start, &g_kernel_buffers);
 
 	g_frame_allocator.free(
-		get_physical_page(kargs.page_tables),
+		get_physical_page(kargs.page_tables.pointer),
-		kargs.table_count);
+		kargs.page_tables.count);
 }
 
 static void