Document methods and move them into kernel cpu folder

3 files changed, 58 insertions, 59 deletions

diff --git a/arch/x86_64/src/context_switching/main.cpp b/arch/x86_64/src/context_switching/main.cpp
index 0f2ec93..5cac878 100644
--- a/arch/x86_64/src/context_switching/main.cpp
+++ b/arch/x86_64/src/context_switching/main.cpp
@@ -12,46 +12,6 @@ namespace teachos::arch::context_switching
 {
   namespace
   {
-
-    /**
-     * @brief Switch context into the mode defined in the segment selectors.
-     *
-     * Setup the stack IRETQ expects to switch the mode:
-     * 1. push data selector
-     * 2. push current stack pointer
-     * 3. push eflags
-     * 4. push code segment selector
-     * 5. push return address
-     *
-     * @param data_segment
-     * @param code_segment
-     * @param address
-     */
-    [[gnu::naked]]
-    auto far_return(context_switching::interrupt_descriptor_table::segment_selector data_segment,
-                    context_switching::interrupt_descriptor_table::segment_selector code_segment, uint64_t address)
-        -> void
-    {
-      asm volatile("mov %[data_segment], %%rax\n"
-                   "mov %%rax, %%ds\n"
-                   "mov %%rax, %%es\n"
-                   "mov %%rax, %%fs\n"
-                   "mov %%rax, %%gs\n"
-                   "mov %%rsp, %%rax\n"
-
-                   "push %[data_segment]\n"
-                   "push %%rax\n"
-                   "pushfq\n"
-                   "push %[code_segment]\n"
-                   "mov %[return_function], %%rax\n"
-                   "push %%rax\n"
-
-                   "iretq\n"
-                   :
-                   : [data_segment] "m"(data_segment), [code_segment] "m"(code_segment), [return_function] "r"(address)
-                   : "rax");
-    }
-
     constexpr context_switching::interrupt_descriptor_table::segment_selector USER_CODE_SEGMENT_SELECTOR{
         3U, context_switching::interrupt_descriptor_table::segment_selector::REQUEST_LEVEL_USER};
     constexpr context_switching::interrupt_descriptor_table::segment_selector USER_DATA_SEGMENT_SELECTOR{
@@ -61,6 +21,14 @@ namespace teachos::arch::context_switching
 
   auto initialize_descriptor_tables() -> descriptor_tables
   {
+    static bool initalized = false;
+    if (initalized)
+    {
+      descriptor_tables tables = {segment_descriptor_table::get_or_create_global_descriptor_table(),
+                                  interrupt_descriptor_table::get_or_create_interrupt_descriptor_table()};
+      return tables;
+    }
+
     kernel::cpu::clear_interrupt_flag();
 
     segment_descriptor_table::update_global_descriptor_table_register();
@@ -68,7 +36,7 @@ namespace teachos::arch::context_switching
 
     interrupt_descriptor_table::segment_selector segment_selector{
         1U, interrupt_descriptor_table::segment_selector::REQUEST_LEVEL_KERNEL};
-    kernel::cpu::far_pointer pointer{&kernel::cpu::reload_segment_registers, segment_selector};
+    kernel::cpu::far_pointer pointer{&kernel::cpu::reload_data_segment_registers, segment_selector};
     kernel::cpu::call(pointer);
 
     segment_descriptor_table::update_task_state_segment_register();
@@ -77,17 +45,13 @@ namespace teachos::arch::context_switching
 
     descriptor_tables tables = {segment_descriptor_table::get_or_create_global_descriptor_table(),
                                 interrupt_descriptor_table::get_or_create_interrupt_descriptor_table()};
+    initalized = true;
     return tables;
   }
 
   auto user_mode_main() -> void
   {
-    auto current_segment = kernel::cpu::read_code_segment_register();
-    exception_handling::assert(USER_CODE_SEGMENT_SELECTOR == current_segment,
-                               "[Context Switching] Context switch into user mode not successful");
-    exception_handling::assert(USER_DATA_SEGMENT_SELECTOR == kernel::cpu::validate_data_segment_registers(),
-                               "[Context Switching] Context switch into user mode not successful");
-
+    kernel::cpu::validate_segment_registers(USER_DATA_SEGMENT_SELECTOR, USER_CODE_SEGMENT_SELECTOR);
     video::vga::text::write("Successfully entered user mode!", video::vga::text::common_attributes::green_on_black);
   }
 
@@ -99,8 +63,9 @@ namespace teachos::arch::context_switching
   auto switch_context(interrupt_descriptor_table::segment_selector data_segment,
                       interrupt_descriptor_table::segment_selector code_segment, void (*return_function)()) -> void
   {
-    kernel::cpu::set_segment_registers(data_segment);
-    far_return(data_segment, code_segment, reinterpret_cast<uint64_t>(return_function));
+    (void)initialize_descriptor_tables();
+    kernel::cpu::set_data_segment_registers(data_segment);
+    kernel::cpu::set_code_segment_register(data_segment, code_segment, reinterpret_cast<uint64_t>(return_function));
   }
 
 }  // namespace teachos::arch::context_switching
diff --git a/arch/x86_64/src/kernel/cpu/segment_register.cpp b/arch/x86_64/src/kernel/cpu/segment_register.cpp
index cb367b6..44b4eff 100644
--- a/arch/x86_64/src/kernel/cpu/segment_register.cpp
+++ b/arch/x86_64/src/kernel/cpu/segment_register.cpp
@@ -6,7 +6,7 @@
 namespace teachos::arch::kernel::cpu
 {
   [[gnu::naked]]
-  auto reload_segment_registers() -> void
+  auto reload_data_segment_registers() -> void
   {
     asm volatile("xor %rax, %rax\n"
                  "mov %rax, %ss\n"
@@ -18,7 +18,8 @@ namespace teachos::arch::kernel::cpu
   }
 
   [[gnu::naked]]
-  auto set_segment_registers(context_switching::interrupt_descriptor_table::segment_selector segment_selector) -> void
+  auto set_data_segment_registers(context_switching::interrupt_descriptor_table::segment_selector segment_selector)
+      -> void
   {
     asm volatile("xor %%rax, %%rax\n"
                  "mov %[input], %%ax\n"
@@ -38,7 +39,8 @@ namespace teachos::arch::kernel::cpu
     return current_value;
   }
 
-  auto validate_data_segment_registers() -> context_switching::interrupt_descriptor_table::segment_selector
+  auto validate_data_segment_registers(context_switching::interrupt_descriptor_table::segment_selector data_segment)
+      -> void
   {
     context_switching::interrupt_descriptor_table::segment_selector ss;
     context_switching::interrupt_descriptor_table::segment_selector ds;
@@ -54,10 +56,46 @@ namespace teachos::arch::kernel::cpu
         "mov %%gs, %[gs_output]\n"
         : [ss_output] "=r"(ss), [ds_output] "=r"(ds), [es_output] "=r"(es), [fs_output] "=r"(fs), [gs_output] "=r"(gs));
 
-    auto result = ss == ds && ss == es && ss == fs && ss == gs;
+    auto result = (ss == ds && ss == es && ss == fs && ss == gs);
     exception_handling::assert(result, "[Segment Register] Values in data register are not the same.");
+    result = (ss == data_segment);
+    exception_handling::assert(
+        result, "[Segment Register] Expected Data Segment is not the same as the value in the Stack Segment register.");
+  }
+
+  auto validate_code_segment_register(context_switching::interrupt_descriptor_table::segment_selector code_segment)
+      -> void
+  {
+    auto const cs = read_code_segment_register();
+    exception_handling::assert(
+        cs == code_segment,
+        "[Segment Register] Expected Code Segment is not the same as the value in the Code Segment register.");
+  }
+
+  auto validate_segment_registers(context_switching::interrupt_descriptor_table::segment_selector data_segment,
+                                  context_switching::interrupt_descriptor_table::segment_selector code_segment) -> void
+  {
+    validate_data_segment_registers(data_segment);
+    validate_code_segment_register(code_segment);
+  }
+
+  [[gnu::naked]]
+  auto set_code_segment_register(context_switching::interrupt_descriptor_table::segment_selector data_segment,
+                                 context_switching::interrupt_descriptor_table::segment_selector code_segment,
+                                 uint64_t address) -> void
+  {
+    asm volatile("mov %%rsp, %%rax\n"
+                 "push %[data_segment]\n"
+                 "push %%rax\n"
+                 "pushfq\n"
+                 "push %[code_segment]\n"
+                 "mov %[return_function], %%rax\n"
+                 "push %%rax\n"
 
-    return ss;
+                 "iretq\n"
+                 :
+                 : [data_segment] "m"(data_segment), [code_segment] "m"(code_segment), [return_function] "r"(address)
+                 : "rax");
   }
 
 }  // namespace teachos::arch::kernel::cpu
diff --git a/arch/x86_64/src/kernel/main.cpp b/arch/x86_64/src/kernel/main.cpp
index b69064d..05c879e 100644
--- a/arch/x86_64/src/kernel/main.cpp
+++ b/arch/x86_64/src/kernel/main.cpp
@@ -61,11 +61,7 @@ namespace teachos::arch::kernel
     // stack_overflow_test(0);
 
     memory::heap::global_heap_allocator::register_heap_allocator(memory::heap::heap_allocator_type::LINKED_LIST);
-
-    heap_test();
-
-    decltype(auto) descriptor_tables = context_switching::initialize_descriptor_tables();
-    (void)descriptor_tables;
+    // heap_test();
 
     context_switching::switch_to_user_mode();
   }