1 files changed, 13 insertions, 11 deletions

diff --git a/ui.h b/ui.h
index 4c54d6915..636c2ff70 100644
--- a/ui.h
+++ b/ui.h
@@ -21,6 +21,7 @@
 #include <pthread.h>
 #include <time.h>
 
+#include <functional>
 #include <string>
 
 // Abstract class for controlling the user interface during recovery.
@@ -128,17 +129,18 @@ class RecoveryUI {
 
   // --- menu display ---
 
-  // Display some header text followed by a menu of items, which appears at the top of the screen
-  // (in place of any scrolling ui_print() output, if necessary).
-  virtual void StartMenu(const char* const* headers, const char* const* items,
-                         int initial_selection) = 0;
-
-  // Sets the menu highlight to the given index, wrapping if necessary. Returns the actual item
-  // selected.
-  virtual int SelectMenu(int sel) = 0;
-
-  // Ends menu mode, resetting the text overlay so that ui_print() statements will be displayed.
-  virtual void EndMenu() = 0;
+  // Displays a menu with the given 'headers' and 'items'. The supplied 'key_handler' callback,
+  // which is typically bound to Device::HandleMenuKey(), should return the expected action for the
+  // given key code and menu visibility (e.g. to move the cursor or to select an item). Caller sets
+  // 'menu_only' to true to ensure only a menu item gets selected and returned. Otherwise if
+  // 'menu_only' is false, ShowMenu() will forward any non-negative value returned from the
+  // key_handler, which may be beyond the range of menu items. This could be used to trigger a
+  // device-specific action, even without that being listed in the menu. Caller needs to handle
+  // such a case accordingly (e.g. by calling Device::InvokeMenuItem() to process the action).
+  // Returns a non-negative value (the chosen item number or device-specific action code), or -1 if
+  // timed out waiting for input.
+  virtual int ShowMenu(const char* const* headers, const char* const* items, int initial_selection,
+                       bool menu_only, const std::function<int(int, bool)>& key_handler) = 0;
 
  protected:
   void EnqueueKey(int key_code);