2 files changed, 16 insertions(+), 6 deletions(-)

M efi/mmap.ha
M efi/systab.ha

@@ 33,6 33,7 @@ export type MEMORY_DESCRIPTOR = struct {
 // in practice we may have less than 128 items here.
 def MMAP_STATIC: size = 128 * size(MEMORY_DESCRIPTOR);
 let mmap_buf: [MMAP_STATIC]u8 = [0...];
+let mmap_key: u64 = 0;
 
 export type mmap_iter = struct {
 	i: size,


@@ 43,11 44,10 @@ export type mmap_iter = struct {
 // Iterates over the EFI memory map.
 export fn iter_mmap() (mmap_iter | error) = {
 	let msize = 0z;
-	let mkey = 0u64;
 	let dsize = 0z;
 	let dver = 0u32;
 	const err = systab.BootServices.GetMemoryMap(&msize,
-		null, &mkey, &dsize, &dver);
+		null, &mmap_key, &dsize, &dver);
 	assert(err == STATUS::BUFFER_TOO_SMALL);
 
 	if (msize > len(mmap_buf)) {


@@ 56,7 56,7 @@ export fn iter_mmap() (mmap_iter | error) = {
 	assert(dver == 1);
 
 	efi_call(systab.BootServices.GetMemoryMap(&msize,
-			&mmap_buf[0], &mkey, &dsize, &dver))?;
+			&mmap_buf[0], &mmap_key, &dsize, &dver))?;
 	return mmap_iter {
 		i = 0,
 		n = msize / dsize,

@@ 110,12 110,22 @@ export type CONFIGURATION_TABLE = struct {
 
 // Instructs the EFI boot services to exit.
 export fn exit_boot_services() void = {
-	let map_key = 0u64;
 	let msize = 0z, dsize = 0z, dver = 0u32;
+
+	const err = systab.BootServices.ExitBootServices(image, mmap_key);
+
+	if(err == STATUS::SUCCESS){
+		return;
+	};
+
 	const err = systab.BootServices.GetMemoryMap(
-		&msize, null, &map_key, &dsize, &dver);
+		&msize, null, &mmap_key, &dsize, &dver);
+
 	assert(err == STATUS::BUFFER_TOO_SMALL);
 
-	const err = systab.BootServices.ExitBootServices(image, map_key);
+	const err = systab.BootServices.GetMemoryMap(
+		&msize, &mmap_buf[0], &mmap_key, &dsize, &dver);
+
+	const err = systab.BootServices.ExitBootServices(image, mmap_key);
 	assert(err == STATUS::SUCCESS);
 };