Move TaskContext into TCB instead of kstack.

os/src/loader.rs

@@ -28,14 +28,10 @@ impl KernelStack {
     fn get_sp(&self) -> usize {
         self.data.as_ptr() as usize + KERNEL_STACK_SIZE
     }
-    pub fn push_context(&self, trap_cx: TrapContext, task_cx: TaskContext) -> &'static mut TaskContext {
-        unsafe {
-            let trap_cx_ptr = (self.get_sp() - core::mem::size_of::<TrapContext>()) as *mut TrapContext;
-            *trap_cx_ptr = trap_cx;
-            let task_cx_ptr = (trap_cx_ptr as usize - core::mem::size_of::<TaskContext>()) as *mut TaskContext;
-            *task_cx_ptr = task_cx;
-            task_cx_ptr.as_mut().unwrap()
-        }
+    pub fn push_context(&self, trap_cx: TrapContext) -> usize {
+        let trap_cx_ptr = (self.get_sp() - core::mem::size_of::<TrapContext>()) as *mut TrapContext;
+        unsafe { *trap_cx_ptr = trap_cx; }
+        trap_cx_ptr as usize
     }
 }
 
@@ -81,9 +77,8 @@ pub fn load_apps() {
     }
 }
 
-pub fn init_app_cx(app_id: usize) -> &'static TaskContext {
+pub fn init_app_cx(app_id: usize) -> usize {
     KERNEL_STACK[app_id].push_context(
         TrapContext::app_init_context(get_base_i(app_id), USER_STACK[app_id].get_sp()),
-        TaskContext::goto_restore(),
     )
 }

os/src/task/context.rs

@@ -1,14 +1,24 @@
+#[derive(Copy, Clone)]
 #[repr(C)]
 pub struct TaskContext {
     ra: usize,
+    sp: usize,
     s: [usize; 12],
 }
 
 impl TaskContext {
-    pub fn goto_restore() -> Self {
+    pub fn zero_init() -> Self {
+        Self {
+            ra: 0,
+            sp: 0,
+            s: [0; 12],
+        }
+    }
+    pub fn goto_restore(kstack_ptr: usize) -> Self {
         extern "C" { fn __restore(); }
         Self {
             ra: __restore as usize,
+            sp: kstack_ptr,
             s: [0; 12],
         }
     }

os/src/task/mod.rs

@@ -25,11 +25,14 @@ lazy_static! {
     pub static ref TASK_MANAGER: TaskManager = {
         let num_app = get_num_app();
         let mut tasks = [
-            TaskControlBlock { task_cx_ptr: 0, task_status: TaskStatus::UnInit };
+            TaskControlBlock {
+                task_cx: TaskContext::zero_init(),
+                task_status: TaskStatus::UnInit 
+            };
             MAX_APP_NUM
         ];
         for i in 0..num_app {
-            tasks[i].task_cx_ptr = init_app_cx(i) as * const _ as usize;
+            tasks[i].task_cx = TaskContext::goto_restore(init_app_cx(i));            
             tasks[i].task_status = TaskStatus::Ready;
         }
         TaskManager {
@@ -46,12 +49,12 @@ impl TaskManager {
     fn run_first_task(&self) {
         let task0 = &mut self.inner.upsafe_access().tasks[0];
         task0.task_status = TaskStatus::Running;
-        let next_task_cx_ptr2 = task0.get_task_cx_ptr2();
-        let _unused: usize = 0;
+        let next_task_cx_ptr = &task0.task_cx as *const TaskContext; 
+        let mut _unused = TaskContext::zero_init();
         unsafe {
             __switch(
-                &_unused as *const _,
-                next_task_cx_ptr2,
+                &mut _unused as *mut TaskContext,
+                next_task_cx_ptr,
             );
         }
     }
@@ -84,13 +87,13 @@ impl TaskManager {
             let current = inner.current_task;
             inner.tasks[next].task_status = TaskStatus::Running;
             inner.current_task = next;
-            let current_task_cx_ptr2 = inner.tasks[current].get_task_cx_ptr2();
-            let next_task_cx_ptr2 = inner.tasks[next].get_task_cx_ptr2();
-            core::mem::drop(inner);
+            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);
             unsafe {
                 __switch(
-                    current_task_cx_ptr2,
-                    next_task_cx_ptr2,
+                    current_task_cx_ptr,
+                    next_task_cx_ptr,
                 );
             }
         } else {

os/src/task/switch.S

@@ -1,37 +1,34 @@
 .altmacro
 .macro SAVE_SN n
-    sd s\n, (\n+1)*8(sp)
+    sd s\n, (\n+2)*8(a0)
 .endm
 .macro LOAD_SN n
-    ld s\n, (\n+1)*8(sp)
+    ld s\n, (\n+2)*8(a1)
 .endm
     .section .text
     .globl __switch
 __switch:
     # __switch(
-    #     current_task_cx_ptr2: &*const TaskContext,
-    #     next_task_cx_ptr2: &*const TaskContext
+    #     current_task_cx_ptr: *mut TaskContext,
+    #     next_task_cx_ptr: *const TaskContext
     # )
-    # push TaskContext to current sp and save its address to where a0 points to
-    addi sp, sp, -13*8
-    sd sp, 0(a0)
-    # fill TaskContext with ra & s0-s11
-    sd ra, 0(sp)
+    # save kernel stack of current task
+    sd sp, 8(a0)
+    # save ra & s0~s11 of current execution
+    sd ra, 0(a0)
     .set n, 0
     .rept 12
         SAVE_SN %n
         .set n, n + 1
     .endr
-    # ready for loading TaskContext a1 points to
-    ld sp, 0(a1)
-    # load registers in the TaskContext
-    ld ra, 0(sp)
+    # restore ra & s0~s11 of next execution
+    ld ra, 0(a1)
     .set n, 0
     .rept 12
         LOAD_SN %n
         .set n, n + 1
     .endr
-    # pop TaskContext
-    addi sp, sp, 13*8
+    # restore kernel stack of next task
+    ld sp, 8(a1)
     ret
 

os/src/task/switch.rs

@@ -1,8 +1,10 @@
 global_asm!(include_str!("switch.S"));
 
+use super::TaskContext;
+
 extern "C" {
     pub fn __switch(
-        current_task_cx_ptr2: *const usize,
-        next_task_cx_ptr2: *const usize
+        current_task_cx_ptr: *mut TaskContext,
+        next_task_cx_ptr: *const TaskContext
     );
 }

os/src/task/task.rs

@@ -1,13 +1,9 @@
+use super::TaskContext;
+
 #[derive(Copy, Clone)]
 pub struct TaskControlBlock {
-    pub task_cx_ptr: usize,
     pub task_status: TaskStatus,
-}
-
-impl TaskControlBlock {
-    pub fn get_task_cx_ptr2(&self) -> *const usize {
-        &self.task_cx_ptr as *const usize
-    }
+    pub task_cx: TaskContext,
 }
 
 #[derive(Copy, Clone, PartialEq)]