Refactor and rewrite

src/multitasking/context_switch.rs

@@ -0,0 +1,93 @@
+use alloc::boxed::Box;
+use x86_64::VirtAddr;
+use core::raw::TraitObject;
+use crate::multitasking::thread::ThreadId;
+use core::mem;
+use super::with_scheduler;
+
+pub struct Stack {
+    pointer: VirtAddr,
+}
+
+impl Stack {
+    pub unsafe fn new(stack_pointer: VirtAddr) -> Self {
+        Stack { pointer: stack_pointer, }
+    }
+
+    pub fn get_stack_pointer(self) -> VirtAddr {
+        self.pointer
+    }
+
+    pub fn set_up_for_closure(&mut self, closure: Box<dyn FnOnce() -> !>) {
+        let trait_object: TraitObject = unsafe { mem::transmute(closure) };
+        unsafe { self.push(trait_object.data) };
+        unsafe { self.push(trait_object.vtable) };
+    
+        self.set_up_for_entry_point(call_closure_entry);
+    }
+
+    pub fn set_up_for_entry_point(&mut self, entry_point: fn() -> !) {
+        unsafe { self.push(entry_point) };
+        let rflags: u64 = 0x200;
+        unsafe { self.push(rflags) };
+    }
+
+    unsafe fn push<T>(&mut self, value: T) {
+        self.pointer -= core::mem::size_of::<T>();
+        let ptr: *mut T = self.pointer.as_mut_ptr();
+        ptr.write(value);
+    }
+}
+
+pub unsafe fn context_switch_to(thread_id: ThreadId, stack_pointer: VirtAddr) {
+    asm!(
+        "call asm_context_switch"
+        :
+        : "{rdi}"(stack_pointer), "{rsi}"(thread_id)
+        : "rax", "rbx", "rcx", "rdx", "rsi", "rdi", "rpb", "r8", "r9", "r10",
+        "r11", "r12", "r13", "r14", "r15", "rflags", "memory"
+        : "intel", "volatile"
+    );
+}
+
+global_asm!(
+    "
+    .intel_syntax noprefix
+
+    // asm_context_switch(stack_pointer: u64, thread_id: u64)
+    asm_context_switch:
+        pushfq
+
+        mov rax, rsp
+        mov rsp, rdi
+
+        mov rdi, rax
+        call add_paused_thread
+
+        popfq
+        ret
+"
+);
+
+#[no_mangle]
+pub extern "C" fn add_paused_thread(paused_stack_pointer: VirtAddr, new_thread_id: ThreadId) {
+    with_scheduler(|s| s.add_paused_thread(paused_stack_pointer, new_thread_id));
+}
+
+#[naked]
+fn call_closure_entry() -> ! {
+    unsafe { asm!("
+        pop rsi
+        pop rdi
+        call call_closure
+    " ::: "mem" : "intel", "volatile") };
+    unreachable!();
+}
+
+// no_mangle required because of https://github.com/rust-lang/rust/issues/68136
+#[no_mangle]
+extern "C" fn call_closure(data: *mut (), vtable: *mut ()) -> ! {
+    let trait_object = TraitObject { data, vtable };
+    let f: Box<dyn FnOnce() -> !> = unsafe { mem::transmute(trait_object) };
+    f()
+}

src/multitasking/mod.rs

@@ -0,0 +1,20 @@
+use scheduler::Scheduler;
+
+pub mod thread;
+pub mod scheduler;
+pub mod context_switch;
+
+static SCHEDULER: spin::Mutex<Option<Scheduler>> = spin::Mutex::new(None);
+
+pub fn invoke_scheduler() {
+    let next = SCHEDULER
+        .try_lock()
+        .and_then(|mut scheduler| scheduler.as_mut().and_then(|s| s.schedule()));
+    if let Some((next_id, next_stack_pointer)) = next {
+        unsafe { context_switch::context_switch_to(next_id, next_stack_pointer) };
+    }
+}
+
+pub fn with_scheduler<F, T>(f: F) -> T where F: FnOnce(&mut Scheduler) -> T {
+    f(SCHEDULER.lock().get_or_insert_with(Scheduler::new))
+}

src/multitasking/scheduler.rs

@@ -0,0 +1,57 @@
+use crate::multitasking::thread::{Thread, ThreadId};
+use alloc::collections::{VecDeque, BTreeMap};
+use x86_64::VirtAddr;
+use core::mem;
+
+pub struct Scheduler {
+    threads: BTreeMap<ThreadId, Thread>,
+    current_thread_id: ThreadId,
+    paused_threads: VecDeque<ThreadId>,
+}
+
+impl Scheduler {
+    pub fn new() -> Self {
+        let root_thread = Thread::create_root_thread();
+        let root_id = root_thread.id();
+        let mut threads = BTreeMap::new();
+        threads.insert(root_id, root_thread).expect_none("map is not empty after creation");
+        Scheduler {
+            threads,
+            current_thread_id: root_id,
+            paused_threads: VecDeque::new(),
+        }
+    }
+
+    fn next_thread(&mut self) -> Option<ThreadId> {
+        self.paused_threads.pop_front()
+    }
+
+    pub fn schedule(&mut self) -> Option<(ThreadId, VirtAddr)> {
+        if let Some(next_id) = self.next_thread() {
+            let next_thread = self.threads.get_mut(&next_id).expect("next thread does not exist");
+            let next_stack_pointer = next_thread.stack_pointer().take()
+                .expect("paused thread has no stack pointer");
+            Some((next_id, next_stack_pointer))
+        } else {
+            None
+        }
+    }
+
+    pub(super) fn add_paused_thread(&mut self, paused_stack_pointer: VirtAddr, next_thread_id: ThreadId) {
+        let paused_thread_id = mem::replace(&mut self.current_thread_id, next_thread_id);
+        let paused_thread = self.threads.get_mut(&paused_thread_id).expect("paused thread does not exist");
+        paused_thread.stack_pointer().replace(paused_stack_pointer)
+            .expect_none("running thread should have stack pointer set to None");
+        self.paused_threads.push_back(paused_thread_id);
+    }
+
+    pub fn add_new_thread(&mut self, thread: Thread) {
+        let thread_id = thread.id();
+        self.threads.insert(thread_id, thread).expect_none("thread already exists");
+        self.paused_threads.push_back(thread_id);
+    }
+
+    pub fn current_thread_id(&self) -> ThreadId {
+        self.current_thread_id
+    }
+}

src/multitasking/thread.rs

@@ -0,0 +1,76 @@
+use x86_64::{VirtAddr, structures::paging::{Mapper, mapper, FrameAllocator, Size4KiB}};
+use alloc::boxed::Box;
+use crate::memory::{alloc_stack, StackBounds};
+use crate::multitasking::context_switch::Stack;
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
+pub struct ThreadId(u64);
+
+impl ThreadId {
+    pub fn as_u64(&self) -> u64 {
+        self.0
+    }
+    
+    fn new() -> Self {
+        use core::sync::atomic::{AtomicU64, Ordering};
+        static NEXT_THREAD_ID: AtomicU64 = AtomicU64::new(1);
+        ThreadId(NEXT_THREAD_ID.fetch_add(1, Ordering::SeqCst))
+    }
+}
+
+#[derive(Debug)]
+pub struct Thread {
+    id: ThreadId,
+    stack_pointer: Option<VirtAddr>,
+    stack_bounds: Option<StackBounds>,
+}
+
+impl Thread {
+    pub fn create(
+        entry_point: fn() -> !,
+        stack_size: u64,
+        mapper: &mut impl Mapper<Size4KiB>,
+        frame_allocator: &mut impl FrameAllocator<Size4KiB>,
+    ) -> Result<Self, mapper::MapToError> {
+        let stack_bounds = alloc_stack(stack_size, mapper, frame_allocator)?;        
+        let mut stack = unsafe { Stack::new(stack_bounds.end()) };
+        stack.set_up_for_entry_point(entry_point);
+        Ok(Self::new(stack.get_stack_pointer(), stack_bounds))
+    }
+
+    pub fn create_from_closure<F>(
+        closure: F,
+        stack_size: u64,
+        mapper: &mut impl Mapper<Size4KiB>,
+        frame_allocator: &mut impl FrameAllocator<Size4KiB>,
+    ) -> Result<Self, mapper::MapToError> where F: FnOnce() -> ! + 'static + Send + Sync {
+        let stack_bounds = alloc_stack(stack_size, mapper, frame_allocator)?;        
+        let mut stack = unsafe { Stack::new(stack_bounds.end()) };
+        stack.set_up_for_closure(Box::new(closure));
+        Ok(Self::new(stack.get_stack_pointer(), stack_bounds))
+    }
+    
+    fn new(stack_pointer: VirtAddr, stack_bounds: StackBounds) -> Self {
+        Thread {
+            id: ThreadId::new(),
+            stack_pointer: Some(stack_pointer),
+            stack_bounds: Some(stack_bounds),
+        }
+    }
+    
+    pub(super) fn create_root_thread() -> Self {
+        Thread {
+            id: ThreadId(0),
+            stack_pointer: None,
+            stack_bounds: None,
+        }
+    }
+
+    pub fn id(&self) -> ThreadId {
+        self.id
+    }
+
+    pub(super) fn stack_pointer(&mut self) -> &mut Option<VirtAddr> {
+        &mut self.stack_pointer
+    }
+}