From 24187b5f80eef36c318c82a22aa512ac6364ed64 Mon Sep 17 00:00:00 2001
From: bjorn3 <17426603+bjorn3@users.noreply.github.com>
Date: Sun, 25 Feb 2024 20:10:49 +0100
Subject: [PATCH] Move set_tss_stack call to switch_to

---
 src/context/arch/x86.rs    | 4 ++++
 src/context/arch/x86_64.rs | 4 ++++
 src/context/switch.rs      | 6 ------
 3 files changed, 8 insertions(+), 6 deletions(-)

diff --git a/src/context/arch/x86.rs b/src/context/arch/x86.rs
index ff2f5bfba7..4a67efe732 100644
--- a/src/context/arch/x86.rs
+++ b/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}]
diff --git a/src/context/arch/x86_64.rs b/src/context/arch/x86_64.rs
index 6eaa3ca074..8d5f9d3827 100644
--- a/src/context/arch/x86_64.rs
+++ b/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",
diff --git a/src/context/switch.rs b/src/context/switch.rs
index 8870e97168..4cc17248f3 100644
--- a/src/context/switch.rs
+++ b/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);