From dfd66be8abced1bba1aa09bf117f028c63c926ad Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Bj=C3=B6rn=20Quentin?= <bjoernQ@users.noreply.github.com>
Date: Wed, 30 Jul 2025 20:31:05 +0200
Subject: [PATCH] Simplify riscv trap handler (#3875)

* use `riscv::interrupt::nested` for nesting

* Only save/restore caller saved registers in trap-frame (risc-v)

* Replace RISC-V scheduler

* Cleanup

* Clippy

* CHANGELOG.md

* Make sure to reset the runlevel to current, not prio-1

* Clippy

* Fix it for real

* Address review comments

* Address review comments

* More docs, optimize away one instruction

* Address review comments

* Address review comments

* Use a bool for the nested flag again

* Fix

* Clippy
---
 esp-hal/CHANGELOG.md                          |   1 +
 esp-hal/src/interrupt/mod.rs                  |  37 +-
 esp-hal/src/interrupt/riscv.rs                | 109 +-----
 esp-radio-preempt-baremetal/Cargo.toml        |   1 +
 esp-radio-preempt-baremetal/src/lib.rs        |  21 ++
 esp-radio-preempt-baremetal/src/task/mod.rs   |  17 +-
 esp-radio-preempt-baremetal/src/task/riscv.rs | 331 ++++++++++++++----
 esp-radio-preempt-baremetal/src/timer/mod.rs  |  27 +-
 .../src/timer/riscv.rs                        |  23 +-
 .../src/timer/xtensa.rs                       |  15 +
 esp-riscv-rt/src/lib.rs                       | 179 ++--------
 11 files changed, 425 insertions(+), 336 deletions(-)

diff --git a/esp-hal/CHANGELOG.md b/esp-hal/CHANGELOG.md
index 57fe0433e..dd912960d 100644
--- a/esp-hal/CHANGELOG.md
+++ b/esp-hal/CHANGELOG.md
@@ -11,6 +11,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 - A reimplemntation of the `assign_resources!` macro (#3809)
 - `TrngSource` to manage random number generator entropy (#3829)
+- On RISC-V you can opt-out of nested interrupts for an interrupt handler by using `new_not_nested` (#3875)
 
 ### Changed
 
diff --git a/esp-hal/src/interrupt/mod.rs b/esp-hal/src/interrupt/mod.rs
index 77a6d19ff..4084cd51d 100644
--- a/esp-hal/src/interrupt/mod.rs
+++ b/esp-hal/src/interrupt/mod.rs
@@ -135,19 +135,50 @@ pub trait InterruptConfigurable: crate::private::Sealed {
 pub struct InterruptHandler {
     f: extern "C" fn(),
     prio: Priority,
+    #[cfg(riscv)]
+    nested: bool,
 }
 
 impl InterruptHandler {
     /// Creates a new [InterruptHandler] which will call the given function at
     /// the given priority.
     pub const fn new(f: extern "C" fn(), prio: Priority) -> Self {
-        Self { f, prio }
+        Self {
+            f,
+            prio,
+            #[cfg(riscv)]
+            nested: true,
+        }
     }
 
-    /// The function to be called
+    /// Creates a new [InterruptHandler] which will call the given function at
+    /// the given priority with disabled interrupt nesting.
+    ///
+    /// Usually higher priority interrupts get served while handling an interrupt.
+    /// Using this the interrupt handler won't get preempted by higher priority interrupts.
+    #[cfg(riscv)]
+    pub fn new_not_nested(f: extern "C" fn(), prio: Priority) -> Self {
+        Self {
+            f,
+            prio,
+            nested: false,
+        }
+    }
+
+    /// The function to be called.
+    #[cfg_attr(
+        riscv,
+        doc = "\n\nNote that the function pointer might be misaligned. Don't ever just call the function."
+    )]
     #[inline]
     pub fn handler(&self) -> extern "C" fn() {
-        self.f
+        cfg_if::cfg_if! {
+            if #[cfg(riscv)] {
+                unsafe { core::mem::transmute::<usize, extern "C" fn()>((self.f as usize) | !self.nested as usize) }
+            } else {
+                self.f
+            }
+        }
     }
 
     /// Priority to be used when registering the interrupt
diff --git a/esp-hal/src/interrupt/riscv.rs b/esp-hal/src/interrupt/riscv.rs
index 552533192..44b413170 100644
--- a/esp-hal/src/interrupt/riscv.rs
+++ b/esp-hal/src/interrupt/riscv.rs
@@ -542,45 +542,6 @@ mod classic {
 
         prev_interrupt_priority
     }
-
-    #[cfg(feature = "rt")]
-    mod rt {
-        use super::*;
-
-        #[unsafe(link_section = ".trap")]
-        #[unsafe(no_mangle)]
-        pub(super) unsafe extern "C" fn _handle_priority() -> u32 {
-            // Both C6 and H2 have 5 bits of code. The riscv crate masks 31 bits, which then
-            // causes a bounds check to be present.
-            let interrupt_id: usize = riscv::register::mcause::read().bits() & 0x1f;
-            let intr = INTERRUPT_CORE0::regs();
-            let interrupt_priority = intr.cpu_int_pri(interrupt_id).read().bits();
-
-            let prev_interrupt_priority = intr.cpu_int_thresh().read().bits();
-            if interrupt_priority < 15 {
-                // leave interrupts disabled if interrupt is of max priority.
-                intr.cpu_int_thresh()
-                    .write(|w| unsafe { w.bits(interrupt_priority + 1) }); // set the prio threshold to 1 more than current interrupt prio
-                unsafe { riscv::interrupt::enable() };
-            }
-            prev_interrupt_priority
-        }
-
-        #[unsafe(link_section = ".trap")]
-        #[unsafe(no_mangle)]
-        unsafe extern "C" fn _restore_priority(stored_prio: u32) {
-            riscv::interrupt::disable();
-            let intr = INTERRUPT_CORE0::regs();
-            intr.cpu_int_thresh()
-                .write(|w| unsafe { w.bits(stored_prio) });
-        }
-
-        /// Globally exported so assembly code can call it.
-        #[unsafe(no_mangle)]
-        unsafe extern "C" fn priority(cpu_interrupt: CpuInterrupt) -> Priority {
-            super::priority(cpu_interrupt)
-        }
-    }
 }
 
 #[cfg(plic)]
@@ -723,53 +684,6 @@ mod plic {
 
         prev_interrupt_priority
     }
-
-    #[cfg(feature = "rt")]
-    mod rt {
-        use super::*;
-        /// Get interrupt priority - called by assembly code
-        #[unsafe(no_mangle)]
-        pub(super) unsafe extern "C" fn priority(cpu_interrupt: CpuInterrupt) -> Priority {
-            super::priority(cpu_interrupt)
-        }
-
-        #[unsafe(no_mangle)]
-        #[unsafe(link_section = ".trap")]
-        pub(super) unsafe extern "C" fn _handle_priority() -> u32 {
-            let interrupt_id: usize = riscv::register::mcause::read().code(); // MSB is whether its exception or interrupt.
-            let interrupt_priority = PLIC_MX::regs()
-                .mxint_pri(interrupt_id)
-                .read()
-                .cpu_mxint_pri()
-                .bits();
-
-            let prev_interrupt_priority = PLIC_MX::regs()
-                .mxint_thresh()
-                .read()
-                .cpu_mxint_thresh()
-                .bits();
-            if interrupt_priority < 15 {
-                // leave interrupts disabled if interrupt is of max priority.
-                PLIC_MX::regs()
-                    .mxint_thresh()
-                    .write(|w| unsafe { w.cpu_mxint_thresh().bits(interrupt_priority + 1) });
-
-                unsafe {
-                    riscv::interrupt::enable();
-                }
-            }
-            prev_interrupt_priority as u32
-        }
-
-        #[unsafe(no_mangle)]
-        #[unsafe(link_section = ".trap")]
-        pub(super) unsafe extern "C" fn _restore_priority(stored_prio: u32) {
-            riscv::interrupt::disable();
-            PLIC_MX::regs()
-                .mxint_thresh()
-                .write(|w| unsafe { w.cpu_mxint_thresh().bits(stored_prio as u8) });
-        }
-    }
 }
 
 #[cfg(feature = "rt")]
@@ -849,12 +763,25 @@ mod rt {
         let configured_interrupts = vectored::configured_interrupts(core, status, prio);
 
         for interrupt_nr in configured_interrupts.iterator() {
-            let handler = unsafe { pac::__EXTERNAL_INTERRUPTS[interrupt_nr as usize]._handler };
+            let handler =
+                unsafe { pac::__EXTERNAL_INTERRUPTS[interrupt_nr as usize]._handler } as usize;
+            let not_nested = (handler & 1) == 1;
+            let handler = handler & !1;
 
-            let handler: fn(&mut TrapFrame) = unsafe {
-                core::mem::transmute::<unsafe extern "C" fn(), fn(&mut TrapFrame)>(handler)
-            };
-            handler(context);
+            let handler: fn(&mut TrapFrame) =
+                unsafe { core::mem::transmute::<usize, fn(&mut TrapFrame)>(handler) };
+
+            if not_nested || prio == Priority::Priority15 {
+                handler(context);
+            } else {
+                let elevated = prio as u8;
+                unsafe {
+                    let level =
+                        change_current_runlevel(unwrap!(Priority::try_from(elevated as u32)));
+                    riscv::interrupt::nested(|| handler(context));
+                    change_current_runlevel(level);
+                }
+            }
         }
     }
 
diff --git a/esp-radio-preempt-baremetal/Cargo.toml b/esp-radio-preempt-baremetal/Cargo.toml
index 4a2c995b5..e67d059fc 100644
--- a/esp-radio-preempt-baremetal/Cargo.toml
+++ b/esp-radio-preempt-baremetal/Cargo.toml
@@ -31,6 +31,7 @@ document-features  = "0.2.11"
 esp-alloc = { version = "0.8.0", path = "../esp-alloc", optional = true }
 esp-config = { version = "0.5.0", path = "../esp-config" }
 esp-radio-preempt-driver = { version = "0.0.1", path = "../esp-radio-preempt-driver" }
+portable-atomic = { version = "1.11.0", default-features = false }
 
 # Logging interfaces, they are mutually exclusive so they need to be behind separate features.
 defmt            = { version = "1.0", optional = true }
diff --git a/esp-radio-preempt-baremetal/src/lib.rs b/esp-radio-preempt-baremetal/src/lib.rs
index f59632ff3..57b650049 100644
--- a/esp-radio-preempt-baremetal/src/lib.rs
+++ b/esp-radio-preempt-baremetal/src/lib.rs
@@ -153,6 +153,7 @@ impl SchedulerState {
         }
     }
 
+    #[cfg(xtensa)]
     fn switch_task(&mut self, trap_frame: &mut TrapFrame) {
         task::save_task_context(unsafe { &mut *self.current_task }, trap_frame);
 
@@ -166,6 +167,26 @@ impl SchedulerState {
         task::restore_task_context(unsafe { &mut *self.current_task }, trap_frame);
     }
 
+    #[cfg(riscv)]
+    fn switch_task(&mut self, _trap_frame: &mut TrapFrame) {
+        if !self.to_delete.is_null() {
+            let task_to_delete = core::mem::take(&mut self.to_delete);
+            self.delete_task(task_to_delete);
+        }
+
+        let task = self.current_task;
+        let context = unsafe { &mut (*task).trap_frame };
+        let old_ctx = core::ptr::addr_of_mut!(*context);
+
+        let task = unsafe { (*self.current_task).next };
+        let context = unsafe { &mut (*task).trap_frame };
+        let new_ctx = core::ptr::addr_of_mut!(*context);
+
+        if crate::task::arch_specific::task_switch(old_ctx, new_ctx) {
+            unsafe { self.current_task = (*self.current_task).next };
+        }
+    }
+
     fn schedule_task_deletion(&mut self, task: *mut Context) -> bool {
         if task.is_null() {
             self.to_delete = self.current_task;
diff --git a/esp-radio-preempt-baremetal/src/task/mod.rs b/esp-radio-preempt-baremetal/src/task/mod.rs
index d2399e561..6a0dcee61 100644
--- a/esp-radio-preempt-baremetal/src/task/mod.rs
+++ b/esp-radio-preempt-baremetal/src/task/mod.rs
@@ -1,17 +1,23 @@
-#[cfg_attr(target_arch = "riscv32", path = "riscv.rs")]
-#[cfg_attr(target_arch = "xtensa", path = "xtensa.rs")]
-mod arch_specific;
+#[cfg_attr(riscv, path = "riscv.rs")]
+#[cfg_attr(xtensa, path = "xtensa.rs")]
+pub(crate) mod arch_specific;
 
 use core::{ffi::c_void, mem::MaybeUninit};
 
 use allocator_api2::boxed::Box;
+#[cfg(riscv)]
+use arch_specific::Registers;
 pub(crate) use arch_specific::*;
 use esp_hal::trapframe::TrapFrame;
 
 use crate::{InternalMemory, SCHEDULER_STATE, task, timer};
 
+#[repr(C)]
 pub(crate) struct Context {
-    trap_frame: TrapFrame,
+    #[cfg(riscv)]
+    pub trap_frame: Registers,
+    #[cfg(xtensa)]
+    pub trap_frame: TrapFrame,
     pub thread_semaphore: u32,
     pub next: *mut Context,
     pub _allocated_stack: Box<[MaybeUninit<u8>], InternalMemory>,
@@ -42,6 +48,9 @@ pub(super) fn allocate_main_task() {
     // This context will be filled out by the first context switch.
     let context = Box::new_in(
         Context {
+            #[cfg(riscv)]
+            trap_frame: Registers::default(),
+            #[cfg(xtensa)]
             trap_frame: TrapFrame::default(),
             thread_semaphore: 0,
             next: core::ptr::null_mut(),
diff --git a/esp-radio-preempt-baremetal/src/task/riscv.rs b/esp-radio-preempt-baremetal/src/task/riscv.rs
index 04bc336e0..e34aaae26 100644
--- a/esp-radio-preempt-baremetal/src/task/riscv.rs
+++ b/esp-radio-preempt-baremetal/src/task/riscv.rs
@@ -1,18 +1,108 @@
 use core::ffi::c_void;
 
-pub use esp_hal::trapframe::TrapFrame;
+unsafe extern "C" {
+    fn sys_switch();
+}
 
-use crate::task::Context;
+static _CURRENT_CTX_PTR: portable_atomic::AtomicPtr<Registers> =
+    portable_atomic::AtomicPtr::new(core::ptr::null_mut());
+
+static _NEXT_CTX_PTR: portable_atomic::AtomicPtr<Registers> =
+    portable_atomic::AtomicPtr::new(core::ptr::null_mut());
+
+/// Registers saved / restored
+#[derive(Debug, Default, Clone)]
+#[repr(C)]
+pub struct Registers {
+    /// Return address, stores the address to return to after a function call or
+    /// interrupt.
+    pub ra: usize,
+    /// Temporary register t0, used for intermediate values.
+    pub t0: usize,
+    /// Temporary register t1, used for intermediate values.
+    pub t1: usize,
+    /// Temporary register t2, used for intermediate values.
+    pub t2: usize,
+    /// Temporary register t3, used for intermediate values.
+    pub t3: usize,
+    /// Temporary register t4, used for intermediate values.
+    pub t4: usize,
+    /// Temporary register t5, used for intermediate values.
+    pub t5: usize,
+    /// Temporary register t6, used for intermediate values.
+    pub t6: usize,
+    /// Argument register a0, typically used to pass the first argument to a
+    /// function.
+    pub a0: usize,
+    /// Argument register a1, typically used to pass the second argument to a
+    /// function.
+    pub a1: usize,
+    /// Argument register a2, typically used to pass the third argument to a
+    /// function.
+    pub a2: usize,
+    /// Argument register a3, typically used to pass the fourth argument to a
+    /// function.
+    pub a3: usize,
+    /// Argument register a4, typically used to pass the fifth argument to a
+    /// function.
+    pub a4: usize,
+    /// Argument register a5, typically used to pass the sixth argument to a
+    /// function.
+    pub a5: usize,
+    /// Argument register a6, typically used to pass the seventh argument to a
+    /// function.
+    pub a6: usize,
+    /// Argument register a7, typically used to pass the eighth argument to a
+    /// function.
+    pub a7: usize,
+    /// Saved register s0, used to hold values across function calls.
+    pub s0: usize,
+    /// Saved register s1, used to hold values across function calls.
+    pub s1: usize,
+    /// Saved register s2, used to hold values across function calls.
+    pub s2: usize,
+    /// Saved register s3, used to hold values across function calls.
+    pub s3: usize,
+    /// Saved register s4, used to hold values across function calls.
+    pub s4: usize,
+    /// Saved register s5, used to hold values across function calls.
+    pub s5: usize,
+    /// Saved register s6, used to hold values across function calls.
+    pub s6: usize,
+    /// Saved register s7, used to hold values across function calls.
+    pub s7: usize,
+    /// Saved register s8, used to hold values across function calls.
+    pub s8: usize,
+    /// Saved register s9, used to hold values across function calls.
+    pub s9: usize,
+    /// Saved register s10, used to hold values across function calls.
+    pub s10: usize,
+    /// Saved register s11, used to hold values across function calls.
+    pub s11: usize,
+    /// Global pointer register, holds the address of the global data area.
+    pub gp: usize,
+    /// Thread pointer register, holds the address of the thread-local storage
+    /// area.
+    pub tp: usize,
+    /// Stack pointer register, holds the address of the top of the stack.
+    pub sp: usize,
+    /// Program counter, stores the address of the next instruction to be
+    /// executed.
+    pub pc: usize,
+
+    /// The mstatus which will be loaded before MRET
+    pub mstatus: usize,
+}
 
 pub(crate) fn new_task_context(
     task: extern "C" fn(*mut c_void),
     param: *mut c_void,
     stack_top: *mut (),
-) -> TrapFrame {
+) -> Registers {
     let stack_top = stack_top as usize;
     let stack_top = stack_top - (stack_top % 16);
 
-    TrapFrame {
+    Registers {
         pc: task as usize,
         a0: param as usize,
         sp: stack_top,
@@ -20,72 +110,171 @@ pub(crate) fn new_task_context(
     }
 }
 
-pub(crate) fn restore_task_context(ctx: &mut Context, trap_frame: &mut TrapFrame) {
-    trap_frame.ra = ctx.trap_frame.ra;
-    trap_frame.sp = ctx.trap_frame.sp;
-    trap_frame.a0 = ctx.trap_frame.a0;
-    trap_frame.a1 = ctx.trap_frame.a1;
-    trap_frame.a2 = ctx.trap_frame.a2;
-    trap_frame.a3 = ctx.trap_frame.a3;
-    trap_frame.a4 = ctx.trap_frame.a4;
-    trap_frame.a5 = ctx.trap_frame.a5;
-    trap_frame.a6 = ctx.trap_frame.a6;
-    trap_frame.a7 = ctx.trap_frame.a7;
-    trap_frame.t0 = ctx.trap_frame.t0;
-    trap_frame.t1 = ctx.trap_frame.t1;
-    trap_frame.t2 = ctx.trap_frame.t2;
-    trap_frame.t3 = ctx.trap_frame.t3;
-    trap_frame.t4 = ctx.trap_frame.t4;
-    trap_frame.t5 = ctx.trap_frame.t5;
-    trap_frame.t6 = ctx.trap_frame.t6;
-    trap_frame.s0 = ctx.trap_frame.s0;
-    trap_frame.s1 = ctx.trap_frame.s1;
-    trap_frame.s2 = ctx.trap_frame.s2;
-    trap_frame.s3 = ctx.trap_frame.s3;
-    trap_frame.s4 = ctx.trap_frame.s4;
-    trap_frame.s5 = ctx.trap_frame.s5;
-    trap_frame.s6 = ctx.trap_frame.s6;
-    trap_frame.s7 = ctx.trap_frame.s7;
-    trap_frame.s8 = ctx.trap_frame.s8;
-    trap_frame.s9 = ctx.trap_frame.s9;
-    trap_frame.s10 = ctx.trap_frame.s10;
-    trap_frame.s11 = ctx.trap_frame.s11;
-    trap_frame.gp = ctx.trap_frame.gp;
-    trap_frame.tp = ctx.trap_frame.tp;
-    trap_frame.pc = ctx.trap_frame.pc;
+/// Switch to next task
+///
+/// *MUST* be called from inside an ISR without interrupt nesting.
+///
+/// The task switching relies on MEPC and MSTATUS to not get clobbered.
+/// We save MEPC as the current task's PC and change MEPC to an assembly function
+/// which will save the current CPU state for the current task (excluding PC) and
+/// restoring the CPU state from the next task.
+pub fn task_switch(old_ctx: *mut Registers, new_ctx: *mut Registers) -> bool {
+    // Check that there isn't a switch "in-progress"
+    //
+    // While this shouldn't happen: from observation it does!
+    //
+    // This happens if the timer tick interrupt gets asserted while the task switch is in
+    // progress (i.e. the sw-int is served).
+    //
+    // In that case returning via `mret` will first service the pending interrupt before actually
+    // ending up in `sys_switch`.
+    //
+    // Setting MPIE to 0 _should_ prevent that from happening.
+    if !_NEXT_CTX_PTR
+        .load(portable_atomic::Ordering::SeqCst)
+        .is_null()
+    {
+        return false;
+    }
+
+    _CURRENT_CTX_PTR.store(old_ctx, portable_atomic::Ordering::SeqCst);
+    _NEXT_CTX_PTR.store(new_ctx, portable_atomic::Ordering::SeqCst);
+
+    let old = esp_hal::riscv::register::mepc::read();
+    unsafe {
+        (*old_ctx).pc = old;
+    }
+
+    // set MSTATUS for the switched to task
+    // MIE will be set from MPIE
+    // MPP will be used to determine the privilege-level
+    let mstatus = esp_hal::riscv::register::mstatus::read().bits();
+    unsafe {
+        (*new_ctx).mstatus = mstatus;
+    }
+
+    unsafe {
+        // set MPIE in MSTATUS to 0 to disable interrupts while task switching
+        esp_hal::riscv::register::mstatus::write(
+            esp_hal::riscv::register::mstatus::Mstatus::from_bits(mstatus & !(1 << 7)),
+        );
+
+        // load address of sys_switch into MEPC - will run after all registers are restored
+        esp_hal::riscv::register::mepc::write(sys_switch as usize);
+    }
+
+    true
 }
 
-pub(crate) fn save_task_context(ctx: &mut Context, trap_frame: &TrapFrame) {
-    ctx.trap_frame.ra = trap_frame.ra;
-    ctx.trap_frame.sp = trap_frame.sp;
-    ctx.trap_frame.a0 = trap_frame.a0;
-    ctx.trap_frame.a1 = trap_frame.a1;
-    ctx.trap_frame.a2 = trap_frame.a2;
-    ctx.trap_frame.a3 = trap_frame.a3;
-    ctx.trap_frame.a4 = trap_frame.a4;
-    ctx.trap_frame.a5 = trap_frame.a5;
-    ctx.trap_frame.a6 = trap_frame.a6;
-    ctx.trap_frame.a7 = trap_frame.a7;
-    ctx.trap_frame.t0 = trap_frame.t0;
-    ctx.trap_frame.t1 = trap_frame.t1;
-    ctx.trap_frame.t2 = trap_frame.t2;
-    ctx.trap_frame.t3 = trap_frame.t3;
-    ctx.trap_frame.t4 = trap_frame.t4;
-    ctx.trap_frame.t5 = trap_frame.t5;
-    ctx.trap_frame.t6 = trap_frame.t6;
-    ctx.trap_frame.s0 = trap_frame.s0;
-    ctx.trap_frame.s1 = trap_frame.s1;
-    ctx.trap_frame.s2 = trap_frame.s2;
-    ctx.trap_frame.s3 = trap_frame.s3;
-    ctx.trap_frame.s4 = trap_frame.s4;
-    ctx.trap_frame.s5 = trap_frame.s5;
-    ctx.trap_frame.s6 = trap_frame.s6;
-    ctx.trap_frame.s7 = trap_frame.s7;
-    ctx.trap_frame.s8 = trap_frame.s8;
-    ctx.trap_frame.s9 = trap_frame.s9;
-    ctx.trap_frame.s10 = trap_frame.s10;
-    ctx.trap_frame.s11 = trap_frame.s11;
-    ctx.trap_frame.gp = trap_frame.gp;
-    ctx.trap_frame.tp = trap_frame.tp;
-    ctx.trap_frame.pc = trap_frame.pc;
-}
+core::arch::global_asm!(
+    r#"
+.section .trap, "ax"
+
+.globl sys_switch
+.align 4
+sys_switch:
+    # put some regs on the stack since we will need those regs
+    addi sp, sp, -16
+    sw t0, 0(sp)
+    sw t1, 4(sp)
+
+    # t0 => current context
+    la t0, {_CURRENT_CTX_PTR}
+    lw t0, 0(t0)
+
+    # store registers to old context - PC needs to be set by the "caller"
+    sw ra, 0*4(t0)
+
+    lw t1, 0(sp)
+    sw t1, 1*4(t0)
+
+    lw t1, 4(sp)
+    sw t1, 2*4(t0)
+
+    sw t2, 3*4(t0)
+    sw t3, 4*4(t0)
+    sw t4, 5*4(t0)
+    sw t5, 6*4(t0)
+    sw t6, 7*4(t0)
+    sw a0, 8*4(t0)
+    sw a1, 9*4(t0)
+    sw a2, 10*4(t0)
+    sw a3, 11*4(t0)
+    sw a4, 12*4(t0)
+    sw a5, 13*4(t0)
+    sw a6, 14*4(t0)
+    sw a7, 15*4(t0)
+    sw s0, 16*4(t0)
+    sw s1, 17*4(t0)
+    sw s2, 18*4(t0)
+    sw s3, 19*4(t0)
+    sw s4, 20*4(t0)
+    sw s5, 21*4(t0)
+    sw s6, 22*4(t0)
+    sw s7, 23*4(t0)
+    sw s8, 24*4(t0)
+    sw s9, 25*4(t0)
+    sw s10, 26*4(t0)
+    sw s11, 27*4(t0)
+    sw gp, 28*4(t0)
+    sw tp, 29*4(t0)
+
+    addi t1, sp, 16
+    sw t1, 30*4(t0)
+
+    # t0 => next context
+    la t1, {_NEXT_CTX_PTR}
+    lw t0, 0(t1)
+
+    # signal that the task switch is done - safe to do it already now - interrupts are disabled
+    sw x0, 0(t1)
+
+    # set the next task's PC as MEPC
+    lw t1, 31*4(t0)
+    csrrw x0, mepc, t1
+
+    # set MSTATUS from next context
+    lw t1, 32*4(t0)
+    csrrw x0, mstatus, t1
+
+    # restore registers from old context
+    lw ra, 0*4(t0)
+    lw t1, 2*4(t0)
+    lw t2, 3*4(t0)
+    lw t3, 4*4(t0)
+    lw t4, 5*4(t0)
+    lw t5, 6*4(t0)
+    lw t6, 7*4(t0)
+    lw a0, 8*4(t0)
+    lw a1, 9*4(t0)
+    lw a2, 10*4(t0)
+    lw a3, 11*4(t0)
+    lw a4, 12*4(t0)
+    lw a5, 13*4(t0)
+    lw a6, 14*4(t0)
+    lw a7, 15*4(t0)
+    lw s0, 16*4(t0)
+    lw s1, 17*4(t0)
+    lw s2, 18*4(t0)
+    lw s3, 19*4(t0)
+    lw s4, 20*4(t0)
+    lw s5, 21*4(t0)
+    lw s6, 22*4(t0)
+    lw s7, 23*4(t0)
+    lw s8, 24*4(t0)
+    lw s9, 25*4(t0)
+    lw s10, 26*4(t0)
+    lw s11, 27*4(t0)
+    lw gp, 28*4(t0)
+    lw tp, 29*4(t0)
+    lw sp, 30*4(t0)
+    lw t0, 1*4(t0)
+
+
+    # jump to next task's PC
+    mret
+
+    "#, 
+    _CURRENT_CTX_PTR = sym _CURRENT_CTX_PTR,
+    _NEXT_CTX_PTR = sym _NEXT_CTX_PTR,
+);
diff --git a/esp-radio-preempt-baremetal/src/timer/mod.rs b/esp-radio-preempt-baremetal/src/timer/mod.rs
index ab737e59a..1e169b23d 100644
--- a/esp-radio-preempt-baremetal/src/timer/mod.rs
+++ b/esp-radio-preempt-baremetal/src/timer/mod.rs
@@ -1,12 +1,11 @@
-#[cfg_attr(target_arch = "riscv32", path = "riscv.rs")]
-#[cfg_attr(target_arch = "xtensa", path = "xtensa.rs")]
+#[cfg_attr(riscv, path = "riscv.rs")]
+#[cfg_attr(xtensa, path = "xtensa.rs")]
 mod arch_specific;
 
 pub(crate) use arch_specific::*;
 use esp_hal::{
     interrupt::{InterruptHandler, Priority},
     sync::Locked,
-    trapframe::TrapFrame,
 };
 
 use crate::TimeBase;
@@ -22,6 +21,13 @@ pub(crate) fn setup_timebase(mut timer: TimeBase) {
     // The timer needs to tick at Priority 1 to prevent accidentally interrupting
     // priority 1 limited locks.
     let cb: extern "C" fn() = unsafe { core::mem::transmute(timer_tick_handler as *const ()) };
+
+    // Register the interrupt handler without nesting to satisfy the requirements of the task
+    // switching code
+    #[cfg(riscv)]
+    let handler = InterruptHandler::new_not_nested(cb, Priority::Priority1);
+
+    #[cfg(xtensa)]
     let handler = InterruptHandler::new(cb, Priority::Priority1);
 
     timer.set_interrupt_handler(handler);
@@ -44,18 +50,3 @@ pub(crate) fn disable_timebase() {
         unwrap!(timer.cancel());
     });
 }
-
-#[esp_hal::ram]
-extern "C" fn timer_tick_handler(context: &mut TrapFrame) {
-    clear_timer_interrupt();
-
-    // `task_switch` must be called on a single interrupt priority level only.
-    // Because on ESP32 the software interrupt is triggered at priority 3 but
-    // the timer interrupt is triggered at priority 1, we need to trigger the
-    // software interrupt manually.
-    if cfg!(esp32) {
-        yield_task();
-    } else {
-        crate::task::task_switch(context);
-    }
-}
diff --git a/esp-radio-preempt-baremetal/src/timer/riscv.rs b/esp-radio-preempt-baremetal/src/timer/riscv.rs
index 18ea5a08e..16a4059f3 100644
--- a/esp-radio-preempt-baremetal/src/timer/riscv.rs
+++ b/esp-radio-preempt-baremetal/src/timer/riscv.rs
@@ -1,24 +1,27 @@
 use esp_hal::{
-    interrupt::{self, TrapFrame, software::SoftwareInterrupt},
+    interrupt::{self, software::SoftwareInterrupt},
     peripherals::Interrupt,
 };
 
 use crate::task::task_switch;
 
 pub(crate) fn setup_multitasking() {
-    unwrap!(interrupt::enable(
-        Interrupt::FROM_CPU_INTR2,
-        interrupt::Priority::Priority1,
-    ));
+    // Register the interrupt handler without nesting to satisfy the requirements of the task
+    // switching code
+    let swint2_handler = esp_hal::interrupt::InterruptHandler::new_not_nested(
+        unsafe { core::mem::transmute::<*const (), extern "C" fn()>(swint2_handler as *const ()) },
+        esp_hal::interrupt::Priority::Priority1,
+    );
+
+    unsafe { SoftwareInterrupt::<2>::steal() }.set_interrupt_handler(swint2_handler);
 }
 
 pub(crate) fn disable_multitasking() {
     interrupt::disable(esp_hal::system::Cpu::ProCpu, Interrupt::FROM_CPU_INTR2);
 }
 
-#[unsafe(no_mangle)]
 #[esp_hal::ram]
-extern "C" fn FROM_CPU_INTR2(trap_frame: &mut TrapFrame) {
+extern "C" fn swint2_handler(trap_frame: &mut esp_hal::interrupt::TrapFrame) {
     // clear FROM_CPU_INTR2
     let swi = unsafe { SoftwareInterrupt::<2>::steal() };
     swi.reset();
@@ -32,3 +35,9 @@ pub(crate) fn yield_task() {
     let swi = unsafe { SoftwareInterrupt::<2>::steal() };
     swi.raise();
 }
+
+#[esp_hal::ram]
+pub(crate) extern "C" fn timer_tick_handler(trap_frame: &mut esp_hal::interrupt::TrapFrame) {
+    super::clear_timer_interrupt();
+    crate::task::task_switch(trap_frame);
+}
diff --git a/esp-radio-preempt-baremetal/src/timer/xtensa.rs b/esp-radio-preempt-baremetal/src/timer/xtensa.rs
index af1408cfd..a1c926444 100644
--- a/esp-radio-preempt-baremetal/src/timer/xtensa.rs
+++ b/esp-radio-preempt-baremetal/src/timer/xtensa.rs
@@ -38,3 +38,18 @@ pub(crate) fn yield_task() {
     let intr = SW_INTERRUPT;
     unsafe { core::arch::asm!("wsr.intset  {0}", in(reg) intr, options(nostack)) };
 }
+
+#[esp_hal::ram]
+pub(crate) extern "C" fn timer_tick_handler(_context: &mut TrapFrame) {
+    super::clear_timer_interrupt();
+
+    // `task_switch` must be called on a single interrupt priority level only.
+    // Because on ESP32 the software interrupt is triggered at priority 3 but
+    // the timer interrupt is triggered at priority 1, we need to trigger the
+    // software interrupt manually.
+    if cfg!(esp32) {
+        yield_task();
+    } else {
+        crate::task::task_switch(_context);
+    }
+}
diff --git a/esp-riscv-rt/src/lib.rs b/esp-riscv-rt/src/lib.rs
index c05092db0..6e2afe451 100644
--- a/esp-riscv-rt/src/lib.rs
+++ b/esp-riscv-rt/src/lib.rs
@@ -113,49 +113,6 @@ pub struct TrapFrame {
     /// Argument register a7, typically used to pass the eighth argument to a
     /// function.
     pub a7: usize,
-    /// Saved register s0, used to hold values across function calls.
-    pub s0: usize,
-    /// Saved register s1, used to hold values across function calls.
-    pub s1: usize,
-    /// Saved register s2, used to hold values across function calls.
-    pub s2: usize,
-    /// Saved register s3, used to hold values across function calls.
-    pub s3: usize,
-    /// Saved register s4, used to hold values across function calls.
-    pub s4: usize,
-    /// Saved register s5, used to hold values across function calls.
-    pub s5: usize,
-    /// Saved register s6, used to hold values across function calls.
-    pub s6: usize,
-    /// Saved register s7, used to hold values across function calls.
-    pub s7: usize,
-    /// Saved register s8, used to hold values across function calls.
-    pub s8: usize,
-    /// Saved register s9, used to hold values across function calls.
-    pub s9: usize,
-    /// Saved register s10, used to hold values across function calls.
-    pub s10: usize,
-    /// Saved register s11, used to hold values across function calls.
-    pub s11: usize,
-    /// Global pointer register, holds the address of the global data area.
-    pub gp: usize,
-    /// Thread pointer register, holds the address of the thread-local storage
-    /// area.
-    pub tp: usize,
-    /// Stack pointer register, holds the address of the top of the stack.
-    pub sp: usize,
-    /// Program counter, stores the address of the next instruction to be
-    /// executed.
-    pub pc: usize,
-    /// Machine status register, holds the current status of the processor,
-    /// including interrupt enable bits and privilege mode.
-    pub mstatus: usize,
-    /// Machine cause register, contains the reason for the trap (e.g.,
-    /// exception or interrupt number).
-    pub mcause: usize,
-    /// Machine trap value register, contains additional information about the
-    /// trap (e.g., faulting address).
-    pub mtval: usize,
 }
 
 /// Trap entry point rust (_start_trap_rust)
@@ -514,162 +471,162 @@ _start_trap: // Handle exceptions in vectored mode
     csrr t0, mscratch
     // now SP is in RAM - continue
 
-    addi sp, sp, -40*4
+    addi sp, sp, -16*4
     sw ra, 0(sp)
     la ra, _start_trap_rust_hal /* Load the HAL trap handler */
     j _start_trap_direct
 _start_trap1:
-    addi sp, sp, -40*4
+    addi sp, sp, -16*4
     sw ra, 0(sp)
     la ra, interrupt1
     j _start_trap_direct
 _start_trap2:
-    addi sp, sp, -40*4
+    addi sp, sp, -16*4
     sw ra, 0(sp)
     la ra, interrupt2
     j _start_trap_direct
 _start_trap3:
-    addi sp, sp, -40*4
+    addi sp, sp, -16*4
     sw ra, 0(sp)
     la ra, interrupt3
     j _start_trap_direct
 _start_trap4:
-    addi sp, sp, -40*4
+    addi sp, sp, -16*4
     sw ra, 0(sp)
     la ra, interrupt4
     j _start_trap_direct
 _start_trap5:
-    addi sp, sp, -40*4
+    addi sp, sp, -16*4
     sw ra, 0(sp)
     la ra, interrupt5
     j _start_trap_direct
 _start_trap6:
-    addi sp, sp, -40*4
+    addi sp, sp, -16*4
     sw ra, 0(sp)
     la ra, interrupt6
     j _start_trap_direct
 _start_trap7:
-    addi sp, sp, -40*4
+    addi sp, sp, -16*4
     sw ra, 0(sp)
     la ra, interrupt7
     j _start_trap_direct
 _start_trap8:
-    addi sp, sp, -40*4
+    addi sp, sp, -16*4
     sw ra, 0(sp)
     la ra, interrupt8
     j _start_trap_direct
 _start_trap9:
-    addi sp, sp, -40*4
+    addi sp, sp, -16*4
     sw ra, 0(sp)
     la ra, interrupt9
     j _start_trap_direct
 _start_trap10:
-    addi sp, sp, -40*4
+    addi sp, sp, -16*4
     sw ra, 0(sp)
     la ra, interrupt10
     j _start_trap_direct
 _start_trap11:
-    addi sp, sp, -40*4
+    addi sp, sp, -16*4
     sw ra, 0(sp)
     la ra, interrupt11
     j _start_trap_direct
 _start_trap12:
-    addi sp, sp, -40*4
+    addi sp, sp, -16*4
     sw ra, 0(sp)
     la ra, interrupt12
     j _start_trap_direct
 _start_trap13:
-    addi sp, sp, -40*4
+    addi sp, sp, -16*4
     sw ra, 0(sp)
     la ra, interrupt13
     j _start_trap_direct
 _start_trap14:
-    addi sp, sp, -40*4
+    addi sp, sp, -16*4
     sw ra, 0(sp)
     la ra, interrupt14
     j _start_trap_direct
 _start_trap15:
-    addi sp, sp, -40*4
+    addi sp, sp, -16*4
     sw ra, 0(sp)
     la ra, interrupt15
     j _start_trap_direct
 _start_trap16:
-    addi sp, sp, -40*4
+    addi sp, sp, -16*4
     sw ra, 0(sp)
     la ra, interrupt16
     j _start_trap_direct
 _start_trap17:
-    addi sp, sp, -40*4
+    addi sp, sp, -16*4
     sw ra, 0(sp)
     la ra, interrupt17
     j _start_trap_direct
 _start_trap18:
-    addi sp, sp, -40*4
+    addi sp, sp, -16*4
     sw ra, 0(sp)
     la ra, interrupt18
     j _start_trap_direct
 _start_trap19:
-    addi sp, sp, -40*4
+    addi sp, sp, -16*4
     sw ra, 0(sp)
     la ra, interrupt19
     j _start_trap_direct
 _start_trap20:
-    addi sp, sp, -40*4
+    addi sp, sp, -16*4
     sw ra, 0(sp)
     la ra, interrupt20
     j _start_trap_direct
 _start_trap21:
-    addi sp, sp, -40*4
+    addi sp, sp, -16*4
     sw ra, 0(sp)
     la ra, interrupt21
     j _start_trap_direct
 _start_trap22:
-    addi sp, sp, -40*4
+    addi sp, sp, -16*4
     sw ra, 0(sp)
     la ra, interrupt22
     j _start_trap_direct
 _start_trap23:
-    addi sp, sp, -40*4
+    addi sp, sp, -16*4
     sw ra, 0(sp)
     la ra, interrupt23
     j _start_trap_direct
 _start_trap24:
-    addi sp, sp, -40*4
+    addi sp, sp, -16*4
     sw ra, 0(sp)
     la ra, interrupt24
     j _start_trap_direct
 _start_trap25:
-    addi sp, sp, -40*4
+    addi sp, sp, -16*4
     sw ra, 0(sp)
     la ra, interrupt25
     j _start_trap_direct
 _start_trap26:
-    addi sp, sp, -40*4
+    addi sp, sp, -16*4
     sw ra, 0(sp)
     la ra, interrupt26
     j _start_trap_direct
 _start_trap27:
-    addi sp, sp, -40*4
+    addi sp, sp, -16*4
     sw ra, 0(sp)
     la ra, interrupt27
     j _start_trap_direct
 _start_trap28:
-    addi sp, sp, -40*4
+    addi sp, sp, -16*4
     sw ra, 0(sp)
     la ra, interrupt28
     j _start_trap_direct
 _start_trap29:
-    addi sp, sp, -40*4
+    addi sp, sp, -16*4
     sw ra, 0(sp)
     la ra, interrupt29
     j _start_trap_direct
 _start_trap30:
-    addi sp, sp, -40*4
+    addi sp, sp, -16*4
     sw ra, 0(sp)
     la ra, interrupt30
     j _start_trap_direct
 _start_trap31:
-    addi sp, sp, -40*4
+    addi sp, sp, -16*4
     sw ra, 0(sp)
     la ra, interrupt31
     j _start_trap_direct
@@ -692,59 +649,10 @@ r#"
     sw a5, 13*4(sp)
     sw a6, 14*4(sp)
     sw a7, 15*4(sp)
-    sw s0, 16*4(sp)
-    sw s1, 17*4(sp)
-    sw s2, 18*4(sp)
-    sw s3, 19*4(sp)
-    sw s4, 20*4(sp)
-    sw s5, 21*4(sp)
-    sw s6, 22*4(sp)
-    sw s7, 23*4(sp)
-    sw s8, 24*4(sp)
-    sw s9, 25*4(sp)
-    sw s10, 26*4(sp)
-    sw s11, 27*4(sp)
-    sw gp, 28*4(sp)
-    sw tp, 29*4(sp)
-    csrrs t1, mepc, x0
-    sw t1, 31*4(sp)
-    csrrs t1, mstatus, x0
-    sw t1, 32*4(sp)
-    csrrs t1, mcause, x0
-    sw t1, 33*4(sp)
-    csrrs t1, mtval, x0
-    sw t1, 34*4(sp)
 
-    addi s0, sp, 40*4
-    sw s0, 30*4(sp)
-
-    add a0, sp, zero
-    "#,
-    // store current priority, set threshold, enable interrupts
-    r#"
-    addi sp, sp, -16 #build stack
-    sw ra, 0(sp)
-    jal ra, _handle_priority
-    lw ra, 0(sp)
-    sw a0, 0(sp) #reuse old stack, a0 is return of _handle_priority
-    addi a0, sp, 16 #the proper stack pointer is an argument to the HAL handler
-    "#,
-    // jump to handler loaded in direct handler
-    r#"
-    jalr ra, ra #jump to label loaded in _start_trapx
-    "#,
-    // restore threshold
-    r#"
-    lw a0, 0(sp) #load stored priority
-    jal ra, _restore_priority
-    addi sp, sp, 16 #pop
-    "#,
-    r#"
-    lw t1, 31*4(sp)
-    csrrw x0, mepc, t1
-
-    lw t1, 32*4(sp)
-    csrrw x0, mstatus, t1
+    # jump to handler loaded in direct handler
+    add a0, sp, zero # load trap-frame address in a0
+    jalr ra, ra # jump to label loaded in _start_trapX
 
     lw ra, 0*4(sp)
     lw t0, 1*4(sp)
@@ -762,21 +670,8 @@ r#"
     lw a5, 13*4(sp)
     lw a6, 14*4(sp)
     lw a7, 15*4(sp)
-    lw s0, 16*4(sp)
-    lw s1, 17*4(sp)
-    lw s2, 18*4(sp)
-    lw s3, 19*4(sp)
-    lw s4, 20*4(sp)
-    lw s5, 21*4(sp)
-    lw s6, 22*4(sp)
-    lw s7, 23*4(sp)
-    lw s8, 24*4(sp)
-    lw s9, 25*4(sp)
-    lw s10, 26*4(sp)
-    lw s11, 27*4(sp)
-    lw gp, 28*4(sp)
-    lw tp, 29*4(sp)
-    lw sp, 30*4(sp)
+
+    addi sp, sp, 16*4
 
     # SP was restored from the original SP
     mret