Simple block_on that doesn't just spin

[Dominaezzz]

Would esp-hal be interested in having a simple block_on that uses waiti/wfi to waits for futures?

I've been using embassy-futures::block_on for simple projects and tests, but I think I shouldn't have to pull in another dependency (which spins, 100% CPU, as it's hardware agnostic) to just do run some async code. I also don't want to opt in to the "whole" embassy framework just to wait on a single future (#1035). The HAL should have a simple way to efficiently wait on a single future on its hardware.

I want to contribute this but want to make sure it would be welcome, or if there are different plans, or if embassy is simply the way to do async with esp-hal.

Side note: I was wondering if it'd make sense for the blocking drivers to use some version of this. They shouldn't be spinning when waiti/wfi exists. Now that interrupt handlers can be registered at runtime, perhaps the blocking drivers can register one temporarily to just block efficiently, just an idea.

Another side note: With the right reactors set up, it may be possible to have a Rust program automatically go into proper light sleep if the set of all waiting wakers are waiting for light sleep compatible timers and/or gpio pins. (Though this is mostly out of scope of this issue)

[Dominaezzz]

Xtensa implementation.

pub fn block_on<F: Future>(mut fut: F) -> F::Output {
    let state = TaskState::new();

    let raw_waker = state.raw_waker();
    let waker = unsafe { Waker::from_raw(raw_waker) };

    let res = {
        let mut cx = Context::from_waker(&waker);
        let mut fut = pin!(fut);
        loop {
            if let Poll::Ready(res) = fut.as_mut().poll(&mut cx) {
                break res;
            } else {
                state.wait_until_woken();
            }
        }
    };

    let clone_count = state.clone_count.load(Ordering::SeqCst);
    // This method cannot return if there are clones since TaskState is created on the stack.
    assert_eq!(
        clone_count, 0,
        "Waker was leaked in {} places!",
        clone_count
    );

    res
}

struct TaskState {
    work_available: AtomicBool,
    clone_count: AtomicUsize,
    core: Cpu,
}

impl TaskState {
    fn new() -> Self {
        Self {
            work_available: AtomicBool::new(false),
            clone_count: AtomicUsize::new(0),
            core: get_core(),
        }
    }

    fn clone(&self) {
        self.clone_count.fetch_add(1, Ordering::SeqCst);
    }

    fn wake(&self) {
        self.work_available.store(true, Ordering::SeqCst);
        if self.core != get_core() {
            // This could be implemented with software interrupts.
            // SoftwareInterrupt<..> methods would need to take & instead of &mut.
            unimplemented!()
        }
    }

    fn drop(&self) {
        self.clone_count.fetch_sub(1, Ordering::SeqCst);
    }

    fn wait_until_woken(&self) {
        wait_for_interrupts_until(|| self.work_available.swap(false, Ordering::SeqCst));
    }

    const VTABLE: RawWakerVTable = RawWakerVTable::new(
        Self::raw_clone,
        Self::raw_wake,
        Self::raw_wake_by_ref,
        Self::raw_drop,
    );

    fn raw_waker(&self) -> RawWaker {
        RawWaker::new(self as *const Self as _, &Self::VTABLE)
    }

    fn raw_clone(src: *const ()) -> RawWaker {
        (unsafe { &*(src as *const Self) }).clone();
        RawWaker::new(src, &Self::VTABLE)
    }

    fn raw_wake(src: *const ()) {
        let state = unsafe { &*(src as *const Self) };
        state.wake();
        state.drop();
    }

    fn raw_wake_by_ref(src: *const ()) {
        unsafe { &*(src as *const Self) }.wake();
    }

    fn raw_drop(src: *const ()) {
        unsafe { &*(src as *const Self) }.drop();
    }
}

fn wait_for_interrupts_until(mut condition: impl FnMut() -> bool) {
    use core::arch::asm;

    if condition() {
        // In case of spin, skip any PS.INTLEVEL tinkering code.
        return;
    }

    loop {
        // Previous PS (Processor Status) Register.
        let prev_ps: u32;

        // RSIL - Read and Set Interrupt Level
        //
        // Read the current value of the PS register and set PS.INTLEVEL to 5 to prevent
        // interrupts (that we're about to wait for) from firing.
        // The highest level on the ESP32 is 5. Xtensa ISA supports up to 15.
        unsafe { asm!("rsil {0}, 5", out(reg) prev_ps) };

        if condition() {
            // WSR - Write Special Register
            //
            // Set PS.INTLEVEL back to what it was before.
            unsafe {
                asm!(
                    "wsr.ps {0}",
                    "rsync",
                    in(reg) prev_ps
                );
            }

            // We can stop waiting, there is work available.
            break;
        } else {
            // WAITI - Wait for Interrupt
            //
            // Set PS.INTLEVEL back to what it previously was and suspend this core until an
            // interrupt above the existing PS.INTLEVEL occurs.

            let int_level = prev_ps & 0x00000F;
            match int_level {
                0 => unsafe { asm!("waiti 0") },
                1 => unsafe { asm!("waiti 1") },
                2 => unsafe { asm!("waiti 2") },
                3 => unsafe { asm!("waiti 3") },
                4 => unsafe { asm!("waiti 4") },
                5 => unsafe { asm!("waiti 5") },
                _ => unreachable!(),
            }
        }
    }
}