yggdrasil/kernel/arch/riscv64/src/context.rs

use core::{arch::global_asm, cell::UnsafeCell, marker::PhantomData};

use kernel_arch_interface::{
    mem::{KernelTableManager, PhysicalMemoryAllocator},
    task::{StackBuilder, TaskContext, UserContextInfo},
    Architecture,
};
use libk_mm_interface::address::PhysicalAddress;
use tock_registers::{
    interfaces::{Readable, Writeable},
    registers::InMemoryRegister,
};
use yggdrasil_abi::error::Error;

use crate::{
    mem::{self, KERNEL_VIRT_OFFSET},
    registers::SATP,
    ArchitectureImpl, PerCpuData,
};

pub const CONTEXT_SIZE: usize = 14 * size_of::<usize>();

#[repr(C, align(0x10))]
struct TaskContextInner {
    // 0x00
    sp: usize,

    satp: InMemoryRegister<u64, SATP::Register>,
}

pub struct TaskContextImpl<
    K: KernelTableManager,
    PA: PhysicalMemoryAllocator<Address = PhysicalAddress>,
> {
    inner: UnsafeCell<TaskContextInner>,
    // fp_context: UnsafeCell<FpContext>,
    stack_base_phys: PhysicalAddress,
    stack_top: usize,
    stack_size: usize,

    _pd: PhantomData<(K, PA)>,
}

impl<K: KernelTableManager, PA: PhysicalMemoryAllocator<Address = PhysicalAddress>>
    TaskContextImpl<K, PA>
{
    unsafe fn load_state(&self) {
        // TODO load new SATP value
        let inner = unsafe { &*self.inner.get() };
        let cpu = unsafe { &mut *ArchitectureImpl::local_cpu().cast::<PerCpuData>() };

        // Copy new SATP
        let satp = inner.satp.get();
        let asid = inner.satp.read(SATP::ASID);
        if satp != SATP.get() {
            mem::tlb_flush_asid(asid as usize);
            SATP.set(satp);
        }
        cpu.smode_sp = self.stack_top;
    }

    unsafe fn store_state(&self) {}
}

impl<K: KernelTableManager, PA: PhysicalMemoryAllocator<Address = PhysicalAddress>>
    TaskContext<K, PA> for TaskContextImpl<K, PA>
{
    const USER_STACK_EXTRA_ALIGN: usize = 8;
    const SIGNAL_STACK_EXTRA_ALIGN: usize = 0;

    fn user(context: UserContextInfo) -> Result<Self, Error> {
        const USER_TASK_PAGES: usize = 16;
        let stack_base_phys = PA::allocate_contiguous_pages(USER_TASK_PAGES)?;
        let stack_base = stack_base_phys.raw_virtualize::<K>();

        let mut stack = StackBuilder::new(stack_base, USER_TASK_PAGES * 0x1000);

        log::debug!(
            "Set up user task: pc={:#x}, sp={:#x}, tp={:#x}",
            context.entry,
            context.stack_pointer,
            context.thread_pointer
        );
        stack.push(context.stack_pointer);
        stack.push(context.thread_pointer);
        stack.push(context.entry);
        stack.push(context.argument);

        setup_common_context(&mut stack, __rv64_task_enter_user as _);

        let sp = stack.build();
        let satp = InMemoryRegister::new(0);
        satp.write(
            SATP::MODE::Sv39
                + SATP::ASID.val(context.asid)
                + SATP::PPN.val(context.address_space >> 12),
        );

        Ok(Self {
            inner: UnsafeCell::new(TaskContextInner { sp, satp }),
            // fp_context: UnsafeCell::new(FpContext::new()),
            stack_base_phys,
            stack_top: stack_base + USER_TASK_PAGES * 0x1000,
            stack_size: USER_TASK_PAGES * 0x1000,

            _pd: PhantomData,
        })
    }

    fn kernel(entry: extern "C" fn(usize) -> !, arg: usize) -> Result<Self, Error> {
        const KERNEL_TASK_PAGES: usize = 8;
        let stack_base_phys = PA::allocate_contiguous_pages(KERNEL_TASK_PAGES)?;
        let stack_base = stack_base_phys.raw_virtualize::<K>();

        let mut stack = StackBuilder::new(stack_base, KERNEL_TASK_PAGES * 0x1000);

        // Entry and argument
        stack.push(entry as _);
        stack.push(arg);

        setup_common_context(&mut stack, __rv64_task_enter_kernel as _);

        let sp = stack.build();

        // TODO stack is leaked
        let satp = InMemoryRegister::new(0);
        let kernel_table_phys = ((&raw const mem::FIXED_L1).addr() - KERNEL_VIRT_OFFSET) as u64;
        satp.write(SATP::MODE::Sv39 + SATP::ASID.val(0) + SATP::PPN.val(kernel_table_phys >> 12));

        Ok(Self {
            inner: UnsafeCell::new(TaskContextInner { sp, satp }),
            // fp_context: UnsafeCell::new(FpContext::new()),
            stack_base_phys,
            stack_top: 0,
            stack_size: KERNEL_TASK_PAGES * 0x1000,

            _pd: PhantomData,
        })
    }

    fn set_thread_pointer(&self, tp: usize) {
        let _ = tp;
        todo!()
    }

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

    unsafe fn enter(&self) -> ! {
        unsafe {
            self.load_state();
            __rv64_enter_task(self.inner.get())
        }
    }

    unsafe fn switch(&self, from: &Self) {
        if core::ptr::addr_eq(self, from) {
            return;
        }

        unsafe {
            from.store_state();
            self.load_state();
            __rv64_switch_task(self.inner.get(), from.inner.get())
        }
    }

    unsafe fn switch_and_drop(&self, thread: *const ()) {
        unsafe {
            self.load_state();
            __rv64_switch_task_and_drop(self.inner.get(), thread)
        }
    }
}

impl<K: KernelTableManager, PA: PhysicalMemoryAllocator<Address = PhysicalAddress>> Drop
    for TaskContextImpl<K, PA>
{
    fn drop(&mut self) {
        assert_eq!(self.stack_size % 0x1000, 0);

        for offset in (0..self.stack_size).step_by(0x1000) {
            unsafe {
                PA::free_page(self.stack_base_phys.add(offset));
            }
        }
    }
}

fn setup_common_context(builder: &mut StackBuilder, entry: usize) {
    builder.push(0); // x8/s0/fp
    builder.push(0); // x9/s1
    builder.push(0); // x18/s2
    builder.push(0); // x19/s3
    builder.push(0); // x20/s4
    builder.push(0); // x21/s5
    builder.push(0); // x22/s6
    builder.push(0); // x23/s7
    builder.push(0); // x24/s8
    builder.push(0); // x25/s9
    builder.push(0); // x26/s10
    builder.push(0); // x27/s11
    builder.push(0); // x4/gp
    builder.push(entry); // x1/ra return address
}

unsafe extern "C" {
    fn __rv64_enter_task(to: *mut TaskContextInner) -> !;
    fn __rv64_switch_task(to: *mut TaskContextInner, from: *mut TaskContextInner);
    fn __rv64_switch_task_and_drop(to: *mut TaskContextInner, thread: *const ()) -> !;
    fn __rv64_task_enter_kernel();
    fn __rv64_task_enter_user();
    // fn __rv64_fp_store_context(to: *mut c_void);
    // fn __rv64_fp_restore_context(from: *const c_void);
}

global_asm!(
    include_str!("context.S"),
    context_size = const CONTEXT_SIZE,
);