ui: port menu component to async

diff --git a/src/rust/bitbox02-rust/src/workflow/cancel.rs b/src/rust/bitbox02-rust/src/workflow/cancel.rs
index cb6a901..77cbe6e 100644
--- a/src/rust/bitbox02-rust/src/workflow/cancel.rs
+++ b/src/rust/bitbox02-rust/src/workflow/cancel.rs
@@ -1,54 +1,6 @@
 // SPDX-License-Identifier: Apache-2.0
 
-use crate::bb02_async::option_no_screensaver;
-use core::cell::RefCell;
-
-use super::confirm;
-
 #[derive(Debug)]
 pub enum Error {
     Cancelled,
 }
-
-pub type ResultCell<R> = RefCell<Option<Result<R, Error>>>;
-
-/// Resolves the `with_cancel` future as cancelled.
-pub fn cancel<R>(result_cell: &ResultCell<R>) {
-    *result_cell.borrow_mut() = Some(Err(Error::Cancelled));
-}
-
-/// Resolves the `with_cancel` future with the given result.
-pub fn set_result<R>(result_cell: &ResultCell<R>, result: R) {
-    *result_cell.borrow_mut() = Some(Ok(result));
-}
-
-/// Blocks on showing/running a component until `cancel` or `result` is
-/// called on the same `result_cell`.
-/// In the former, a prompt with the given title to confirm cancellation is shown.
-///
-/// * `title` - title to show in the cancel confirm prompt.
-/// * `component` - component to process
-/// * `result_cell` - result var to synchronize the result on. Pass the same to `cancel` and
-///   `set_result`.
-pub async fn with_cancel<R>(
-    title: &str,
-    component: &mut bitbox02::ui::Component<'_>,
-    result_cell: &ResultCell<R>,
-) -> Result<R, Error> {
-    component.screen_stack_push();
-    loop {
-        let result = option_no_screensaver(result_cell).await;
-        if let Err(Error::Cancelled) = result {
-            let params = confirm::Params {
-                title,
-                body: "Do you really\nwant to cancel?",
-                ..Default::default()
-            };
-
-            if let Err(confirm::UserAbort) = confirm::confirm(&params).await {
-                continue;
-            }
-        }
-        return result;
-    }
-}
diff --git a/src/rust/bitbox02-rust/src/workflow/menu.rs b/src/rust/bitbox02-rust/src/workflow/menu.rs
index cd22f51..25d21fd 100644
--- a/src/rust/bitbox02-rust/src/workflow/menu.rs
+++ b/src/rust/bitbox02-rust/src/workflow/menu.rs
@@ -2,25 +2,19 @@
 
 pub use super::cancel::Error as CancelError;
 
-use crate::bb02_async::option_no_screensaver;
-
-use alloc::boxed::Box;
-use core::cell::RefCell;
-
 /// Returns the index of the word chosen by the user.
 pub async fn pick(words: &[&str], title: Option<&str>) -> Result<u8, CancelError> {
-    let result = RefCell::new(None as Option<Result<u8, CancelError>>);
-    let mut component = bitbox02::ui::menu_create(bitbox02::ui::MenuParams {
+    match bitbox02::ui::menu(bitbox02::ui::MenuParams {
         words,
         title,
-        select_word_cb: Some(Box::new(|choice_idx| {
-            *result.borrow_mut() = Some(Ok(choice_idx));
-        })),
-        continue_on_last_cb: None,
-        cancel_cb: Some(Box::new(|| {
-            *result.borrow_mut() = Some(Err(CancelError::Cancelled));
-        })),
-    });
-    component.screen_stack_push();
-    option_no_screensaver(&result).await
+        select_word: true,
+        continue_on_last: false,
+        cancel_confirm_title: None,
+    })
+    .await
+    {
+        bitbox02::ui::MenuResponse::SelectWord(choice_idx) => Ok(choice_idx),
+        bitbox02::ui::MenuResponse::ContinueOnLast => panic!("unexpected continue-on-last"),
+        bitbox02::ui::MenuResponse::Cancel => Err(CancelError::Cancelled),
+    }
 }
diff --git a/src/rust/bitbox02-rust/src/workflow/mnemonic.rs b/src/rust/bitbox02-rust/src/workflow/mnemonic.rs
index fddfee4..3615891 100644
--- a/src/rust/bitbox02-rust/src/workflow/mnemonic.rs
+++ b/src/rust/bitbox02-rust/src/workflow/mnemonic.rs
@@ -1,16 +1,13 @@
 // SPDX-License-Identifier: Apache-2.0
 
 pub use super::cancel::Error as CancelError;
-use super::cancel::{cancel, set_result, with_cancel};
 use super::confirm;
 use super::menu;
 use super::trinary_choice::TrinaryChoice;
 use super::trinary_input_string;
 
-use alloc::boxed::Box;
 use alloc::string::String;
 use alloc::vec::Vec;
-use core::cell::RefCell;
 
 use sha2::{Digest, Sha256};
 
@@ -65,36 +62,36 @@ fn create_random_unique_words(word: &str, length: u8) -> (u8, Vec<zeroize::Zeroi
 
 /// Displays all mnemonic words in a scroll-through screen.
 pub async fn show_mnemonic(words: &[&str]) -> Result<(), CancelError> {
-    let result = RefCell::new(None);
-    let mut component = bitbox02::ui::menu_create(bitbox02::ui::MenuParams {
+    match bitbox02::ui::menu(bitbox02::ui::MenuParams {
         words,
         title: None,
-        select_word_cb: None,
-        continue_on_last_cb: Some(Box::new(|| {
-            set_result(&result, ());
-        })),
-        cancel_cb: Some(Box::new(|| {
-            cancel(&result);
-        })),
-    });
-    with_cancel("Recovery\nwords", &mut component, &result).await
+        select_word: false,
+        continue_on_last: true,
+        cancel_confirm_title: Some("Recovery\nwords"),
+    })
+    .await
+    {
+        bitbox02::ui::MenuResponse::ContinueOnLast => Ok(()),
+        bitbox02::ui::MenuResponse::SelectWord(_) => panic!("unexpected select-word"),
+        bitbox02::ui::MenuResponse::Cancel => Err(CancelError::Cancelled),
+    }
 }
 
 /// Displays the `choices` to the user, returning the index of the selected choice.
 pub async fn confirm_word(choices: &[&str], title: &str) -> Result<u8, CancelError> {
-    let result = RefCell::new(None);
-    let mut component = bitbox02::ui::menu_create(bitbox02::ui::MenuParams {
+    match bitbox02::ui::menu(bitbox02::ui::MenuParams {
         words: choices,
         title: Some(title),
-        select_word_cb: Some(Box::new(|idx| {
-            set_result(&result, idx);
-        })),
-        continue_on_last_cb: None,
-        cancel_cb: Some(Box::new(|| {
-            cancel(&result);
-        })),
-    });
-    with_cancel("Recovery\nwords", &mut component, &result).await
+        select_word: true,
+        continue_on_last: false,
+        cancel_confirm_title: Some("Recovery\nwords"),
+    })
+    .await
+    {
+        bitbox02::ui::MenuResponse::SelectWord(choice_idx) => Ok(choice_idx),
+        bitbox02::ui::MenuResponse::ContinueOnLast => panic!("unexpected continue-on-last"),
+        bitbox02::ui::MenuResponse::Cancel => Err(CancelError::Cancelled),
+    }
 }
 
 pub async fn show_and_confirm_mnemonic(
diff --git a/src/rust/bitbox02/src/ui/types.rs b/src/rust/bitbox02/src/ui/types.rs
index 4bcd882..40b2093 100644
--- a/src/rust/bitbox02/src/ui/types.rs
+++ b/src/rust/bitbox02/src/ui/types.rs
@@ -66,12 +66,25 @@ pub struct TrinaryInputStringParams<'a> {
 pub type SelectWordCb<'a> = Box<dyn FnMut(u8) + 'a>;
 pub type ContinueCancelCb<'a> = Box<dyn FnMut() + 'a>;
 
+/// Result of a resolved menu interaction.
+pub enum MenuResponse {
+    /// User chose one of the entries by index.
+    SelectWord(u8),
+    /// User reached the last item and continued.
+    ContinueOnLast,
+    /// User cancelled the menu flow.
+    Cancel,
+}
+
 pub struct MenuParams<'a> {
     pub words: &'a [&'a str],
     pub title: Option<&'a str>,
-    pub select_word_cb: Option<SelectWordCb<'a>>,
-    pub continue_on_last_cb: Option<ContinueCancelCb<'a>>,
-    pub cancel_cb: Option<ContinueCancelCb<'a>>,
+    /// If true, selecting a word is possible and `MenuResponse::SelectWord` is returned.
+    pub select_word: bool,
+    /// If true, user can continue at the last word, and `MenuResponse::ContinueOnLast` is returned.
+    pub continue_on_last: bool,
+    /// `None` means immediate cancel. `Some(title)` asks for cancel confirmation.
+    pub cancel_confirm_title: Option<&'a str>,
 }
 
 pub type TrinaryChoiceCb<'a> = Box<dyn FnMut(TrinaryChoice) + 'a>;
diff --git a/src/rust/bitbox02/src/ui/ui.rs b/src/rust/bitbox02/src/ui/ui.rs
index 3efb699..ad4c329 100644
--- a/src/rust/bitbox02/src/ui/ui.rs
+++ b/src/rust/bitbox02/src/ui/ui.rs
@@ -1,8 +1,8 @@
 // SPDX-License-Identifier: Apache-2.0
 
 pub use super::types::{
-    AcceptRejectCb, ConfirmParams, ContinueCancelCb, Font, MenuParams, SelectWordCb, TrinaryChoice,
-    TrinaryChoiceCb, TrinaryInputStringParams,
+    AcceptRejectCb, ConfirmParams, ContinueCancelCb, Font, MenuParams, MenuResponse, SelectWordCb,
+    TrinaryChoice, TrinaryChoiceCb, TrinaryInputStringParams,
 };
 
 use core::ffi::{c_char, c_void};
@@ -299,52 +299,81 @@ where
     }
 }
 
-pub fn menu_create(params: MenuParams<'_>) -> Component<'_> {
-    unsafe extern "C" fn c_select_word_cb(word_idx: u8, user_data: *mut c_void) {
-        let callback = user_data as *mut SelectWordCb;
-        unsafe { (*callback)(word_idx) };
+pub async fn menu(params: MenuParams<'_>) -> MenuResponse {
+    let _no_screensaver = crate::screen_saver::ScreensaverInhibitor::new();
+    let cancel_confirm_title = params.cancel_confirm_title;
+
+    // Shared between the async context and the c callback
+    struct SharedState {
+        waker: Option<Waker>,
+        result: Option<MenuResponse>,
     }
+    let shared_state = Box::new(RefCell::new(SharedState {
+        waker: None,
+        result: None,
+    }));
+    let shared_state_ptr = shared_state.as_ref() as *const RefCell<SharedState> as *mut c_void;
 
-    unsafe extern "C" fn c_continue_cancel_cb(user_data: *mut c_void) {
-        let callback = user_data as *mut ContinueCancelCb;
-        unsafe { (*callback)() };
+    unsafe extern "C" fn select_word_cb(word_idx: u8, user_data: *mut c_void) {
+        let shared_state = unsafe { &*(user_data as *mut RefCell<SharedState>) };
+        let mut shared_state = shared_state.borrow_mut();
+        if shared_state.result.is_none() {
+            shared_state.result = Some(MenuResponse::SelectWord(word_idx));
+            if let Some(waker) = shared_state.waker.as_ref() {
+                waker.wake_by_ref();
+            }
+        }
+    }
+
+    unsafe extern "C" fn continue_on_last_cb(user_data: *mut c_void) {
+        let shared_state = unsafe { &*(user_data as *mut RefCell<SharedState>) };
+        let mut shared_state = shared_state.borrow_mut();
+        if shared_state.result.is_none() {
+            shared_state.result = Some(MenuResponse::ContinueOnLast);
+            if let Some(waker) = shared_state.waker.as_ref() {
+                waker.wake_by_ref();
+            }
+        }
+    }
+
+    unsafe extern "C" fn cancel_cb(user_data: *mut c_void) {
+        let shared_state = unsafe { &*(user_data as *mut RefCell<SharedState>) };
+        let mut shared_state = shared_state.borrow_mut();
+        if shared_state.result.is_none() {
+            shared_state.result = Some(MenuResponse::Cancel);
+            if let Some(waker) = shared_state.waker.as_ref() {
+                waker.wake_by_ref();
+            }
+        }
     }
 
     // We want to turn &[&str] into a C char**.
     //
-    // Step 1: create the C strings. This var has to be alive until after menu_create() finishes,
+    // Step 1: create the C strings. This var has to be alive until after menu() finishes,
     // otherwise the pointers we send to menu_create() will be invalid.
     let words: Vec<Vec<core::ffi::c_char>> = params
         .words
         .iter()
         .map(|word| util::strings::str_to_cstr_vec(word).unwrap())
         .collect();
-    // Step two: collect pointers. This var also has to be valid until menu_create() finishes, or
+    // Step two: collect pointers. This var also has to be valid until menu() finishes, or
     // the pointer will be invalid.
     let c_words: Vec<*const core::ffi::c_char> =
         words.iter().map(|word| word.as_ptr() as _).collect();
 
-    let (select_word_cb, select_word_user_data) = match params.select_word_cb {
-        None => (None, core::ptr::null_mut()),
-        Some(cb) => (
-            Some(c_select_word_cb as _),
-            Box::into_raw(Box::new(cb)) as *mut c_void,
-        ),
-    };
-
-    let (continue_on_last_cb, continue_on_last_user_data) = match params.continue_on_last_cb {
-        None => (None, core::ptr::null_mut()),
-        Some(cb) => (
-            Some(c_continue_cancel_cb as _),
-            Box::into_raw(Box::new(cb)) as *mut c_void,
+    let (select_word_cb, select_word_user_data) = match params.select_word {
+        false => (None, core::ptr::null_mut()),
+        true => (
+            Some(select_word_cb as _),
+            shared_state_ptr, // passed to select_word_cb as `user_data`.
         ),
     };
 
-    let (cancel_cb, cancel_user_data) = match params.cancel_cb {
-        None => (None, core::ptr::null_mut()),
-        Some(cb) => (
-            Some(c_continue_cancel_cb as _),
-            Box::into_raw(Box::new(cb)) as *mut c_void,
+    let (continue_on_last_cb, continue_on_last_user_data) = match params.continue_on_last {
+        false => (None, core::ptr::null_mut()),
+        true => (
+            Some(continue_on_last_cb as _),
+            shared_state_ptr, // passed to continue_on_last_cb as `user_data`.
         ),
     };
     let title = params
@@ -362,29 +391,55 @@ pub fn menu_create(params: MenuParams<'_>) -> Component<'_> {
                 .map_or_else(core::ptr::null, |title| title.as_ptr()),
             continue_on_last_cb,
             continue_on_last_user_data,
-            cancel_cb,
-            cancel_user_data,
+            Some(cancel_cb as _),
+            shared_state_ptr, // passed to cancel_cb as `user_data`.
             core::ptr::null_mut(),
         )
     };
-    Component {
+    let mut component = Component {
         component,
         is_pushed: false,
-        on_drop: Some(Box::new(move || unsafe {
-            // Drop all callbacks.
-            if !select_word_user_data.is_null() {
-                drop(Box::from_raw(select_word_user_data as *mut SelectWordCb));
-            }
-            if !continue_on_last_user_data.is_null() {
-                drop(Box::from_raw(
-                    continue_on_last_user_data as *mut ContinueCancelCb,
-                ));
-            }
-            if !cancel_user_data.is_null() {
-                drop(Box::from_raw(cancel_user_data as *mut ContinueCancelCb));
-            }
-        })),
+        on_drop: None,
         _p: PhantomData,
+    };
+    component.screen_stack_push();
+
+    loop {
+        let result = core::future::poll_fn({
+            let shared_state = &shared_state;
+            move |cx| {
+                let mut shared_state = shared_state.borrow_mut();
+
+                if let Some(result) = shared_state.result.take() {
+                    Poll::Ready(result)
+                } else {
+                    // Store the waker so the callback can wake up this task
+                    shared_state.waker = Some(cx.waker().clone());
+                    Poll::Pending
+                }
+            }
+        })
+        .await;
+
+        match result {
+            MenuResponse::SelectWord(_) | MenuResponse::ContinueOnLast => return result,
+            MenuResponse::Cancel => match cancel_confirm_title {
+                None => return MenuResponse::Cancel,
+                Some(title) => {
+                    // false means _do not cancel_, stay in the same menu component.
+                    if !confirm(&ConfirmParams {
+                        title,
+                        body: "Do you really\nwant to cancel?",
+                        ..Default::default()
+                    })
+                    .await
+                    {
+                        continue;
+                    }
+                    return MenuResponse::Cancel;
+                }
+            },
+        }
     }
 }
 
diff --git a/src/rust/bitbox02/src/ui/ui_stub.rs b/src/rust/bitbox02/src/ui/ui_stub.rs
index aeedd6b..746dc01 100644
--- a/src/rust/bitbox02/src/ui/ui_stub.rs
+++ b/src/rust/bitbox02/src/ui/ui_stub.rs
@@ -6,8 +6,8 @@
 //! workflows now.
 
 pub use super::types::{
-    AcceptRejectCb, ConfirmParams, ContinueCancelCb, Font, MenuParams, SelectWordCb, TrinaryChoice,
-    TrinaryChoiceCb, TrinaryInputStringParams,
+    AcceptRejectCb, ConfirmParams, ContinueCancelCb, Font, MenuParams, MenuResponse, SelectWordCb,
+    TrinaryChoice, TrinaryChoiceCb, TrinaryInputStringParams,
 };
 
 use core::marker::PhantomData;
@@ -63,7 +63,7 @@ where
     panic!("not used");
 }
 
-pub fn menu_create(_params: MenuParams<'_>) -> Component<'_> {
+pub async fn menu(_params: MenuParams<'_>) -> MenuResponse {
     panic!("not used");
 }
 
diff --git a/src/rust/bitbox02/src/ui/ui_stub_c_unit_tests.rs b/src/rust/bitbox02/src/ui/ui_stub_c_unit_tests.rs
index 5e53147..496bba9 100644
--- a/src/rust/bitbox02/src/ui/ui_stub_c_unit_tests.rs
+++ b/src/rust/bitbox02/src/ui/ui_stub_c_unit_tests.rs
@@ -3,8 +3,8 @@
 //! Stubs for the Bitbox02 simulator and also C unit-tests.
 
 pub use super::types::{
-    AcceptRejectCb, ConfirmParams, ContinueCancelCb, Font, MenuParams, SelectWordCb, TrinaryChoice,
-    TrinaryChoiceCb, TrinaryInputStringParams,
+    AcceptRejectCb, ConfirmParams, ContinueCancelCb, Font, MenuParams, MenuResponse, SelectWordCb,
+    TrinaryChoice, TrinaryChoiceCb, TrinaryInputStringParams,
 };
 
 use core::marker::PhantomData;
@@ -81,7 +81,7 @@ where
     }
 }
 
-pub fn menu_create(_params: MenuParams<'_>) -> Component<'_> {
+pub async fn menu(_params: MenuParams<'_>) -> MenuResponse {
     panic!("not implemented");
 }