    1
    2
    3
    4
    5
    6
    7
    8
    9
    10
    11
    12
    13
    14
    15
    16
    17
    18
    19
    20
    21
    22
    23
    24
    25
    26
    27
    28
    29
    30
    31
    32
    33
    34
    35
    36
    37
    38
    39
    40
    41
    42
    43
    44
    45
    46
    47
    48
    49
    50
    51
    52
    53
    54
    55
    56
    57
    58
    59
    60
    61
    62
    63
    64
    65
    66
    67
    68
    69
    70
    71
    72
    73
    74
    75
    76
    77

    use super::id::TaskUserRes;
    use super::{kstack_alloc, KernelStack, ProcessControlBlock, TaskContext};
    use crate::trap::TrapContext;
    use crate::{mm::PhysPageNum, sync::{UPIntrFreeCell, UPIntrRefMut}};
    use alloc::sync::{Arc, Weak};

    pub struct TaskControlBlock {
        // immutable
        pub process: Weak<ProcessControlBlock>,
        pub kstack: KernelStack,
        // mutable
        pub inner: UPIntrFreeCell<TaskControlBlockInner>,
    }

    impl TaskControlBlock {
        pub fn inner_exclusive_access(&self) -> UPIntrRefMut<'_, TaskControlBlockInner> {
            self.inner.exclusive_access()
        }

        pub fn get_user_token(&self) -> usize {
            let process = self.process.upgrade().unwrap();
            let inner = process.inner_exclusive_access();
            inner.memory_set.token()
        }
    }

    pub struct TaskControlBlockInner {
        pub res: Option<TaskUserRes>,
        pub trap_cx_ppn: PhysPageNum,
        pub task_cx: TaskContext,
        pub task_status: TaskStatus,
        pub exit_code: Option<i32>,
    }

    impl TaskControlBlockInner {
        pub fn get_trap_cx(&self) -> &'static mut TrapContext {
            self.trap_cx_ppn.get_mut()
        }

        #[allow(unused)]
        fn get_status(&self) -> TaskStatus {
            self.task_status
        }
    }

    impl TaskControlBlock {
        pub fn new(
            process: Arc<ProcessControlBlock>,
            ustack_base: usize,
            alloc_user_res: bool,
        ) -> Self {
            let res = TaskUserRes::new(Arc::clone(&process), ustack_base, alloc_user_res);
            let trap_cx_ppn = res.trap_cx_ppn();
            let kstack = kstack_alloc();
            let kstack_top = kstack.get_top();
            Self {
                process: Arc::downgrade(&process),
                kstack,
                inner: unsafe {
                    UPIntrFreeCell::new(TaskControlBlockInner {
                        res: Some(res),
                        trap_cx_ppn,
                        task_cx: TaskContext::goto_trap_return(kstack_top),
                        task_status: TaskStatus::Ready,
                        exit_code: None,
                    })
                },
            }
        }
    }

    #[derive(Copy, Clone, PartialEq)]
    pub enum TaskStatus {
        Ready,
        Running,
        Blocking,
    }