change(trap): improve trap handling process

crates/eonix_hal/eonix_hal_traits/src/trap.rs

@@ -8,9 +8,12 @@ use eonix_mm::address::VAddr;
 /// and will be used in the HAL crates.
 #[doc(notable_trait)]
 pub trait RawTrapContext: Copy {
+    type FIrq: FnOnce(fn(irqno: usize));
+    type FTimer: FnOnce(fn());
+
     fn new() -> Self;
 
-    fn trap_type(&self) -> TrapType;
+    fn trap_type(&self) -> TrapType<Self::FIrq, Self::FTimer>;
 
     fn get_program_counter(&self) -> usize;
     fn get_stack_pointer(&self) -> usize;
@@ -56,11 +59,15 @@ pub trait IrqState {
 }
 
 /// The reason that caused the trap.
-pub enum TrapType {
+pub enum TrapType<FIrq, FTimer>
+where
+    FIrq: FnOnce(fn(irqno: usize)),
+    FTimer: FnOnce(fn()),
+{
     Syscall { no: usize, args: [usize; 6] },
     Fault(Fault),
-    Irq(usize),
-    Timer,
+    Irq { callback: FIrq },
+    Timer { callback: FTimer },
 }
 
 /// A marker type that indicates that the type is a raw trap context.

crates/eonix_hal/src/arch/riscv64/cpu.rs

@@ -2,10 +2,7 @@ use super::{
     interrupt::InterruptControl,
     trap::{setup_trap, TRAP_SCRATCH},
 };
-use crate::arch::{
-    fdt::{FdtExt, FDT},
-    time::set_next_timer,
-};
+use crate::arch::fdt::{FdtExt, FDT};
 use core::{arch::asm, pin::Pin, ptr::NonNull};
 use eonix_preempt::PreemptGuard;
 use eonix_sync_base::LazyLock;
@@ -81,12 +78,6 @@ impl CPU {
         // nothing
     }
 
-    pub fn end_of_interrupt(self: Pin<&mut Self>) {
-        // TODO: only timer interrupt should do this, here may need to change 
-        // if some other interrupt need send end signal
-        set_next_timer();
-    }
-
     pub fn local() -> PreemptGuard<Pin<&'static mut Self>> {
         unsafe {
             // SAFETY: We pass the reference into a `PreemptGuard`, which ensures

crates/eonix_hal/src/arch/riscv64/trap/trap_context.rs

@@ -1,3 +1,4 @@
+use crate::arch::time::set_next_timer;
 use core::arch::asm;
 use eonix_hal_traits::{
     fault::{Fault, PageFaultErrorCode},
@@ -117,6 +118,9 @@ impl TrapContext {
 }
 
 impl RawTrapContext for TrapContext {
+    type FIrq = fn(handler: fn(irqno: usize));
+    type FTimer = fn(handler: fn());
+
     fn new() -> Self {
         let mut sstatus = Sstatus::from_bits(0);
         sstatus.set_fs(FS::Initial);
@@ -130,14 +134,24 @@ impl RawTrapContext for TrapContext {
         }
     }
 
-    fn trap_type(&self) -> TrapType {
+    fn trap_type(&self) -> TrapType<Self::FIrq, Self::FTimer> {
         let cause = self.scause.cause();
         match cause {
             Trap::Interrupt(i) => {
                 match Interrupt::from_number(i).unwrap() {
-                    Interrupt::SupervisorTimer => TrapType::Timer,
-                    // TODO: need to read plic
-                    Interrupt::SupervisorExternal => TrapType::Irq(0),
+                    Interrupt::SupervisorTimer => TrapType::Timer {
+                        callback: |handler| {
+                            handler();
+                            set_next_timer();
+                        },
+                    },
+                    Interrupt::SupervisorExternal => TrapType::Irq {
+                        callback: |handler| {
+                            // TODO: Read PLIC to determine the IRQ number.
+                            //       Use 0xa (serial) for now.
+                            handler(0xa);
+                        },
+                    },
                     // soft interrupt
                     _ => TrapType::Fault(Fault::Unknown(i)),
                 }

crates/eonix_hal/src/arch/x86_64/trap/trap_context.rs

@@ -1,3 +1,4 @@
+use crate::processor::CPU;
 use core::arch::asm;
 use eonix_hal_traits::{
     fault::{Fault, PageFaultErrorCode},
@@ -77,16 +78,24 @@ impl TrapContext {
 }
 
 impl RawTrapContext for TrapContext {
+    type FIrq = impl FnOnce(fn(irqno: usize));
+    type FTimer = fn(handler: fn());
+
     fn new() -> Self {
         Self {
             ..Default::default()
         }
     }
 
-    fn trap_type(&self) -> TrapType {
+    fn trap_type(&self) -> TrapType<Self::FIrq, Self::FTimer> {
         match self.int_no {
             0..0x20 => TrapType::Fault(self.get_fault_type()),
-            0x40 => TrapType::Timer,
+            0x40 => TrapType::Timer {
+                callback: |handler| {
+                    handler();
+                    CPU::local().as_mut().end_of_interrupt();
+                },
+            },
             0x80 => TrapType::Syscall {
                 no: self.rax as usize,
                 args: [
@@ -98,7 +107,22 @@ impl RawTrapContext for TrapContext {
                     self.rbp as usize,
                 ],
             },
-            no => TrapType::Irq(no as usize - 0x20),
+            no => TrapType::Irq {
+                callback: move |handler| {
+                    let irqno = no as usize - 0x20;
+                    handler(irqno);
+
+                    use crate::arch::io::Port8;
+
+                    const PIC1_COMMAND: Port8 = Port8::new(0x20);
+                    const PIC2_COMMAND: Port8 = Port8::new(0xA0);
+
+                    PIC1_COMMAND.write(0x20); // EOI
+                    if irqno >= 8 {
+                        PIC2_COMMAND.write(0x20); // EOI
+                    }
+                },
+            },
         }
     }
 

crates/eonix_hal/src/lib.rs

@@ -1,6 +1,7 @@
 #![no_std]
 #![feature(allocator_api)]
 #![feature(doc_notable_trait)]
+#![feature(impl_trait_in_assoc_type)]
 
 pub(crate) mod arch;
 

src/kernel/interrupt.rs

@@ -3,7 +3,6 @@ use super::timer::{should_reschedule, timer_interrupt};
 use crate::kernel::constants::EINVAL;
 use crate::prelude::*;
 use alloc::sync::Arc;
-use eonix_hal::processor::CPU;
 use eonix_hal::traits::fault::Fault;
 use eonix_hal::traits::trap::{RawTrapContext, TrapType};
 use eonix_hal::trap::TrapContext;
@@ -24,19 +23,6 @@ pub fn default_irq_handler(irqno: usize) {
             handler();
         }
     }
-
-    #[cfg(target_arch = "x86_64")]
-    {
-        use eonix_hal::arch_exported::io::Port8;
-
-        const PIC1_COMMAND: Port8 = Port8::new(0x20);
-        const PIC2_COMMAND: Port8 = Port8::new(0xA0);
-
-        PIC1_COMMAND.write(0x20); // EOI
-        if irqno >= 8 {
-            PIC2_COMMAND.write(0x20); // EOI
-        }
-    }
 }
 
 pub fn default_fault_handler(fault_type: Fault, trap_ctx: &mut TrapContext) {
@@ -64,9 +50,9 @@ pub fn interrupt_handler(trap_ctx: &mut TrapContext) {
     match trap_ctx.trap_type() {
         TrapType::Syscall { no, .. } => unreachable!("Syscall {} in kernel space.", no),
         TrapType::Fault(fault) => default_fault_handler(fault, trap_ctx),
-        TrapType::Irq(no) => default_irq_handler(no),
-        TrapType::Timer => {
-            timer_interrupt();
+        TrapType::Irq { callback } => callback(default_irq_handler),
+        TrapType::Timer { callback } => {
+            callback(timer_interrupt);
 
             if eonix_preempt::count() == 0 && should_reschedule() {
                 // To make scheduler satisfied.
@@ -88,7 +74,3 @@ where
     IRQ_HANDLERS.lock_irq()[irqno as usize].push(Arc::new(handler));
     Ok(())
 }
-
-pub fn end_of_interrupt() {
-    CPU::local().as_mut().end_of_interrupt();
-}

src/kernel/task/thread.rs

@@ -383,9 +383,9 @@ impl Thread {
                     self.signal_list.raise(Signal::SIGILL);
                 }
                 TrapType::Fault(Fault::Unknown(_)) => unimplemented!("Unhandled fault"),
-                TrapType::Irq(irqno) => default_irq_handler(irqno),
-                TrapType::Timer => {
-                    timer_interrupt();
+                TrapType::Irq { callback } => callback(default_irq_handler),
+                TrapType::Timer { callback } => {
+                    callback(timer_interrupt);
 
                     if should_reschedule() {
                         yield_now().await;

src/kernel/timer.rs

@@ -1,4 +1,3 @@
-use super::interrupt::end_of_interrupt;
 use alloc::{collections::BinaryHeap, vec, vec::Vec};
 use core::{
     cell::RefCell,
@@ -110,7 +109,6 @@ impl Add<Duration> for Instant {
 }
 
 pub fn timer_interrupt() {
-    end_of_interrupt();
     if CPU::local().cpuid() != 0 {
         // Only the BSP should handle the timer interrupt.
         return;