From 67d5a4026509f8207497aa7701b93e83a951a688 Mon Sep 17 00:00:00 2001
From: Kaspar Schleiser <kaspar@schleiser.de>
Date: Fri, 15 Mar 2024 11:06:29 +0100
Subject: [PATCH] riot-rs-threads: cortex_m: make `sched()` less hacky

---
 src/riot-rs-threads/src/arch/cortex_m.rs | 95 +++++++++++-------------
 1 file changed, 42 insertions(+), 53 deletions(-)

diff --git a/src/riot-rs-threads/src/arch/cortex_m.rs b/src/riot-rs-threads/src/arch/cortex_m.rs
index 559977044..21283937a 100644
--- a/src/riot-rs-threads/src/arch/cortex_m.rs
+++ b/src/riot-rs-threads/src/arch/cortex_m.rs
@@ -2,7 +2,6 @@ use super::Arch;
 use core::arch::asm;
 use core::ptr::write_volatile;
 use cortex_m::peripheral::SCB;
-use critical_section::CriticalSection;
 
 use crate::{cleanup, THREADS};
 
@@ -75,9 +74,7 @@ unsafe extern "C" fn PendSV() {
     unsafe {
         asm!(
             "
-            cpsid i
             bl {sched}
-            cpsie i
             cmp r0, #0
             /* label rules:
              * - number only
@@ -108,9 +105,7 @@ unsafe extern "C" fn PendSV() {
     unsafe {
         asm!(
             "
-            cpsid i
             bl sched
-            cpsie i
             cmp r0, #0
             beq 99f
 
@@ -168,55 +163,49 @@ unsafe extern "C" fn PendSV() {
 /// This function is called in PendSV.
 // TODO: make arch independent, or move to arch
 #[no_mangle]
-unsafe fn sched() -> usize {
-    // SAFETY: interrupts are disabled by caller
-    let cs = unsafe { CriticalSection::new() };
-    let next_pid;
-
+unsafe fn sched() -> u128 {
     loop {
-        {
-            if let Some(pid) = (unsafe { &*THREADS.as_ptr(cs) }).runqueue.get_next() {
-                next_pid = pid;
-                break;
-            }
-        }
-        //pm_set_lowest();
-        cortex_m::asm::wfi();
-        unsafe { cortex_m::interrupt::enable() };
-        cortex_m::asm::isb();
-        // pending interrupts would now get to run their ISRs
-        cortex_m::interrupt::disable();
-    }
-
-    let threads = unsafe { &mut *THREADS.as_ptr(cs) };
-    let current_high_regs;
-
-    if let Some(current_pid) = threads.current_pid() {
-        if next_pid == current_pid {
-            return 0;
+        if let Some(res) = critical_section::with(|cs| {
+            let threads = unsafe { &mut *THREADS.as_ptr(cs) };
+            let next_pid = match threads.runqueue.get_next() {
+                Some(pid) => pid,
+                None => {
+                    cortex_m::asm::wfi();
+                    return None;
+                }
+            };
+
+            let current_high_regs;
+            if let Some(current_pid) = threads.current_pid() {
+                if next_pid == current_pid {
+                    return Some(0);
+                }
+
+                threads.threads[current_pid as usize].sp = cortex_m::register::psp::read() as usize;
+                threads.current_thread = Some(next_pid);
+
+                current_high_regs = threads.threads[current_pid as usize].data.as_ptr();
+            } else {
+                threads.current_thread = Some(next_pid);
+                current_high_regs = core::ptr::null();
+            };
+
+            let next = &threads.threads[next_pid as usize];
+            let next_sp = next.sp as usize;
+            let next_high_regs = next.data.as_ptr() as usize;
+
+            // PendSV expects these three pointers in r0, r1 and r2:
+            // r0 = &next.sp
+            // r1 = &current.high_regs
+            // r2 = &next.high_regs
+            // On Cortex-M, a u128 as return value is passed in registers r0-r3.
+            // So let's use that.
+            let res: u128 =
+                //  (r0)                     (r1)                        (r2)
+                (next_sp as u128) |  ((current_high_regs as u128) << 32) | ((next_high_regs as u128) << 64);
+            Some(res)
+        }) {
+            break res;
         }
-        //println!("current: {} next: {}", current_pid, next_pid);
-        threads.threads[current_pid as usize].sp = cortex_m::register::psp::read() as usize;
-        threads.current_thread = Some(next_pid);
-        current_high_regs = threads.threads[current_pid as usize].data.as_ptr();
-    } else {
-        threads.current_thread = Some(next_pid);
-        current_high_regs = core::ptr::null();
     }
-
-    let next = &threads.threads[next_pid as usize];
-    let next_sp = next.sp;
-    let next_high_regs = next.data.as_ptr();
-
-    //println!("old_sp: {:x} next.sp: {:x}", old_sp, next_sp);
-
-    // PendSV expects these three pointers in r0, r1 and r2:
-    // r1= &current.high_regs
-    // r2= &next.high_regs
-    // r0 = &next.sp (implicitly done here via return value)
-    //
-    // write to registers manually, as ABI would return the values via stack
-    unsafe { asm!("", in("r1") current_high_regs, in("r2") next_high_regs) };
-
-    next_sp
 }