// os/src/task/context.rs

pub struct TaskContext {
    ra: usize,
    sp: usize,
    s: [usize; 12],
}

// os/src/task/switch.rs

global_asm!(include_str!("switch.S"));

use super::TaskContext;

extern "C" {
    pub fn __switch(
        current_task_cx_ptr: *mut TaskContext,
        next_task_cx_ptr: *const TaskContext
    );
}

let current_task_cx_ptr = &mut inner.tasks[current].task_cx as *mut TaskContext;
let next_task_cx_ptr = &inner.tasks[next].task_cx as *const TaskContext;
drop(inner);
// before this, we should drop local variables that must be dropped manually
unsafe {
	__switch(
		current_task_cx_ptr,
		next_task_cx_ptr,
	);
}rust

pub fn sys_yield() -> isize {
    syscall(SYSCALL_YIELD, [0, 0, 0])
}

// user/src/lib.rs

pub fn yield_() -> isize { sys_yield() }

const SBI_SET_TIMER: usize = 0;

pub fn set_timer(timer: usize) {
    sbi_call(SBI_SET_TIMER, timer, 0, 0);
}

// os/src/timer.rs

use crate::config::CLOCK_FREQ;
const TICKS_PER_SEC: usize = 100;

pub fn set_next_trigger() {
    set_timer(get_time() + CLOCK_FREQ / TICKS_PER_SEC);
}

// os/src/trap/mod.rs

match scause.cause() {
    Trap::Interrupt(Interrupt::SupervisorTimer) => {
        set_next_trigger();
        suspend_current_and_run_next();
    }
}

// enable supervisor clock
// os/src/trap/mod.rs

use riscv::register::sie;

pub fn enable_timer_interrupt() {
    unsafe { sie::set_stimer(); }
}

// stack struct ...

// buttom to top
fn get_sp(&self) -> usize {
    self.data.as_ptr() as usize + KERNEL_STACK_SIZE
}
pub fn push_context(&self, cx: TrapContext) -> &'static mut TrapContext {
    let cx_ptr = (self.get_sp() - core::mem::size_of::<TrapContext>()) as *mut TrapContext;
    unsafe {
        *cx_ptr = cx;
    }
    unsafe { cx_ptr.as_mut().unwrap() }
}

// os/src/trap/context.rs

#[repr(C)]
pub struct TrapContext {
    pub x: [usize; 32], // General register
    pub sstatus: Sstatus,
    pub sepc: usize,
}
/// set stack pointer to x_2 reg (sp)
pub fn set_sp(&mut self, sp: usize) {
    self.x[2] = sp;
}
/// init app context
pub fn app_init_context(entry: usize, sp: usize) -> Self {
    let mut sstatus = sstatus::read(); // CSR sstatus
    sstatus.set_spp(SPP::User); //previous privilege mode: user mode
    let mut cx = Self {
        x: [0; 32],
        sstatus,
        sepc: entry, // entry point of app
    };
    cx.set_sp(sp); // app's user stack pointer
    cx // return initial Trap Context of app
}

// os/Cargo.toml

[dependencies]
riscv = { git = "https://github.com/rcore-os/riscv", features = ["inline-asm"] }

// os/src/trap/mod.rs
global_asm!(include_str!("trap.S"));

pub fn init() {
    extern "C" { fn __alltraps(); }
    // write to stvec
    unsafe {
        stvec::write(__alltraps as usize, TrapMode::Direct);
    }
}

#[no_mangle]
pub fn trap_handler(cx: &mut TrapContext) -> &mut TrapContext {
    ...
}

extern "C" { fn __restore(cx_addr: usize); }
unsafe {
    __restore(KERNEL_STACK.push_context(
        TrapContext::app_init_context(APP_BASE_ADDRESS, USER_STACK.get_sp())
// This context store the ptr to UserStack for restoration
    ) as *const _ as usize);
}

// os/src/batch.rs

struct AppManager {
    num_app: usize,
    current_app: usize,
    app_start: [usize; MAX_APP_NUM + 1],
}

// part of read in static init of AppManager
let num_app_ptr = _num_app as usize as *const usize;
let num_app = num_app_ptr.read_volatile();
let mut app_start: [usize; MAX_APP_NUM + 1] = [0; MAX_APP_NUM + 1];
let app_start_raw: &[usize] =  core::slice::from_raw_parts(
    num_app_ptr.add(1), num_app + 1
);
app_start[..=num_app].copy_from_slice(app_start_raw);

// part of code of copying to kernel
asm!("fence.i");
// clear app area
core::slice::from_raw_parts_mut(
    APP_BASE_ADDRESS as *mut u8,
    APP_SIZE_LIMIT
).fill(0);
let app_src = core::slice::from_raw_parts(
    self.app_start[app_id] as *const u8,
    self.app_start[app_id + 1] - self.app_start[app_id]
);
let app_dst = core::slice::from_raw_parts_mut(
    APP_BASE_ADDRESS as *mut u8,
    app_src.len()
);
app_dst.copy_from_slice(app_src);

// os/src/batch.rs

pub fn run_next_app() -> ! {
    let mut app_manager = APP_MANAGER.exclusive_access();
    let current_app = app_manager.get_current_app();
    unsafe {
        app_manager.load_app(current_app);
    }
    app_manager.move_to_next_app();
    drop(app_manager);
    // before this we have to drop local variables related to resources manually
    // and release the resources
    extern "C" { fn __restore(cx_addr: usize); }
    unsafe {
        __restore(KERNEL_STACK.push_context(
            TrapContext::app_init_context(APP_BASE_ADDRESS, USER_STACK.get_sp())
        ) as *const _ as usize);
    }
    panic!("Unreachable in batch::run_current_app!");
}

// main logic:
// os/src/main.rs

...above code
// load entry for trap
trap::init()
// load app
batch::run_next_app()

const PAGE_SIZE: usize = 4096
const PAGE_SIZE_BIT: usize = 12

// os/src/mm/address.rs

impl PhysAddr {
    pub fn floor(&self) -> PhysPageNum { PhysPageNum(self.0 / PAGE_SIZE) }
    pub fn ceil(&self) -> PhysPageNum { PhysPageNum((self.0 + PAGE_SIZE - 1) / PAGE_SIZE) }
}

pub fn ppn(&self) -> PhysPageNum {
	(self.bits >> 10 & ((1usize << 44) - 1)).into()
}
pub fn flags(&self) -> PTEFlags {
	PTEFlags::from_bits(self.bits as u8).unwrap()
}

pub struct StackFrameAllocator {
    current: usize,  //空闲内存的起始物理页号
    end: usize,      //空闲内存的结束物理页号
    recycled: Vec<usize>,
}

// os/src/mm/frame_allocator.rs

pub struct FrameTracker {
    pub ppn: PhysPageNum,
}

impl FrameTracker {
    pub fn new(ppn: PhysPageNum) -> Self {
        // page cleaning
        let bytes_array = ppn.get_bytes_array();
        for i in bytes_array {
            *i = 0;
        }
        Self { ppn }
    }
}

// os/src/mm/page_table.rs

pub struct PageTable {
    root_ppn: PhysPageNum,
    frames: Vec<FrameTracker>,
}

// for each idxs represent index slices of virtual addr.
for i in 0..3 {
	let pte = &mut ppn.get_pte_array()[idxs[i]];
	if i == 2 {
		result = Some(pte);
		break;
	}
	if !pte.is_valid() {
		let frame = frame_alloc().unwrap();
		// create or return None
		*pte = PageTableEntry::new(frame.ppn, PTEFlags::V);
		self.frames.push(frame);
	}
	ppn = pte.ppn();
}

// os/src/mm/memory_set.rs

pub struct MapArea {
    vpn_range: VPNRange,
    data_frames: BTreeMap<VirtPageNum, FrameTracker>,
    map_type: MapType,
    map_perm: MapPermission,
}

// os/src/mm/memory_set.rs

pub struct MemorySet {
    page_table: PageTable,
    areas: Vec<MapArea>,
}    

// impl MemorySet
fn push(&mut self, mut map_area: MapArea, data: Option<&[u8]>) {
	// map range of virtual addr in allocation in page_table
	map_area.map(&mut self.page_table);
	if let Some(data) = data {
		map_area.copy_data(&mut self.page_table, data);
	}
	self.areas.push(map_area);
}

// os/src/mm/page_table.rs

pub fn token(&self) -> usize {
    8usize << 60 | self.root_ppn.0
}

// os/src/mm/memory_set.rs

impl MemorySet {
    pub fn activate(&self) {
        let satp = self.page_table.token();
        unsafe {
            satp::write(satp);
            asm!("sfence.vma" :::: "volatile");
			// sfence.vma is a special instruction to clear `Translaton Lookaside Buffer` which is used for quick search for memory addr to reduce performance expenses.
        }
    }
}

let mut memory_set = Self::new_bare();
// map trampoline
memory_set.map_trampoline();
memory_set.push(MapArea::new(
	(stext as usize).into(),
	(etext as usize).into(),
	MapType::Identical,
	MapPermission::R | MapPermission::X,
), None);
println!("mapping .rodata section");
// other layout ...

let mut max_end_vpn = VirtPageNum(0);
let max_end_va: VirtAddr = max_end_vpn.into();
let mut user_stack_bottom: usize = max_end_va.into();
// guard page
user_stack_bottom += PAGE_SIZE;
let user_stack_top = user_stack_bottom + USER_STACK_SIZE;
memory_set.push(MapArea::new(
	TRAP_CONTEXT.into(),
	TRAMPOLINE.into(),
	MapType::Framed,
	MapPermission::R | MapPermission::W,
), None);
// ...

// os/src/trap/context.rs

#[repr(C)]
pub struct TrapContext {
    pub x: [usize; 32],
    pub sstatus: Sstatus,
    pub sepc: usize,
	// new:
    pub kernel_satp: usize, 
    pub kernel_sp: usize,
    pub trap_handler: usize,
}

#[no_mangle]
pub fn trap_handler() -> ! {
    set_kernel_trap_entry();
    let cx = current_trap_cx();
    let scause = scause::read();
    let stval = stval::read();
    match scause.cause() {
        ...
    }
    trap_return();
}

#[no_mangle]
pub fn trap_return() -> ! {
    set_user_trap_entry();
    let trap_cx_ptr = TRAP_CONTEXT;
    let user_satp = current_user_token();
    extern "C" {
        fn __alltraps();
        fn __restore();
    }
    let restore_va = __restore as usize - __alltraps as usize + TRAMPOLINE;
    unsafe {
        asm!(
            "fence.i",
            "jr {restore_va}",
            restore_va = in(reg) restore_va,
            in("a0") trap_cx_ptr,
            in("a1") user_satp,
            options(noreturn)
        );
    }
    panic!("Unreachable in back_to_user!");
}

// os/src/config.rs

pub const TRAMPOLINE: usize = usize::MAX - PAGE_SIZE + 1;

// os/src/mm/memory_set.rs

impl MemorySet {
    /// Mention that trampoline is not collected by areas.
    fn map_trampoline(&mut self) {
        self.page_table.map(
            VirtAddr::from(TRAMPOLINE).into(),
            PhysAddr::from(strampoline as usize).into(),
            PTEFlags::R | PTEFlags::X,
        );
    }
}

// os/src/task/task.rs

pub struct TaskControlBlock {
    pub task_cx: TaskContext,
    pub task_status: TaskStatus,
    pub memory_set: MemorySet,
    pub trap_cx_ppn: PhysPageNum,
    pub base_size: usize,
}

impl TaskContext {
    pub fn goto_trap_return() -> Self {
        Self {
            ra: trap_return as usize,
            s: [0; 12],
        }
    }
}

// fn new(elf_data: &[u8], app_id: usize) -> Self 
let task_control_block = Self {
	task_status,
	task_cx: TaskContext::goto_trap_return(kernel_stack_top),
	memory_set,
	trap_cx_ppn,
	base_size: user_sp,
};
// prepare TrapContext in user space
let trap_cx = task_control_block.get_trap_cx();
*trap_cx = TrapContext::app_init_context(
	entry_point,
	user_sp,
	KERNEL_SPACE.exclusive_access().token(),
	kernel_stack_top,
	trap_handler as usize,
);
task_control_block

#![no_std]
#![no_main]

const SBI_CONSOLE_PUTCHAR: usize = 1;

fn sbi_call(...) -> usize {
	...
}

pub fn console_putchar(c:usize) {
	sbi_call(SBI_CONSOLE_PUTCHAR,c,0,0)
}

struct Stdout;

impl Write for Stdout {
    fn write_str(&mut self, s: &str) -> fmt::Result {
        for c in s.chars() {
            console_putchar(c as usize);
        }
        Ok(())
    }
}

pub fn print(args: fmt::Arguments) {
    Stdout.write_fmt(args).unwrap();
}

pub struct PidHandle(pub usize);

struct PidAllocator {
    current: usize,
    recycled: Vec<usize>,
}

// impl PidAllocator
pub fn alloc(&mut self) -> PidHandle {
	if let Some(pid) = self.recycled.pop() {
		PidHandle(pid)
	} else {
		self.current += 1;
		PidHandle(self.current - 1)
	}
}

// os/src/task/pid.rs

pub struct KernelStack {
    pid: usize,
}

pub fn kernel_stack_position(app_id: usize) -> (usize, usize) {
    let top = TRAMPOLINE - app_id * (KERNEL_STACK_SIZE + PAGE_SIZE);
    let bottom = top - KERNEL_STACK_SIZE;
    (bottom, top)
}

// impl KernelStack
pub fn push_on_top<T>(&self, value: T) -> *mut T where
	T: Sized, {
	let kernel_stack_top = self.get_top();
	let ptr_mut = (kernel_stack_top - core::mem::size_of::<T>()) as *mut T;
	unsafe { *ptr_mut = value; }
	ptr_mut
}
pub fn get_top(&self) -> usize {
	let (_, kernel_stack_top) = kernel_stack_position(self.pid);
	kernel_stack_top
}

impl Drop for KernelStack {
    fn drop(&mut self) {
        let (kernel_stack_bottom, _) = kernel_stack_position(self.pid);
        let kernel_stack_bottom_va: VirtAddr = kernel_stack_bottom.into();
        KERNEL_SPACE
            .exclusive_access()
            .remove_area_with_start_vpn(kernel_stack_bottom_va.into());
    }
}

pub struct TaskControlBlockInner {
	...
	// new:
    pub parent: Option<Weak<TaskControlBlock>>,
    pub children: Vec<Arc<TaskControlBlock>>,
    pub exit_code: i32,
}

// impl TaskControlBlockInner
pub fn is_zombie(&self) -> bool {
	self.get_status() == TaskStatus::Zombie
}

// os/src/task/manager.rs

pub struct TaskManager {
    ready_queue: VecDeque<Arc<TaskControlBlock>>,
}

lazy_static! {
    pub static ref TASK_MANAGER: UPSafeCell<TaskManager> = unsafe {
        UPSafeCell::new(TaskManager::new())
    };
}

pub fn add_task(task: Arc<TaskControlBlock>) {
    TASK_MANAGER.exclusive_access().add(task);
}

pub fn fetch_task() -> Option<Arc<TaskControlBlock>> {
    TASK_MANAGER.exclusive_access().fetch()
}

pub struct Processor {
    current: Option<Arc<TaskControlBlock>>,
	// idle process of current cpu
    idle_task_cx: TaskContext,
}

lazy_static! {
    pub static ref PROCESSOR: UPSafeCell<Processor> = unsafe {
        UPSafeCell::new(Processor::new())
    };
}

// run_tasks() 
// loop to fetch task and switch possible task
if let Some(task) = fetch_task() {
    let idle_task_cx_ptr = processor.get_idle_task_cx_ptr();
    let mut task_inner = task.inner_exclusive_access();
    let next_task_cx_ptr = &task_inner.task_cx as *const TaskContext;
    task_inner.task_status = TaskStatus::Running;
    drop(task_inner);
    processor.current = Some(task);
    drop(processor);
    unsafe {
        __switch(
            idle_task_cx_ptr,
            next_task_cx_ptr,
        );
    }
}

// switch to idle process if one task run out of its time.
pub fn schedule(switched_task_cx_ptr: *mut TaskContext) {
    let mut processor = PROCESSOR.exclusive_access();
    let idle_task_cx_ptr = processor.get_idle_task_cx_ptr();
    drop(processor);
    unsafe {
        __switch(
            switched_task_cx_ptr,
            idle_task_cx_ptr,
        );
    }
}

// os/src/task/mod.rs

pub fn suspend_current_and_run_next() {
    let task = take_current_task().unwrap();

    // ---- access current TCB exclusively
    let mut task_inner = task.inner_exclusive_access();
    let task_cx_ptr = &mut task_inner.task_cx as *mut TaskContext;
    task_inner.task_status = TaskStatus::Ready;
    drop(task_inner);
    // ---- stop exclusively accessing current PCB

    // add to task deque
    add_task(task);
    // change current to idle process
    schedule(task_cx_ptr);
}

// fn trap_handler() -> !
Trap::Exception(Exception::UserEnvCall) => {
    // jump to next instruction anyway
    let mut cx = current_trap_cx();
    cx.sepc += 4;
    // syscall may create new process and change trap context.
    let result = syscall(cx.x[17], [cx.x[10], cx.x[11], cx.x[12]]);
    // wether cx is changed or not, we will refetch it.
    cx = current_trap_cx();
    cx.x[10] = result as usize;
}

// impl MemorySet
pub fn from_existed_user(user_space: &MemorySet) -> MemorySet {
    let mut memory_set = Self::new_bare();
    // map trampoline
    memory_set.map_trampoline();
    // copy data sections/trap_context/user_stack
    for area in user_space.areas.iter() {
        let new_area = MapArea::from_another(area);
        memory_set.push(new_area, None);
        // copy data from another space
        for vpn in area.vpn_range {
            let src_ppn = user_space.translate(vpn).unwrap().ppn();
            let dst_ppn = memory_set.translate(vpn).unwrap().ppn();
            dst_ppn.get_bytes_array().copy_from_slice(src_ppn.get_bytes_array());
        }
    }
    memory_set
}

// impl TaskControlBlock
// fn fork
let trap_cx_ppn = memory_set
    .translate(VirtAddr::from(TRAP_CONTEXT).into())
    .unwrap()
    .ppn();
// alloc a pid and a kernel stack in kernel space
let pid_handle = pid_alloc();
let kernel_stack = KernelStack::new(&pid_handle);
let kernel_stack_top = kernel_stack.get_top();
let task_control_block = Arc::new(TaskControlBlock {
    pid: pid_handle,
    kernel_stack,
    inner: unsafe { UPSafeCell::new(TaskControlBlockInner {
        trap_cx_ppn,
        base_size: parent_inner.base_size,
        task_cx: TaskContext::goto_trap_return(kernel_stack_top),
        task_status: TaskStatus::Ready,
        memory_set,
        parent: Some(Arc::downgrade(self)),
        children: Vec::new(),
        exit_code: 0,
    })},
});
// add child
parent_inner.children.push(task_control_block.clone());
// modify kernel_sp in trap_cx
// **** access children PCB exclusively
let trap_cx = task_control_block.inner_exclusive_access().get_trap_cx();
trap_cx.kernel_sp = kernel_stack_top;

pub fn sys_fork() -> isize {
    let current_task = current_task().unwrap();
    let new_task = current_task.fork();
    let new_pid = new_task.pid.0;
    let trap_cx = new_task.inner_exclusive_access().get_trap_cx();

    // for child process, fork returns 0
    trap_cx.x[10] = 0;  //x[10] is a0 reg

    add_task(new_task);
    new_pid as isize
}

// os/build.rs

writeln!(f, r#"
.global _app_names
_app_names:"#)?;
for app in apps.iter() {
    writeln!(f, r#"    .string "{}""#, app)?;
}

// os/src/loader.rs

// APP_NAMES: Vec<&'static str> {...
for _ in 0..num_app {
    let mut end = start;
    while end.read_volatile() != '\0' as u8 {
        end = end.add(1);
    }
    let slice = core::slice::from_raw_parts(start, end as usize - start as usize);
    let str = core::str::from_utf8(slice).unwrap();
    v.push(str);
    start = end.add(1);
}

// os/src/task/task.rs

impl TaskControlBlock {
    pub fn exec(&self, elf_data: &[u8]) {
        let (memory_set, user_sp, entry_point) = MemorySet::from_elf(elf_data);
        let trap_cx_ppn = memory_set
            .translate(VirtAddr::from(TRAP_CONTEXT).into())
            .unwrap()
            .ppn();

        // **** access inner exclusively
        let mut inner = self.inner_exclusive_access();
        inner.memory_set = memory_set;
        inner.trap_cx_ppn = trap_cx_ppn;

        let trap_cx = inner.get_trap_cx();
        *trap_cx = TrapContext::app_init_context(
            entry_point,
            user_sp,
            KERNEL_SPACE.exclusive_access().token(),
            self.kernel_stack.get_top(),
            trap_handler as usize,
        );
        // **** stop exclusively accessing inner automatically
    }
}

// os/src/syscall/process.rs

pub fn sys_exec(path: *const u8) -> isize {
    let token = current_user_token();
    let path = translated_str(token, path);
    if let Some(data) = get_app_data_by_name(path.as_str()) {
        let task = current_task().unwrap();
        task.exec(data);
        0
    } else {
        -1
    }
}

// fn exit_current_and_run_next(exit_code:i32)
    inner.task_status = TaskStatus::Zombie;
    inner.exit_code = exit_code;

    // move all its children to the initial process
    // ++++++ access initproc TCB exclusively
    {
        let mut initproc_inner = INITPROC.inner_exclusive_access();
        for child in inner.children.iter() {
            child.inner_exclusive_access().parent = Some(Arc::downgrade(&INITPROC));
            initproc_inner.children.push(child.clone());
        }
    }
    // ++++++ stop exclusively accessing parent PCB

    // clear all memory.
    inner.children.clear();
    inner.memory_set.recycle_data_pages();
    drop(inner);
    // **** stop exclusively accessing current PCB
    // drop task manually to maintain rc correctly
    drop(task);
    // use _unused replace original context, which will be recycled by rust.
    let mut _unused = TaskContext::zero_init();
    schedule(&mut _unused as *mut _);

// os/src/syscall/process.rs

pub fn sys_exit(exit_code: i32) -> ! {
    exit_current_and_run_next(exit_code);
    panic!("Unreachable in sys_exit!");
}

// sys_waitpid(pid:uisze, exit_code_ptr:*mut i32) -> isize

// -1 case
// search task manager and find (task_block)
|p| {pid == -1 || pid as usize == p.getpid()}
return -1;

let pair = // search task managers and find (idx,task_block)
p.inner_exclusive_access().is_zombie() && (pid == -1 || pid as usize == p.getpid())

if let Some((idx,_)) = pair {
    let child = inner.children.remove(idx);
    // confirm that child will be deallocated after removing from children list
    assert_eq!(Arc::strong_count(&child), 1);
    let found_pid = child.getpid();
    // ++++ temporarily access child TCB exclusively
    let exit_code = child.inner_exclusive_access().exit_code;
    // ++++ stop exclusively accessing child PCB
    *translated_refmut(inner.memory_set.token(), exit_code_ptr) = exit_code;
    found_pid as isize
} else {
    // pid process is running
    -2
}

// user/src/lib.rs

pub fn wait(exit_code: &mut i32) -> isize {
    loop {
        match sys_waitpid(-1, exit_code as *mut _) {
            -2 => { yield_(); }
            // -1 or a real pid
            exit_pid => return exit_pid,
        }
    }
}