Move set_tss_stack call to switch_to

src/context/arch/x86.rs

@@ -129,6 +129,10 @@ pub unsafe fn empty_cr3() -> rmm::PhysicalAddress {
 
 /// Switch to the next context by restoring its stack and registers
 pub unsafe fn switch_to(prev: &mut super::Context, next: &mut super::Context) {
+    if let Some(ref stack) = next.kstack {
+        crate::gdt::set_tss_stack(stack.as_ptr() as usize + stack.len());
+    }
+
     core::arch::asm!("
         fxsave [{prev_fx}]
         fxrstor [{next_fx}]

src/context/arch/x86_64.rs

@@ -143,6 +143,10 @@ pub unsafe fn empty_cr3() -> rmm::PhysicalAddress {
 
 /// Switch to the next context by restoring its stack and registers
 pub unsafe fn switch_to(prev: &mut super::Context, next: &mut super::Context) {
+    if let Some(ref stack) = next.kstack {
+        crate::gdt::set_tss_stack(stack.as_ptr() as usize + stack.len());
+    }
+
     core::arch::asm!(
         alternative2!(
             feature1: "xsaveopt",

src/context/switch.rs

@@ -198,12 +198,6 @@ pub unsafe fn switch() -> bool {
         next_context.cpu_id = Some(cpu_id);
         next_context.switch_time = switch_time;
 
-        #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
-        {
-            if let Some(ref stack) = next_context.kstack {
-                crate::gdt::set_tss_stack(stack.as_ptr() as usize + stack.len());
-            }
-        }
         let percpu = PercpuBlock::current();
         percpu.switch_internals.context_id.set(next_context.id);