yggdrasil/kernel/arch/interface/src/task.rs

use core::fmt;

use alloc::boxed::Box;
use yggdrasil_abi::{arch::SavedFrame, error::Error, process::ExitCode};

use crate::mem::{KernelTableManager, PhysicalMemoryAllocator};

pub trait Scheduler {
    type ThreadId: Copy;

    fn for_cpu(index: usize) -> &'static Self;
    fn for_affinity_mask(mask: u64) -> &'static Self;
    fn local() -> &'static Self;

    fn is_local(&self) -> bool;
    fn push(&self, task: Self::ThreadId);

    /// Selects a new thread from the queue and performs a context switch if necessary.
    ///
    /// # Safety
    ///
    /// Only meant to be called from within the timer handler or the thread impl.
    unsafe fn yield_cpu(&self) -> bool;
}

/// Conversion trait to allow multiple kernel closure return types
pub trait Termination {
    /// Converts the closure return type into [ExitCode]
    fn into_exit_code(self) -> ExitCode;
}

/// Interface for task state save/restore mechanisms
pub trait TaskFrame {
    /// Creates a "snapshot" of a exception/syscall frame
    fn store(&self) -> SavedFrame;

    /// Restores the exception/syscall frame from its saved state
    fn restore(&mut self, saved: &SavedFrame);

    /// Replaces the return value in the frame (or does nothing, if the frame is not a part of a
    /// syscall signal handler)
    fn set_return_value(&mut self, value: u64);

    /// Replaces the userspace stack pointer in the frame
    fn set_user_sp(&mut self, value: usize);

    /// Replaces the userspace instruction pointer in the frame
    fn set_user_ip(&mut self, value: usize);

    /// Replaces the argument in the frame
    fn set_argument(&mut self, value: u64);

    /// Returns the argument (if any) of the frame being processed
    fn argument(&self) -> u64;

    /// Returns the userspace stack pointer
    fn user_sp(&self) -> usize;
    /// Returns the userspace instruction pointer
    fn user_ip(&self) -> usize;

    fn set_single_step(&mut self, step: bool);
}

/// Interface for performing context fork operations
pub trait ForkFrame<K: KernelTableManager, PA: PhysicalMemoryAllocator>: Sized {
    type Context: TaskContext<K, PA>;

    /// Constructs a "forked" task context by copying the registers from this one and supplying a
    /// new address space to it.
    ///
    /// # Safety
    ///
    /// Unsafe: accepts raw frames and address space address.
    unsafe fn fork(&self, address_space: u64) -> Result<Self::Context, Error>;

    /// Replaces the return value inside the frame with a new one
    fn set_return_value(&mut self, value: u64);
}

pub struct UserContextInfo {
    pub entry: usize,
    pub argument: usize,
    pub stack_pointer: usize,
    pub thread_pointer: usize,
    pub address_space: u64,
    pub asid: u64,
    pub single_step: bool,
}

/// Platform-specific task context implementation
pub trait TaskContext<K: KernelTableManager, PA: PhysicalMemoryAllocator>: Sized {
    /// Number of bytes to offset the signal stack pointer by
    const SIGNAL_STACK_EXTRA_ALIGN: usize;
    /// Number of bytes to offset the user stack pointer by
    const USER_STACK_EXTRA_ALIGN: usize;

    /// Constructs a kernel-space task context
    fn kernel(entry: extern "C" fn(usize) -> !, arg: usize) -> Result<Self, Error>;

    /// Constructs a user thread context. The caller is responsible for allocating the userspace
    /// stack and setting up a valid address space for the context.
    fn user(context: UserContextInfo) -> Result<Self, Error>;

    /// Performs an entry into a context.
    ///
    /// # Safety
    ///
    /// Only meant to be called from the scheduler code.
    unsafe fn enter(&self) -> !;

    /// Performs a context switch between two contexts.
    ///
    /// # Safety
    ///
    /// Only meant to be called from the scheduler code.
    unsafe fn switch(&self, from: &Self);

    /// Performs a context switch and drops the source thread.
    ///
    /// # Safety
    ///
    /// Only meant to be called from the scheduler code after the `thread` has terminated.
    unsafe fn switch_and_drop(&self, thread: *const ());

    /// Replaces the current thread pointer with the provided one.
    fn set_thread_pointer(&self, tp: usize);

    // XXX
    /// Constructs a safe wrapper process to execute a kernel-space closure
    fn kernel_closure<F: FnOnce() -> ! + Send + 'static>(f: F) -> Result<Self, Error> {
        extern "C" fn closure_wrapper<F: FnOnce() -> ! + Send + 'static>(
            mut closure_addr: usize,
        ) -> ! {
            if closure_addr == 0 {
                closure_addr = align_of::<F>();
            }
            let closure = unsafe { Box::from_raw(closure_addr as *mut F) };
            closure()
        }

        let closure = Box::new(f);
        let ptr = Box::into_raw(closure) as usize;
        Self::kernel(closure_wrapper::<F>, ptr)
    }

    fn align_stack_for_entry(sp: usize) -> usize {
        sp
    }
}

pub struct StackBuilder {
    base: usize,
    sp: usize,
}

impl StackBuilder {
    pub fn new(base: usize, size: usize) -> Self {
        Self {
            base,
            sp: base + size,
        }
    }

    pub fn push(&mut self, value: usize) {
        if self.sp == self.base {
            panic!();
        }
        self.sp -= size_of::<usize>();
        unsafe {
            (self.sp as *mut usize).write_volatile(value);
        }
    }

    pub fn build(self) -> usize {
        self.sp
    }
}

impl<T, E: fmt::Debug> Termination for Result<T, E> {
    fn into_exit_code(self) -> ExitCode {
        match self {
            Ok(_) => ExitCode::SUCCESS,
            Err(_err) => {
                // XXX
                // log::warn!("Kernel thread failed: {:?}", err);
                ExitCode::Exited(1)
            }
        }
    }
}

impl Termination for ExitCode {
    fn into_exit_code(self) -> ExitCode {
        self
    }
}

impl Termination for () {
    fn into_exit_code(self) -> ExitCode {
        ExitCode::SUCCESS
    }
}