x86_64/memory: extract PML4 injection

1 files changed, 16 insertions, 2 deletions

diff --git a/arch/x86_64/src/memory/paging_root.cpp b/arch/x86_64/src/memory/paging_root.cpp
index 6b65f60..a29304e 100644
--- a/arch/x86_64/src/memory/paging_root.cpp
+++ b/arch/x86_64/src/memory/paging_root.cpp
@@ -40,13 +40,27 @@ namespace teachos::memory::x86_64
 
     auto handle_huge_page = [&] -> std::optional<frame> {
       auto pml3_entry = pml3.transform([&](auto pml3) -> auto { return (*pml3)[pml_index<3>(page)]; });
-      if (pml3_entry && pml3_entry->huge())
+      if (!pml3_entry)
+      {
+        return std::nullopt;
+      }
+      else if (pml3_entry->huge())
       {
         auto pml3_entry_frame = *pml3_entry->frame();
         return frame{pml3_entry_frame.number() + pml_index<2>(page) * entry_count + pml_index<1>(page)};
       }
 
-      // auto pml3_entry = (**pml3)[page.start_address().raw() >> 39 & 0x1ff];
+      auto pml2 = (*pml3)->next(pml_index<3>(page));
+      auto pml2_entry = pml2.transform([&](auto pml2) -> auto { return (*pml2)[pml_index<2>(page)]; });
+      if (!pml2_entry)
+      {
+        return std::nullopt;
+      }
+      else if (pml2_entry->huge())
+      {
+        auto pml2_entry_frame = *pml2_entry->frame();
+        return frame{pml2_entry_frame.number() + pml_index<1>(page)};
+      }
 
       return std::nullopt;
     };