From af4ab10505605c0185a68fc53848c6ca03005fa0 Mon Sep 17 00:00:00 2001
From: Julien CLEMENT <julien.clement@epita.fr>
Date: Sat, 31 Dec 2022 12:23:27 +0100
Subject: [PATCH] Minor thread and scheduler refactoring

Signed-off-by: Julien CLEMENT <julien.clement@epita.fr>
---
 src/proc/scheduler/mod.rs |  8 ++++++++
 src/proc/thread/mod.rs    | 28 ++++++++++++++++------------
 2 files changed, 24 insertions(+), 12 deletions(-)

diff --git a/src/proc/scheduler/mod.rs b/src/proc/scheduler/mod.rs
index 46b69af..d621167 100644
--- a/src/proc/scheduler/mod.rs
+++ b/src/proc/scheduler/mod.rs
@@ -56,4 +56,12 @@ impl Scheduler {
             .push(thread_id)
             .expect("Thread queue full");
     }
+
+    pub fn get_thread(&mut self, id: ThreadId) -> Option<Threadt> {
+        if let Some(thread) = self.threads.get_mut(&id) {
+            Some(thread.clone())
+        } else {
+            None
+        }
+    }
 }
diff --git a/src/proc/thread/mod.rs b/src/proc/thread/mod.rs
index 77d3c1e..6d2570d 100644
--- a/src/proc/thread/mod.rs
+++ b/src/proc/thread/mod.rs
@@ -27,14 +27,18 @@ impl ThreadId {
 
 pub fn exit() {
     println!("Exiting thread");
-    let mut scheduler = SCHEDULER.try_lock().unwrap();
-    let mut thread = scheduler
-        .threads
-        .get_mut(&ThreadId(0))
-        .unwrap()
-        .borrow_mut();
-    SCHEDULER.force_unlock();
-    thread.run();
+    let thread: *mut Thread;
+    {
+        let mut scheduler = SCHEDULER.try_lock().unwrap();
+        thread = scheduler
+            .get_thread(ThreadId(0))
+            .unwrap()
+            .as_ptr();
+    } // Drop scheduler mutex guard
+
+    unsafe {
+        (&mut* thread).run();
+    }
 }
 
 pub struct Thread {
@@ -51,7 +55,7 @@ impl Thread {
                 id: ThreadId::new(),
                 entry_point: entry_point,
                 started: false,
-                rsp: alloc(Layout::new::<[u8; STACK_SIZE]>()) as u64 + STACK_SIZE as u64 - 0x80,
+                rsp: alloc(Layout::new::<[u8; STACK_SIZE]>()) as u64 + STACK_SIZE as u64,
             }
         }
     }
@@ -69,7 +73,7 @@ impl Thread {
             );
 
             let mut scheduler = SCHEDULER.try_lock().unwrap();
-            let current_thread = scheduler.threads.get_mut(&*current_thread_guard).unwrap();
+            let current_thread = scheduler.get_thread(*current_thread_guard).unwrap();
             current_thread.borrow_mut().rsp = current_rsp;
 
             *current_thread_guard = self.id; // change running thread
@@ -109,9 +113,9 @@ impl Thread {
                     "push rsi",
                     "push rdi",
 
-                    "push {rsp}",
+                    "push {rsp}",    // Set stack pointer to the new thread
                     "pop rsp",
-                    "jmp {rip}",
+                    "jmp {rip}",     // Jump to thread routine
                     rsp = in(reg) self.rsp,
                     rip = in(reg) self.entry_point,
                 );