diff --git a/src/interrupt.rs b/src/interrupt.rs
index 8aae2c1..a009811 100644
--- a/src/interrupt.rs
+++ b/src/interrupt.rs
@@ -1,19 +1,20 @@
-use log::info;
-
 use crate::{
-    boot::sbi::{EextensionID, TimerFunctionID}, clear_csr, generate_trap_handler, process::ExecutionContext, read_csr, riscv::disable_interrupt, set_csr, time::{IRQ_M_TIMER, setup_next_timer_interrupt}, write_csr
+    boot::sbi::{EextensionID, TimerFunctionID},
+    clear_csr, generate_trap_handler,
+    process::ExecutionContext,
+    read_csr,
+    riscv::disable_interrupt,
+    scheduler::scheduler,
+    set_csr,
+    time::{setup_next_timer_interrupt, IRQ_M_TIMER},
+    write_csr,
 };
 use core::arch::naked_asm;
 
 use crate::time::{setup_timer_interrupt, timer_interrupt};
 
 #[unsafe(no_mangle)]
-unsafe extern "C" fn machine_trap_handler(
-    _interrupt_state: *const ExecutionContext,
-    mcause: u64,
-    mie: u64,
-    mip: u64,
-) {
+unsafe extern "C" fn machine_trap_handler(mcause: u64, mie: u64, mip: u64) {
     let mepc = read_csr!(mepc);
     let mtval = read_csr!(mtval);
     if mcause & (1 << 63) == 0 {
@@ -89,7 +90,7 @@ unsafe extern "C" fn supervisor_trap_handler(
     scause: u64,
     _sie: u64,
     _sip: u64,
-) {
+) -> *const ExecutionContext {
     #[allow(clippy::single_match)]
     match scause & !(1 << 63) {
         5 => {
@@ -101,10 +102,12 @@ unsafe extern "C" fn supervisor_trap_handler(
                     in("a7") EextensionID::Time as u64,
                 );
             }
-            timer_interrupt(unsafe { *interrupt_state });
+            timer_interrupt();
+            return scheduler(unsafe { *interrupt_state });
         }
         _ => {}
     }
+    interrupt_state
 }
 
 pub unsafe fn setup_machine_trap_handler() {
@@ -116,8 +119,57 @@ pub unsafe fn setup_supervisor_trap_handler() {
     setup_timer_interrupt();
 }
 
-generate_trap_handler! {
-    _machine_mode_trap, machine_trap_handler, m
+#[unsafe(naked)]
+#[unsafe(no_mangle)]
+unsafe extern "C" fn _machine_mode_trap() {
+    naked_asm!(
+        "
+        addi sp, sp, -128
+
+        # Store the current frame
+        sd ra, 120(sp)
+        sd a0, 0(sp)
+        sd a1, 8(sp)
+        sd a2, 16(sp)
+        sd a3, 24(sp)
+        sd a4, 32(sp)
+        sd a5, 40(sp)
+        sd a6, 48(sp)
+        sd a7, 56(sp)
+        sd t0, 64(sp)
+        sd t1, 72(sp)
+        sd t2, 80(sp)
+        sd t3, 88(sp)
+        sd t4, 96(sp)
+        sd t5, 104(sp)
+        sd t6, 112(sp)
+
+        csrr a0, mcause
+        csrr a1, mie
+        csrr a2, mip
+        jal machine_trap_handler
+
+        # Restore registers
+        ld ra, 120(sp)
+        ld a0, 0(sp)
+        ld a1, 8(sp)
+        ld a2, 16(sp)
+        ld a3, 24(sp)
+        ld a4, 32(sp)
+        ld a5, 40(sp)
+        ld a6, 48(sp)
+        ld a7, 56(sp)
+        ld t0, 64(sp)
+        ld t1, 72(sp)
+        ld t2, 80(sp)
+        ld t3, 88(sp)
+        ld t4, 96(sp)
+        ld t5, 104(sp)
+        ld t6, 112(sp)
+
+        addi sp, sp, 128
+        mret"
+    )
 }
 generate_trap_handler! {
     _supervisor_mode_trap, supervisor_trap_handler, s
@@ -129,13 +181,14 @@ macro_rules! generate_trap_handler {
         #[unsafe(naked)]
         #[unsafe(no_mangle)]
         unsafe extern "C" fn $name() {
-            naked_asm!(
-                concat!("
+            naked_asm!(concat!(
+                "
+            mv t0, sp
             addi sp, sp, -264
 
             # Store the current frame
             sd ra, 0(sp)
-            sd sp, 8(sp)
+            sd t0, 8(sp) // sp
             sd gp, 16(sp)
             sd tp, 24(sp)
             sd a0, 32(sp)
@@ -170,92 +223,64 @@ macro_rules! generate_trap_handler {
             csrr t0, sstatus
             sd t0, 256(sp)
 
-            // sd ra, 120(sp)
-    
-            // sd a0, 0(sp)
-            // sd a1, 8(sp)
-            // sd a2, 16(sp)
-            // sd a3, 24(sp)
-            // sd a4, 32(sp)
-            // sd a5, 40(sp)
-            // sd a6, 48(sp)
-            // sd a7, 56(sp)
-
-            // sd t0, 64(sp)
-            // sd t1, 72(sp)
-            // sd t2, 80(sp)
-            // sd t3, 88(sp)
-            // sd t4, 96(sp)
-            // sd t5, 104(sp)
-            // sd t6, 112(sp)
-
             mv a0, sp
-            csrr a1, ", stringify!($mode),"cause
-            csrr a2, ", stringify!($mode),"ie
-            csrr a3, ", stringify!($mode),"ip
-            jal ", stringify!($jump_to), "
+            csrr a1, ",
+                stringify!($mode),
+                "cause
+            csrr a2, ",
+                stringify!($mode),
+                "ie
+            csrr a3, ",
+                stringify!($mode),
+                "ip
+            jal ",
+                stringify!($jump_to),
+                "
 
             # Restore registers
-            ld t0, 248(sp)
+            ld t0, 248(a0)
             csrw sepc, t0
-            ld t0, 256(sp)
-            csrw sstatus, t0
-            ld ra, 0(sp)
-            ld gp, 16(sp)
-            ld tp, 24(sp)
-            ld a0, 32(sp)
-            ld a1, 40(sp)
-            ld a2, 48(sp)
-            ld a3, 56(sp)
-            ld a4, 64(sp)
-            ld a5, 72(sp)
-            ld a6, 80(sp)
-            ld a7, 88(sp)
-            ld t0, 96(sp)
-            ld t1, 104(sp)
-            ld t2, 112(sp)
-            ld t3, 120(sp)
-            ld t4, 128(sp)
-            ld t5, 136(sp)
-            ld t6, 144(sp)
-            ld s0, 152(sp)
-            ld s1, 160(sp)
-            ld s2, 168(sp)
-            ld s3, 176(sp)
-            ld s4, 184(sp)
-            ld s5, 192(sp)
-            ld s6, 200(sp)
-            ld s7, 208(sp)
-            ld s8, 216(sp)
-            ld s9, 224(sp)
-            ld s10, 232(sp)
-            ld s11, 240(sp)
+            // ld t0, 256(a0)
+            // csrw sstatus, t0
+            ld ra, 0(a0)
+            ld sp, 8(a0)
+            ld gp, 16(a0)
+            ld tp, 24(a0)
+            ld a1, 40(a0)
+            ld a2, 48(a0)
+            ld a3, 56(a0)
+            ld a4, 64(a0)
+            ld a5, 72(a0)
+            ld a6, 80(a0)
+            ld a7, 88(a0)
+            ld t0, 96(a0)
+            ld t1, 104(a0)
+            ld t2, 112(a0)
+            ld t3, 120(a0)
+            ld t4, 128(a0)
+            ld t5, 136(a0)
+            ld t6, 144(a0)
+            ld s0, 152(a0)
+            ld s1, 160(a0)
+            ld s2, 168(a0)
+            ld s3, 176(a0)
+            ld s4, 184(a0)
+            ld s5, 192(a0)
+            ld s6, 200(a0)
+            ld s7, 208(a0)
+            ld s8, 216(a0)
+            ld s9, 224(a0)
+            ld s10, 232(a0)
+            ld s11, 240(a0)
 
-            ld sp, 8(sp)
+            ld a0, 32(a0)
 
-            // ld a0, 0(sp)
-            // ld a1, 8(sp)
-            // ld a2, 16(sp)
-            // ld a3, 24(sp)
-            // ld a4, 32(sp)
-            // ld a5, 40(sp)
-            // ld a6, 48(sp)
-            // ld a7, 56(sp)
+            // addi sp, sp, 264
 
-            // ld t0, 64(sp)
-            // ld t1, 72(sp)
-            // ld t2, 80(sp)
-            // ld t3, 88(sp)
-            // ld t4, 96(sp)
-            // ld t5, 104(sp)
-            // ld t6, 112(sp)
-
-            // ld ra, 120(sp)
-    
-            addi sp, sp, 264
-
-            ", stringify!($mode),"ret")
-            )
+            ",
+                stringify!($mode),
+                "ret"
+            ))
         }
     };
 }
diff --git a/src/main.rs b/src/main.rs
index 1a0d92b..f288aa8 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -12,7 +12,7 @@ use crate::{
     io::init_log,
     process::{create_processus, sleep},
     riscv::enable_supervisor_interrupt,
-    scheduler::scheduler_init,
+    scheduler::{idle, scheduler_init},
     vga::{Color, Vga},
 };
 
@@ -67,7 +67,6 @@ pub extern "C" fn supervisor_mode_entry() {
 
     create_processus(test, "proc1");
     create_processus(proc2, "proc2");
-    loop {
-        unsafe { wfi() }
-    }
+
+    idle();
 }
diff --git a/src/process.rs b/src/process.rs
index 9c26afc..1bdf074 100644
--- a/src/process.rs
+++ b/src/process.rs
@@ -3,7 +3,7 @@ use core::{arch::riscv64::wfi, time::Duration};
 use alloc::{format, string::String};
 
 use crate::{
-    scheduler::{scheduler, ACTIVE_PID, PROCESSUS_COUNT, PROCESS_TABLE},
+    scheduler::{ACTIVE_PID, PROCESSUS_COUNT, PROCESS_TABLE},
     time::elapsed_time_since_startup,
 };
 
@@ -26,8 +26,8 @@ pub struct ExecutionContext {
     pub tp: usize,
     pub a: [usize; 8],
     pub t: [usize; 7],
-    pub s: [usize; 11],
-    pub mepc: usize,
+    pub s: [usize; 12],
+    pub mepc: *const usize,
     pub mstatus: usize,
 }
 
@@ -73,7 +73,8 @@ pub fn create_processus<T: Into<String>>(code: extern "C" fn(), name: T) -> i64
         PROCESS_TABLE[next_pid].state = ProcessState::Activable;
         PROCESS_TABLE[next_pid].entry_point = Some(code);
         PROCESS_TABLE[next_pid].ctx.a[0] = code as usize;
-        PROCESS_TABLE[next_pid].ctx.ra = processus_launcher as *const _;
+        PROCESS_TABLE[next_pid].ctx.mepc = processus_launcher as *const _;
+        PROCESS_TABLE[next_pid].ctx.mstatus = 1 << 11;
         PROCESS_TABLE[next_pid].ctx.sp = &raw const PROCESS_TABLE[next_pid].stack[STACK_SIZE - 1];
     }
 
diff --git a/src/scheduler.rs b/src/scheduler.rs
index 2d34947..8b34cc0 100644
--- a/src/scheduler.rs
+++ b/src/scheduler.rs
@@ -1,16 +1,10 @@
-use core::{
-    arch::{naked_asm, riscv64::wfi},
-    array,
-    cell::LazyCell,
-    time::Duration,
-};
+use core::{arch::riscv64::wfi, array, cell::LazyCell, time::Duration};
 
 use alloc::string::String;
 use log::info;
 
 use crate::{
     process::{create_processus, ExecutionContext, Process, ProcessState},
-    riscv::enable_supervisor_interrupt,
     time,
 };
 
@@ -31,7 +25,7 @@ pub static mut PROCESS_TABLE: LazyCell<[Process; PROCESSUS_COUNT]> = LazyCell::n
             a: [0; _],
             t: [0; _],
             s: [0; _],
-            mepc: 0,
+            mepc: core::ptr::null(),
             mstatus: 0,
         },
         entry_point: None,
@@ -41,7 +35,7 @@ pub static mut PROCESS_TABLE: LazyCell<[Process; PROCESSUS_COUNT]> = LazyCell::n
 
 pub extern "C" fn idle() {
     loop {
-        enable_supervisor_interrupt();
+        // enable_supervisor_interrupt();
         unsafe {
             wfi();
         }
@@ -62,13 +56,14 @@ pub fn scheduler_init() {
     }
 }
 
-pub fn scheduler(interrupt_state: ExecutionContext) -> usize {
+pub fn scheduler(interrupt_state: ExecutionContext) -> *const ExecutionContext {
     // info!("scheduler");
     unsafe {
         let prev_pid = ACTIVE_PID;
+        PROCESS_TABLE[prev_pid].ctx = interrupt_state;
 
-        if PROCESS_TABLE[ACTIVE_PID].state == ProcessState::Active {
-            PROCESS_TABLE[ACTIVE_PID].state = ProcessState::Activable;
+        if PROCESS_TABLE[prev_pid].state == ProcessState::Active {
+            PROCESS_TABLE[prev_pid].state = ProcessState::Activable;
         }
 
         loop {
@@ -84,90 +79,6 @@ pub fn scheduler(interrupt_state: ExecutionContext) -> usize {
         }
         PROCESS_TABLE[ACTIVE_PID].state = ProcessState::Active;
 
-        PROCESS_TABLE[prev_pid].ctx = interrupt_state;
-        // PROCESS_TABLE[prev_pid].ctx.t = interrupt_state.t;
-        // PROCESS_TABLE[prev_pid].ctx.ra = interrupt_state.ra;
-        context_switch(
-            PROCESS_TABLE[ACTIVE_PID].entry_point.unwrap(),
-            &raw mut PROCESS_TABLE[prev_pid].ctx,
-            &raw mut PROCESS_TABLE[ACTIVE_PID].ctx,
-        );
-        prev_pid
+        &raw const PROCESS_TABLE[ACTIVE_PID].ctx
     }
 }
-
-#[unsafe(naked)]
-pub extern "C" fn context_switch(
-    code: extern "C" fn(),
-    current: *mut ExecutionContext,
-    next: *mut ExecutionContext,
-) {
-    naked_asm!(
-        "
-    // sd ra, 0(a1)
-    // sd sp, 8(a1)
-    // sd gp, 16(a1)
-    // sd tp, 24(a1)
-    // sd s0, 152(a1)
-    // sd s1, 160(a1)
-    // sd s2, 168(a1)
-    // sd s3, 176(a1)
-    // sd s4, 184(a1)
-    // sd s5, 192(a1)
-    // sd s6, 200(a1)
-    // sd s7, 208(a1)
-    // sd s8, 216(a1)
-    // sd s9, 224(a1)
-    // sd s10, 232(a1)
-    // sd s11, 240(a1)
-    // csrr t0, sepc
-    // sd t0, 248(a1)
-    // csrr t0, sstatus
-    // sd t0, 256(a1)
-
-    # Load next execution context
-    ld t0, 248(a2)
-    csrw sepc, t0
-    ld t0, 256(a2)
-    csrw sstatus, t0
-
-    ld ra, 0(a2)
-    ld sp, 8(a2)
-    ld gp, 16(a2)
-    ld tp, 24(a2)
-    ld a0, 32(a2)
-    ld a1, 40(a2)
-    // Skip a2 since it used as a pointer
-    ld a3, 56(a2)
-    ld a4, 64(a2)
-    ld a5, 72(a2)
-    ld a6, 80(a2)
-    ld a7, 88(a2)
-    ld t0, 96(a2)
-    ld t1, 104(a2)
-    ld t2, 112(a2)
-    ld t3, 120(a2)
-    ld t4, 128(a2)
-    ld t5, 136(a2)
-    ld t6, 144(a2)
-    ld s0, 152(a2)
-    ld s1, 160(a2)
-    ld s2, 168(a2)
-    ld s3, 176(a2)
-    ld s4, 184(a2)
-    ld s5, 192(a2)
-    ld s6, 200(a2)
-    ld s7, 208(a2)
-    ld s8, 216(a2)
-    ld s9, 224(a2)
-    ld s10, 232(a2)
-    ld s11, 240(a2)
-
-    // Restore a2 at the end
-    ld a2, 48(a2)
-
-    addi sp, sp, 264
-
-    sret"
-    );
-}
diff --git a/src/time.rs b/src/time.rs
index 0d6062c..d9125d7 100644
--- a/src/time.rs
+++ b/src/time.rs
@@ -1,10 +1,10 @@
 use core::time::Duration;
 
 use alloc::format;
-use log::info;
 
 use crate::{
-    process::ExecutionContext, scheduler::scheduler, set_csr, vga::{Color, FONT_WIDTH, Vga, WIDTH}
+    set_csr,
+    vga::{Color, Vga, FONT_WIDTH, WIDTH},
 };
 
 pub const IRQ_M_TIMER: u8 = 1 << 7;
@@ -30,7 +30,7 @@ pub fn setup_next_timer_interrupt() {
         );
     }
 }
-pub fn timer_interrupt(interrupt_state: ExecutionContext) -> usize {
+pub fn timer_interrupt() {
     let current_time = elapsed_time_since_startup();
     let seconds = current_time.as_secs();
     let minutes = seconds / 60 % 60;
@@ -46,7 +46,6 @@ pub fn timer_interrupt(interrupt_state: ExecutionContext) -> usize {
             Color::BLACK,
         )
     };
-    scheduler(interrupt_state)
 }
 
 pub fn elapsed_time_since_startup() -> Duration {