From 86780fc776d05c334449ad5aa72ff4cfc324f8ee Mon Sep 17 00:00:00 2001
From: mintsuki <mintsuki@protonmail.com>
Date: Thu, 2 Sep 2021 02:47:08 +0200
Subject: [PATCH] efi: Improve reclaim memory process

---
 Makefile            |  2 +-
 stage23/mm/pmm.s2.c | 34 ++++++++++++++++++++++++++--------
 2 files changed, 27 insertions(+), 9 deletions(-)

diff --git a/Makefile b/Makefile
index 09211a7d..d12b51c6 100644
--- a/Makefile
+++ b/Makefile
@@ -319,7 +319,7 @@ full-hybrid-test:
 	cp -rv bin/* test/* test_image/boot/
 	xorriso -as mkisofs -b boot/limine-cd.bin -no-emul-boot -boot-load-size 4 -boot-info-table --efi-boot boot/limine-eltorito-efi.bin -efi-boot-part --efi-boot-image --protective-msdos-label test_image/ -o test.iso
 	bin/limine-install test.iso
-	qemu-system-x86_64 -m 512M -M q35 -bios ovmf-x64/OVMF.fd -net none -smp 4   -cdrom test.iso -debugcon stdio
+	qemu-system-x86_64 -M q35 -bios ovmf-x64/OVMF.fd -net none -smp 4   -cdrom test.iso -debugcon stdio
 	qemu-system-x86_64 -m 512M -M q35 -bios ovmf-x64/OVMF.fd -net none -smp 4   -hda test.iso -debugcon stdio
 	qemu-system-x86_64 -m 512M -M q35 -bios ovmf-ia32/OVMF.fd -net none -smp 4   -cdrom test.iso -debugcon stdio
 	qemu-system-x86_64 -m 512M -M q35 -bios ovmf-ia32/OVMF.fd -net none -smp 4   -hda test.iso -debugcon stdio
diff --git a/stage23/mm/pmm.s2.c b/stage23/mm/pmm.s2.c
index 3c63a12d..86caabe7 100644
--- a/stage23/mm/pmm.s2.c
+++ b/stage23/mm/pmm.s2.c
@@ -12,7 +12,7 @@
 #endif
 
 #define PAGE_SIZE   4096
-#define MEMMAP_MAX_ENTRIES 256
+#define MEMMAP_MAX_ENTRIES 384
 
 #if bios == 1
 extern symbol bss_end;
@@ -319,7 +319,9 @@ void init_memmap(void) {
     allocations_disallowed = false;
 
     // Let's leave 64MiB to the firmware
-    ext_mem_alloc_type(0x4000000, MEMMAP_EFI_RECLAIMABLE);
+    for (size_t i = 0; i < 64; i++) {
+        ext_mem_alloc_type(0x100000, MEMMAP_EFI_RECLAIMABLE);
+    }
 
     memcpy(untouched_memmap, memmap, memmap_entries * sizeof(struct e820_entry_t));
     untouched_memmap_entries = memmap_entries;
@@ -372,23 +374,34 @@ void pmm_reclaim_uefi_mem(void) {
         }
     }
 
-    struct e820_entry_t recl;
+    size_t recl_i = 0;
 
-    for (size_t i = 0; ; i++) {
+    for (size_t i = 0; i < memmap_entries; i++) {
         if (memmap[i].type == MEMMAP_EFI_RECLAIMABLE) {
-            recl = memmap[i];
-            break;
+            recl_i++;
         }
     }
 
+    struct e820_entry_t *recl = ext_mem_alloc(recl_i * sizeof(struct e820_entry_t));
+
+    {
+        size_t recl_j = 0;
+        for (size_t i = 0; i < memmap_entries; i++) {
+            if (memmap[i].type == MEMMAP_EFI_RECLAIMABLE) {
+                recl[recl_j++] = memmap[i];
+            }
+        }
+    }
+
+another_recl:;
     // Punch holes in our EFI reclaimable entry for every EFI area which is
     // boot services or conventional that fits within
     size_t efi_mmap_entry_count = efi_mmap_size / efi_desc_size;
     for (size_t i = 0; i < efi_mmap_entry_count; i++) {
         EFI_MEMORY_DESCRIPTOR *entry = (void *)efi_mmap + i * efi_desc_size;
 
-        uintptr_t base = recl.base;
-        uintptr_t top = base + recl.length;
+        uintptr_t base = recl->base;
+        uintptr_t top = base + recl->length;
         uintptr_t efi_base = entry->PhysicalStart;
         uintptr_t efi_size = entry->NumberOfPages * 4096;
 
@@ -430,6 +443,11 @@ void pmm_reclaim_uefi_mem(void) {
         memmap_alloc_range(efi_base, efi_size, our_type, false, true, false, true);
     }
 
+    if (--recl_i > 0) {
+        recl++;
+        goto another_recl;
+    }
+
     sanitise_entries(memmap, &memmap_entries, false);
 }