From 563d97f37227283e3205921af997124a492cbcce Mon Sep 17 00:00:00 2001
From: Julien CLEMENT <julien.clement@epita.fr>
Date: Sat, 31 Dec 2022 13:50:24 +0100
Subject: [PATCH] add cleaner exit, freeing thread's stack

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

diff --git a/src/proc/scheduler/mod.rs b/src/proc/scheduler/mod.rs
index 17e4084..1c7ec52 100644
--- a/src/proc/scheduler/mod.rs
+++ b/src/proc/scheduler/mod.rs
@@ -31,6 +31,13 @@ impl ThreadStream {
             waker: AtomicWaker::new(),
         }
     }
+
+    pub fn register(&mut self, id: ThreadId) {
+        self.ids
+            .push(id)
+            .expect("Thread queue full");
+        self.waker.wake();
+    }
 }
 
 impl Stream for ThreadStream {
@@ -69,6 +76,7 @@ impl Scheduler {
             entry_point: 0,
             started: true,
             rsp: 0,
+            base_stack: 0,
         };
         res.register(Arc::new(RefCell::new(k_thread)));
         res
@@ -77,9 +85,11 @@ impl Scheduler {
 
     pub async fn run(&mut self) {
         while let Some(id) = self.thread_queue.next().await {
-            let thread = self.get_thread(id).unwrap();
-            unsafe {
-                (&mut*thread.as_ptr()).run();
+            if let Some(thread) = self.get_thread(id) { // Thread still exists
+                unsafe {
+                    (&mut*thread.as_ptr()).run();
+                }
+                self.thread_queue.register(id);
             }
         }
     }
@@ -89,11 +99,13 @@ impl Scheduler {
         if self.threads.insert(thread_id, thread).is_some() {
             panic!("Duplicate thread ID")
         }
-        self.thread_queue
-            .ids
-            .push(thread_id)
-            .expect("Thread queue full");
-        self.thread_queue.waker.wake();
+        if thread_id != ThreadId(0) {
+            self.thread_queue.register(thread_id);
+        }
+    }
+
+    pub fn exit(&mut self, id: ThreadId) {
+        self.threads.remove(&id).unwrap().borrow().exit();
     }
 
     pub fn get_thread(&mut self, id: ThreadId) -> Option<Threadt> {
diff --git a/src/proc/thread/mod.rs b/src/proc/thread/mod.rs
index 60f513b..29de36a 100644
--- a/src/proc/thread/mod.rs
+++ b/src/proc/thread/mod.rs
@@ -34,6 +34,7 @@ pub fn exit() {
             .get_thread(ThreadId(0))
             .unwrap()
             .as_ptr();
+            scheduler.exit(*RUNNING_THREAD.try_lock().unwrap());
     } // Drop scheduler mutex guard
 
     unsafe {
@@ -51,36 +52,62 @@ pub struct Thread {
     pub entry_point: u64,
     pub started: bool,
     pub rsp: u64,
+    pub base_stack: u64
 }
 
 impl Thread {
     pub fn new(entry_point: u64) -> Self {
         unsafe {
+            let stack_bottom = alloc(Layout::new::<[u8; STACK_SIZE]>()) as u64;
             Thread {
                 id: ThreadId::new(),
                 entry_point: entry_point,
                 started: false,
-                rsp: alloc(Layout::new::<[u8; STACK_SIZE]>()) as u64 + STACK_SIZE as u64,
+                rsp: stack_bottom + STACK_SIZE as u64,
+                base_stack: stack_bottom,
             }
         }
     }
 
+    pub fn exit(&self) {
+        unsafe {
+            dealloc(self.base_stack as *mut u8, Layout::new::<[u8; STACK_SIZE]>());
+        }
+    }
+
     pub fn run(&mut self) {
         println!("Running thread {:?}", self.id);
         unsafe {
             let mut current_thread_guard = RUNNING_THREAD.try_lock().unwrap();
-            let current_rsp: u64;
-            asm!(
-                "push rsp",    // Recover current rsp
-                "pop {out}",
-                "sub {out}, 56", // Offset to saved registers
-                out = out(reg) current_rsp, // Save thread rsp
-            );
 
             let mut scheduler = SCHEDULER.try_lock().unwrap();
-            // TODO: check if the thread still exists
-            let current_thread = scheduler.get_thread(*current_thread_guard).unwrap();
-            current_thread.borrow_mut().rsp = current_rsp;
+            if let Some(current_thread) = scheduler.get_thread(*current_thread_guard) {
+                let current_rsp: u64;
+                asm!(
+                    "push rsp",    // Recover current rsp
+                    "pop {out}",
+                    "sub {out}, 56", // Offset to saved registers
+                    out = out(reg) current_rsp, // Save thread rsp
+                );
+                current_thread.borrow_mut().rsp = current_rsp;
+            }
+            else { // Thread does not exists anymore
+                *current_thread_guard = self.id; // change running thread
+                asm!( // Just switch to new thead without saving registers
+                    "push {rsp}", // Set stack pointer to the new thread
+                    "pop rsp",
+
+                    "pop rdi",       // Restore new thread regs
+                    "pop rsi",
+                    "pop rbp",
+                    "pop rdx",
+                    "pop rcx",
+                    "pop rbx",
+                    "pop rax",
+                    rsp = in(reg) self.rsp,
+                );
+                return;
+            }
 
             *current_thread_guard = self.id; // change running thread
         } // The scheduler and running thread guards are dropped here