4 files changed, 82 insertions, 32 deletions
diff --git a/arch/x86_64/src/memory/cpu/control_register.cpp b/arch/x86_64/src/memory/cpu/control_register.cpp
new file mode 100644
index 0000000..38a887c
--- /dev/null
+++ b/arch/x86_64/src/memory/cpu/control_register.cpp
@@ -0,0 +1,66 @@
+#include "arch/memory/cpu/control_register.hpp"
+
+#include "arch/exception_handling/assert.hpp"
+
+namespace teachos::arch::memory::cpu
+{
+  auto read_control_register(control_register cr) -> std::size_t
+  {
+    std::size_t current_value;
+    switch (cr)
+    {
+      case control_register::CR0:
+        asm volatile("movq %%cr0, %[output]" : [output] "=r"(current_value));
+        break;
+      case control_register::CR2:
+        asm volatile("movq %%cr2, %[output]" : [output] "=r"(current_value));
+        break;
+      case control_register::CR3:
+        asm volatile("movq %%cr3, %[output]" : [output] "=r"(current_value));
+        break;
+      case control_register::CR4:
+        asm volatile("movq %%cr4, %[output]" : [output] "=r"(current_value));
+        break;
+      default:
+        exception_handling::assert(false,
+                                   "[Control Register] Attempted to read non-existent or reserved control register");
+        break;
+    }
+    return current_value;
+  }
+
+  auto write_control_register(control_register cr, std::size_t new_value) -> void
+  {
+    switch (cr)
+    {
+      case control_register::CR0:
+        asm volatile("movq %[input], %%cr0"
+                     : /* no output from call */
+                     : [input] "r"(new_value)
+                     : "memory");
+        break;
+      case control_register::CR2:
+        asm volatile("movq %[input], %%cr2"
+                     : /* no output from call */
+                     : [input] "r"(new_value)
+                     : "memory");
+        break;
+      case control_register::CR3:
+        asm volatile("movq %[input], %%cr3"
+                     : /* no output from call */
+                     : [input] "r"(new_value)
+                     : "memory");
+        break;
+      case control_register::CR4:
+        asm volatile("movq %[input], %%cr4"
+                     : /* no output from call */
+                     : [input] "r"(new_value)
+                     : "memory");
+        break;
+      default:
+        exception_handling::assert(false,
+                                   "[Control Register] Attempted to write non-existent or reserved control register");
+        break;
+    }
+  }
+}  // namespace teachos::arch::memory::cpu
diff --git a/arch/x86_64/src/memory/cpu/cr3.cpp b/arch/x86_64/src/memory/cpu/cr3.cpp
deleted file mode 100644
index 7e48d40..0000000
--- a/arch/x86_64/src/memory/cpu/cr3.cpp
+++ /dev/null
@@ -1,23 +0,0 @@
-#include "arch/memory/cpu/cr3.hpp"
-
-#include "arch/exception_handling/assert.hpp"
-
-namespace teachos::arch::memory::cpu
-{
-  auto read_cr3_register() -> allocator::physical_address
-  {
-    allocator::physical_address cr3;
-    asm volatile("movq %%cr3, %[output]" : [output] "=r"(cr3) : /* no input into call */ : "memory");
-    return cr3;
-  }
-
-  auto write_cr3_register(allocator::physical_address new_p4_table_address) -> void
-  {
-    exception_handling::assert(new_p4_table_address % allocator::PAGE_FRAME_SIZE == 0U,
-                               "[CR3] Physical address to be written into register must be page aligned");
-    asm volatile("movq %[input], %%cr3"
-                 : /* no output from call */
-                 : [input] "r"(new_p4_table_address)
-                 : "memory");
-  }
-}  // namespace teachos::arch::memory::cpu
diff --git a/arch/x86_64/src/memory/cpu/msr.cpp b/arch/x86_64/src/memory/cpu/msr.cpp
index 3a917c9..6e3d1d3 100644
--- a/arch/x86_64/src/memory/cpu/msr.cpp
+++ b/arch/x86_64/src/memory/cpu/msr.cpp
@@ -2,14 +2,19 @@
 
 namespace teachos::arch::memory::cpu
 {
-  uint64_t read_msr(uint32_t msr)
+  namespace
+  {
+    constexpr uint32_t IA32_EFER_ADDRESS = 0xC0000080;
+  }
+
+  auto read_msr(uint32_t msr) -> uint64_t
   {
     uint32_t low, high;
     asm volatile("rdmsr" : "=a"(low), "=d"(high) : "c"(msr));
     return (static_cast<uint64_t>(high) << 32) | low;
   }
 
-  void write_msr(uint32_t msr, uint64_t value)
+  auto write_msr(uint32_t msr, uint64_t value) -> void
   {
     uint32_t low = value & 0xFFFFFFFF;
     uint32_t high = value >> 32;
@@ -18,12 +23,11 @@ namespace teachos::arch::memory::cpu
                  : "c"(msr), "a"(low), "d"(high));
   }
 
-  void set_msr_bit(msr_flags flag)
+  auto set_efer_bit(efer_flags flag) -> void
   {
-    uint64_t efer = read_msr(IA32_EFER);
-    write_msr(IA32_EFER, static_cast<uint64_t>(flag) | efer);
+    uint64_t const efer = read_msr(IA32_EFER_ADDRESS);
+    write_msr(IA32_EFER_ADDRESS, static_cast<uint64_t>(flag) | efer);
   }
 
-  void enable_nxe_bit() { set_msr_bit(msr_flags::NXE); }
-
+  auto enable_nxe_bit() -> void { set_efer_bit(efer_flags::NXE); }
 }  // namespace teachos::arch::memory::cpu
diff --git a/arch/x86_64/src/memory/cpu/tlb.cpp b/arch/x86_64/src/memory/cpu/tlb.cpp
index 1663e80..591d9fc 100644
--- a/arch/x86_64/src/memory/cpu/tlb.cpp
+++ b/arch/x86_64/src/memory/cpu/tlb.cpp
@@ -1,6 +1,6 @@
 #include "arch/memory/cpu/tlb.hpp"
 
-#include "arch/memory/cpu/cr3.hpp"
+#include "arch/memory/cpu/control_register.hpp"
 
 namespace teachos::arch::memory::cpu
 {
@@ -9,5 +9,8 @@ namespace teachos::arch::memory::cpu
     asm volatile("invlpg (%[input])" : /* no output from call */ : [input] "r"(address) : "memory");
   }
 
-  auto tlb_flush_all() -> void { write_cr3_register(read_cr3_register()); }
+  auto tlb_flush_all() -> void
+  {
+    write_control_register(cpu::control_register::CR3, read_control_register(cpu::control_register::CR3));
+  }
 }  // namespace teachos::arch::memory::cpu