From ca85a69442c85250e8985a57136df640d9bc1d6a Mon Sep 17 00:00:00 2001
From: apio <blobs.trading@gmail.com>
Date: Tue, 7 Mar 2023 20:59:11 +0100
Subject: [PATCH] kernel: Save floating-point environment when switching in and
 out of userspace threads

---
 kernel/src/arch/x86_64/CPU.cpp  | 12 ++++++++++++
 kernel/src/arch/x86_64/CPU.h    | 10 ++++++++++
 kernel/src/thread/Scheduler.cpp |  2 ++
 kernel/src/thread/Thread.h      |  4 ++--
 4 files changed, 26 insertions(+), 2 deletions(-)

diff --git a/kernel/src/arch/x86_64/CPU.cpp b/kernel/src/arch/x86_64/CPU.cpp
index b703071c..174fbbc9 100644
--- a/kernel/src/arch/x86_64/CPU.cpp
+++ b/kernel/src/arch/x86_64/CPU.cpp
@@ -23,6 +23,18 @@ extern void pic_eoi(unsigned char irq);
 extern void pic_eoi(Registers* regs);
 extern void setup_idt();
 
+void FPData::save()
+{
+    asm volatile("fxsave (%0)" : : "r"(m_data));
+    m_already_saved = true;
+}
+
+void FPData::restore()
+{
+    if(!m_already_saved) return;
+    asm volatile("fxrstor (%0)" : : "r"(m_data));
+}
+
 // Interrupt handling
 
 #define FIXME_UNHANDLED_INTERRUPT(name)                                                                                \
diff --git a/kernel/src/arch/x86_64/CPU.h b/kernel/src/arch/x86_64/CPU.h
index 3d69c63d..a2efee77 100644
--- a/kernel/src/arch/x86_64/CPU.h
+++ b/kernel/src/arch/x86_64/CPU.h
@@ -9,6 +9,16 @@ struct Registers // Saved CPU registers for x86-64
     u64 rip, cs, rflags, rsp, ss;
 };
 
+struct FPData
+{
+    void save();
+    void restore();
+
+  private:
+    char m_data[512] alignas(16);
+    bool m_already_saved;
+};
+
 struct [[gnu::packed]] TSS
 {
     u32 reserved0;
diff --git a/kernel/src/thread/Scheduler.cpp b/kernel/src/thread/Scheduler.cpp
index 994f6436..286e1674 100644
--- a/kernel/src/thread/Scheduler.cpp
+++ b/kernel/src/thread/Scheduler.cpp
@@ -219,10 +219,12 @@ namespace Scheduler
         if (old_thread != new_thread)
         {
             switch_context(old_thread, new_thread, regs);
+            if (!old_thread->is_kernel) old_thread->fp_data.save();
             if (!new_thread->is_kernel)
             {
                 MMU::switch_page_directory(new_thread->directory);
                 CPU::switch_kernel_stack(new_thread->kernel_stack.top());
+                new_thread->fp_data.restore();
             }
         }
 
diff --git a/kernel/src/thread/Thread.h b/kernel/src/thread/Thread.h
index 9c902b67..1df015d5 100644
--- a/kernel/src/thread/Thread.h
+++ b/kernel/src/thread/Thread.h
@@ -21,8 +21,6 @@ enum class ThreadState
     Dying
 };
 
-// FIXME: Save floating point state. (SSE registers on x86_64 using FXSAVE and FXRSTOR)
-
 struct Thread : public LinkedListNode<Thread>
 {
     Registers regs;
@@ -41,6 +39,8 @@ struct Thread : public LinkedListNode<Thread>
 
     OwnedPtr<UserVM> vm_allocator;
 
+    FPData fp_data;
+
     ThreadState state = ThreadState::Runnable;
 
     bool is_kernel { true };