8 files changed, 196 insertions, 66 deletions

diff --git a/recovery.cpp b/recovery.cpp
index 3ea282fc0..d7bc6fd2f 100644
--- a/recovery.cpp
+++ b/recovery.cpp
@@ -397,23 +397,22 @@ static bool wipe_data(Device* device) {
 
 static InstallResult prompt_and_wipe_data(Device* device) {
   // Use a single string and let ScreenRecoveryUI handles the wrapping.
-  std::vector<std::string> headers{
+  std::vector<std::string> wipe_data_menu_headers{
     "Can't load Android system. Your data may be corrupt. "
     "If you continue to get this message, you may need to "
     "perform a factory data reset and erase all user data "
     "stored on this device.",
   };
   // clang-format off
-  std::vector<std::string> items {
+  std::vector<std::string> wipe_data_menu_items {
     "Try again",
     "Factory data reset",
   };
   // clang-format on
   for (;;) {
-    size_t chosen_item = ui->ShowMenu(
-        headers, items, 0, true,
+    size_t chosen_item = ui->ShowPromptWipeDataMenu(
+        wipe_data_menu_headers, wipe_data_menu_items,
         std::bind(&Device::HandleMenuKey, device, std::placeholders::_1, std::placeholders::_2));
-
     // If ShowMenu() returned RecoveryUI::KeyError::INTERRUPTED, WaitKey() was interrupted.
     if (chosen_item == static_cast<size_t>(RecoveryUI::KeyError::INTERRUPTED)) {
       return INSTALL_KEY_INTERRUPTED;
@@ -421,6 +420,8 @@ static InstallResult prompt_and_wipe_data(Device* device) {
     if (chosen_item != 1) {
       return INSTALL_SUCCESS;  // Just reboot, no wipe; not a failure, user asked for it
     }
+
+    // TODO(xunchang) localize the confirmation texts also.
     if (ask_to_wipe_data(device)) {
       if (wipe_data(device)) {
         return INSTALL_SUCCESS;
diff --git a/screen_ui.cpp b/screen_ui.cpp
index 181d58ecf..c538815b5 100644
--- a/screen_ui.cpp
+++ b/screen_ui.cpp
@@ -197,12 +197,9 @@ int TextMenu::DrawItems(int x, int y, int screen_width, bool long_press) const {
   return offset;
 }
 
-GraphicMenu::GraphicMenu(size_t max_width, size_t max_height, GRSurface* graphic_headers,
-                         const std::vector<GRSurface*>& graphic_items, size_t initial_selection,
-                         const DrawInterface& draw_funcs)
+GraphicMenu::GraphicMenu(GRSurface* graphic_headers, const std::vector<GRSurface*>& graphic_items,
+                         size_t initial_selection, const DrawInterface& draw_funcs)
     : Menu(initial_selection, draw_funcs),
-      max_width_(max_width),
-      max_height_(max_height),
       graphic_headers_(graphic_headers),
       graphic_items_(graphic_items) {}
 
@@ -223,6 +220,7 @@ int GraphicMenu::Select(int sel) {
 }
 
 int GraphicMenu::DrawHeader(int x, int y) const {
+  draw_funcs_.SetColor(UIElement::HEADER);
   draw_funcs_.DrawTextIcon(x, y, graphic_headers_);
   return graphic_headers_->height;
 }
@@ -253,15 +251,16 @@ int GraphicMenu::DrawItems(int x, int y, int screen_width, bool long_press) cons
   return offset;
 }
 
-bool GraphicMenu::Validate() const {
+bool GraphicMenu::Validate(size_t max_width, size_t max_height, GRSurface* graphic_headers,
+                           const std::vector<GRSurface*>& graphic_items) {
   int offset = 0;
-  if (!ValidateGraphicSurface(offset, graphic_headers_)) {
+  if (!ValidateGraphicSurface(max_width, max_height, offset, graphic_headers)) {
     return false;
   }
-  offset += graphic_headers_->height;
+  offset += graphic_headers->height;
 
-  for (const auto& item : graphic_items_) {
-    if (!ValidateGraphicSurface(offset, item)) {
+  for (const auto& item : graphic_items) {
+    if (!ValidateGraphicSurface(max_width, max_height, offset, item)) {
       return false;
     }
     offset += item->height;
@@ -270,7 +269,8 @@ bool GraphicMenu::Validate() const {
   return true;
 }
 
-bool GraphicMenu::ValidateGraphicSurface(int y, const GRSurface* surface) const {
+bool GraphicMenu::ValidateGraphicSurface(size_t max_width, size_t max_height, int y,
+                                         const GRSurface* surface) {
   if (!surface) {
     fprintf(stderr, "Graphic surface can not be null");
     return false;
@@ -282,11 +282,11 @@ bool GraphicMenu::ValidateGraphicSurface(int y, const GRSurface* surface) const
     return false;
   }
 
-  if (surface->width > max_width_ || surface->height > max_height_ - y) {
+  if (surface->width > max_width || surface->height > max_height - y) {
     fprintf(stderr,
             "Graphic surface doesn't fit into the screen. width: %d, height: %d, max_width: %zu,"
             " max_height: %zu, vertical offset: %d\n",
-            surface->width, surface->height, max_width_, max_height_, y);
+            surface->width, surface->height, max_width, max_height, y);
     return false;
   }
 
@@ -697,7 +697,6 @@ void ScreenRecoveryUI::draw_menu_and_text_buffer_locked(
     const std::vector<std::string>& help_message) {
   int y = margin_height_;
   if (menu_) {
-    static constexpr int kMenuIndent = 4;
     int x = margin_width_ + kMenuIndent;
 
     SetColor(UIElement::INFO);
@@ -836,6 +835,16 @@ bool ScreenRecoveryUI::InitTextParams() {
   return true;
 }
 
+// TODO(xunchang) load localized text icons for the menu. (Init for screenRecoveryUI but
+// not wearRecoveryUI).
+bool ScreenRecoveryUI::LoadWipeDataMenuText() {
+  wipe_data_menu_header_text_ = nullptr;
+  factory_data_reset_text_ = nullptr;
+  try_again_text_ = nullptr;
+
+  return true;
+}
+
 bool ScreenRecoveryUI::Init(const std::string& locale) {
   RecoveryUI::Init(locale);
 
@@ -876,6 +885,8 @@ bool ScreenRecoveryUI::Init(const std::string& locale) {
   LoadLocalizedBitmap("no_command_text", &no_command_text);
   LoadLocalizedBitmap("error_text", &error_text);
 
+  LoadWipeDataMenuText();
+
   LoadAnimation();
 
   // Keep the progress bar updated, even when the process is otherwise busy.
@@ -1104,14 +1115,36 @@ void ScreenRecoveryUI::ShowFile(const std::string& filename) {
   text_row_ = old_text_row;
 }
 
-void ScreenRecoveryUI::StartMenu(const std::vector<std::string>& headers,
-                                 const std::vector<std::string>& items, size_t initial_selection) {
-  std::lock_guard<std::mutex> lg(updateMutex);
+std::unique_ptr<Menu> ScreenRecoveryUI::CreateMenu(GRSurface* graphic_header,
+                                                   const std::vector<GRSurface*>& graphic_items,
+                                                   const std::vector<std::string>& text_headers,
+                                                   const std::vector<std::string>& text_items,
+                                                   size_t initial_selection) const {
+  // horizontal unusable area: margin width + menu indent
+  size_t max_width = ScreenWidth() - margin_width_ - kMenuIndent;
+  // vertical unusable area: margin height + title lines + helper message + high light bar.
+  // It is safe to reserve more space.
+  size_t max_height = ScreenHeight() - margin_height_ - char_height_ * (title_lines_.size() + 3);
+  if (GraphicMenu::Validate(max_width, max_height, graphic_header, graphic_items)) {
+    return std::make_unique<GraphicMenu>(graphic_header, graphic_items, initial_selection, *this);
+  }
+
+  fprintf(stderr, "Failed to initialize graphic menu, falling back to use the text menu.\n");
+
+  return CreateMenu(text_headers, text_items, initial_selection);
+}
+
+std::unique_ptr<Menu> ScreenRecoveryUI::CreateMenu(const std::vector<std::string>& text_headers,
+                                                   const std::vector<std::string>& text_items,
+                                                   size_t initial_selection) const {
   if (text_rows_ > 0 && text_cols_ > 1) {
-    menu_ = std::make_unique<TextMenu>(scrollable_menu_, text_rows_, text_cols_ - 1, headers, items,
-                                       initial_selection, char_height_, *this);
-    update_screen_locked();
+    return std::make_unique<TextMenu>(scrollable_menu_, text_rows_, text_cols_ - 1, text_headers,
+                                      text_items, initial_selection, char_height_, *this);
   }
+
+  fprintf(stderr, "Failed to create text menu, text_rows %zu, text_cols %zu.\n", text_rows_,
+          text_cols_);
+  return nullptr;
 }
 
 int ScreenRecoveryUI::SelectMenu(int sel) {
@@ -1127,17 +1160,7 @@ int ScreenRecoveryUI::SelectMenu(int sel) {
   return sel;
 }
 
-void ScreenRecoveryUI::EndMenu() {
-  std::lock_guard<std::mutex> lg(updateMutex);
-  if (menu_) {
-    menu_.reset();
-    update_screen_locked();
-  }
-}
-
-size_t ScreenRecoveryUI::ShowMenu(const std::vector<std::string>& headers,
-                                  const std::vector<std::string>& items, size_t initial_selection,
-                                  bool menu_only,
+size_t ScreenRecoveryUI::ShowMenu(std::unique_ptr<Menu>&& menu, bool menu_only,
                                   const std::function<int(int, bool)>& key_handler) {
   // Throw away keys pressed previously, so user doesn't accidentally trigger menu items.
   FlushKeys();
@@ -1146,9 +1169,13 @@ size_t ScreenRecoveryUI::ShowMenu(const std::vector<std::string>& headers,
   // menu.
   if (IsKeyInterrupted()) return static_cast<size_t>(KeyError::INTERRUPTED);
 
-  StartMenu(headers, items, initial_selection);
+  CHECK(menu != nullptr);
+
+  // Starts and displays the menu
+  menu_ = std::move(menu);
+  Redraw();
 
-  int selected = initial_selection;
+  int selected = menu_->selection();
   int chosen_item = -1;
   while (chosen_item < 0) {
     int key = WaitKey();
@@ -1160,7 +1187,8 @@ size_t ScreenRecoveryUI::ShowMenu(const std::vector<std::string>& headers,
         continue;
       } else {
         LOG(INFO) << "Timed out waiting for key input; rebooting.";
-        EndMenu();
+        menu_.reset();
+        Redraw();
         return static_cast<size_t>(KeyError::TIMED_OUT);
       }
     }
@@ -1186,10 +1214,37 @@ size_t ScreenRecoveryUI::ShowMenu(const std::vector<std::string>& headers,
     }
   }
 
-  EndMenu();
+  menu_.reset();
+  Redraw();
+
   return chosen_item;
 }
 
+size_t ScreenRecoveryUI::ShowMenu(const std::vector<std::string>& headers,
+                                  const std::vector<std::string>& items, size_t initial_selection,
+                                  bool menu_only,
+                                  const std::function<int(int, bool)>& key_handler) {
+  auto menu = CreateMenu(headers, items, initial_selection);
+  if (menu == nullptr) {
+    return initial_selection;
+  }
+
+  return ShowMenu(CreateMenu(headers, items, initial_selection), menu_only, key_handler);
+}
+
+size_t ScreenRecoveryUI::ShowPromptWipeDataMenu(const std::vector<std::string>& backup_headers,
+                                                const std::vector<std::string>& backup_items,
+                                                const std::function<int(int, bool)>& key_handler) {
+  auto wipe_data_menu =
+      CreateMenu(wipe_data_menu_header_text_, { try_again_text_, factory_data_reset_text_ },
+                 backup_headers, backup_items, 0);
+  if (wipe_data_menu == nullptr) {
+    return 0;
+  }
+
+  return ShowMenu(std::move(wipe_data_menu), true, key_handler);
+}
+
 bool ScreenRecoveryUI::IsTextVisible() {
   std::lock_guard<std::mutex> lg(updateMutex);
   int visible = show_text;
diff --git a/screen_ui.h b/screen_ui.h
index 861b3605a..b1be100f9 100644
--- a/screen_ui.h
+++ b/screen_ui.h
@@ -167,25 +167,23 @@ class GraphicMenu : public Menu {
  public:
   // Constructs a Menu instance with the given |headers|, |items| and properties. Sets the initial
   // selection to |initial_selection|.
-  GraphicMenu(size_t max_width, size_t max_height, GRSurface* graphic_headers,
-              const std::vector<GRSurface*>& graphic_items, size_t initial_selection,
-              const DrawInterface& draw_funcs);
+  GraphicMenu(GRSurface* graphic_headers, const std::vector<GRSurface*>& graphic_items,
+              size_t initial_selection, const DrawInterface& draw_funcs);
 
   int Select(int sel) override;
   int DrawHeader(int x, int y) const override;
   int DrawItems(int x, int y, int screen_width, bool long_press) const override;
 
   // Checks if all the header and items are valid GRSurfaces; and that they can fit in the area
-  // defined by |max_width_| and |max_height_|.
-  bool Validate() const;
+  // defined by |max_width| and |max_height|.
+  static bool Validate(size_t max_width, size_t max_height, GRSurface* graphic_headers,
+                       const std::vector<GRSurface*>& graphic_items);
 
- private:
   // Returns true if |surface| fits on the screen with a vertical offset |y|.
-  bool ValidateGraphicSurface(int y, const GRSurface* surface) const;
-
-  const size_t max_width_;
-  const size_t max_height_;
+  static bool ValidateGraphicSurface(size_t max_width, size_t max_height, int y,
+                                     const GRSurface* surface);
 
+ private:
   // Pointers to the menu headers and items in graphic icons. This class does not have the ownership
   // of the these objects.
   GRSurface* graphic_headers_;
@@ -238,7 +236,13 @@ class ScreenRecoveryUI : public RecoveryUI, public DrawInterface {
   // the on-device resource files and shows the localized text, for manual inspection.
   void CheckBackgroundTextImages();
 
+  // Displays the localized wipe data menu.
+  size_t ShowPromptWipeDataMenu(const std::vector<std::string>& backup_headers,
+                                const std::vector<std::string>& backup_items,
+                                const std::function<int(int, bool)>& key_handler) override;
+
  protected:
+  static constexpr int kMenuIndent = 4;
   // The margin that we don't want to use for showing texts (e.g. round screen, or screen with
   // rounded corners).
   const int margin_width_;
@@ -252,18 +256,31 @@ class ScreenRecoveryUI : public RecoveryUI, public DrawInterface {
 
   virtual bool InitTextParams();
 
-  // Displays 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 std::vector<std::string>& headers,
-                         const std::vector<std::string>& items, size_t initial_selection);
+  virtual bool LoadWipeDataMenuText();
+
+  // Creates a GraphicMenu with |graphic_header| and |graphic_items|. If the GraphicMenu isn't
+  // valid or it doesn't fit on the screen; falls back to create a TextMenu instead. If succeeds,
+  // returns a unique pointer to the created menu; otherwise returns nullptr.
+  virtual std::unique_ptr<Menu> CreateMenu(GRSurface* graphic_header,
+                                           const std::vector<GRSurface*>& graphic_items,
+                                           const std::vector<std::string>& text_headers,
+                                           const std::vector<std::string>& text_items,
+                                           size_t initial_selection) const;
+
+  // Creates a TextMenu with |text_headers| and |text_items|; and sets the menu selection to
+  // |initial_selection|.
+  virtual std::unique_ptr<Menu> CreateMenu(const std::vector<std::string>& text_headers,
+                                           const std::vector<std::string>& text_items,
+                                           size_t initial_selection) const;
+
+  // Takes the ownership of |menu| and displays it.
+  virtual size_t ShowMenu(std::unique_ptr<Menu>&& menu, bool menu_only,
+                          const std::function<int(int, bool)>& key_handler);
 
   // Sets the menu highlight to the given index, wrapping if necessary. Returns the actual item
   // selected.
   virtual int SelectMenu(int sel);
 
-  // Ends menu mode, resetting the text overlay so that ui_print() statements will be displayed.
-  virtual void EndMenu();
-
   virtual void draw_background_locked();
   virtual void draw_foreground_locked();
   virtual void draw_screen_locked();
@@ -318,6 +335,11 @@ class ScreenRecoveryUI : public RecoveryUI, public DrawInterface {
   GRSurface* installing_text;
   GRSurface* no_command_text;
 
+  // Graphs for the wipe data menu
+  GRSurface* wipe_data_menu_header_text_;
+  GRSurface* try_again_text_;
+  GRSurface* factory_data_reset_text_;
+
   GRSurface** introFrames;
   GRSurface** loopFrames;
 
diff --git a/stub_ui.h b/stub_ui.h
index a3cf12b05..ca137dff3 100644
--- a/stub_ui.h
+++ b/stub_ui.h
@@ -68,6 +68,12 @@ class StubRecoveryUI : public RecoveryUI {
     return initial_selection;
   }
 
+  size_t ShowPromptWipeDataMenu(const std::vector<std::string>& /* backup_headers */,
+                                const std::vector<std::string>& /* backup_items */,
+                                const std::function<int(int, bool)>& /* key_handle */) override {
+    return 0;
+  }
+
   void SetTitle(const std::vector<std::string>& /* lines */) override {}
 };
 
diff --git a/tests/unit/screen_ui_test.cpp b/tests/unit/screen_ui_test.cpp
index ec269503e..dc6e6a896 100644
--- a/tests/unit/screen_ui_test.cpp
+++ b/tests/unit/screen_ui_test.cpp
@@ -229,6 +229,43 @@ TEST_F(ScreenUITest, WearMenuSelectItemsOverflow) {
   ASSERT_EQ(3u, menu.MenuEnd());
 }
 
+TEST_F(ScreenUITest, GraphicMenuSelection) {
+  GRSurface fake_surface = GRSurface{ 50, 50, 50, 1, nullptr };
+  std::vector<GRSurface*> items = { &fake_surface, &fake_surface, &fake_surface };
+  GraphicMenu menu(&fake_surface, items, 0, draw_funcs_);
+
+  ASSERT_EQ(0, menu.selection());
+
+  int sel = 0;
+  for (int i = 0; i < 3; i++) {
+    sel = menu.Select(++sel);
+    ASSERT_EQ((i + 1) % 3, sel);
+    ASSERT_EQ(sel, menu.selection());
+  }
+
+  sel = 0;
+  for (int i = 0; i < 3; i++) {
+    sel = menu.Select(--sel);
+    ASSERT_EQ(2 - i, sel);
+    ASSERT_EQ(sel, menu.selection());
+  }
+}
+
+TEST_F(ScreenUITest, GraphicMenuValidate) {
+  auto fake_surface = GRSurface{ 50, 50, 50, 1, nullptr };
+  std::vector<GRSurface*> items = { &fake_surface, &fake_surface, &fake_surface };
+
+  ASSERT_TRUE(GraphicMenu::Validate(200, 200, &fake_surface, items));
+
+  // Menu exceeds the horizontal boundary.
+  auto wide_surface = GRSurface{ 300, 50, 300, 1, nullptr };
+  ASSERT_FALSE(GraphicMenu::Validate(299, 200, &wide_surface, items));
+
+  // Menu exceeds the vertical boundary.
+  items.push_back(&fake_surface);
+  ASSERT_FALSE(GraphicMenu::Validate(200, 249, &fake_surface, items));
+}
+
 static constexpr int kMagicAction = 101;
 
 enum class KeyCode : int {
diff --git a/ui.h b/ui.h
index e0fb13e40..1e6186a11 100644
--- a/ui.h
+++ b/ui.h
@@ -162,6 +162,13 @@ class RecoveryUI {
                           const std::vector<std::string>& items, size_t initial_selection,
                           bool menu_only, const std::function<int(int, bool)>& key_handler) = 0;
 
+  // Displays the localized wipe data menu with pre-generated graphs. If there's an issue
+  // with the graphs, falls back to use the backup string headers and items instead. The initial
+  // selection is the 0th item in the menu, which is expected to reboot the device without a wipe.
+  virtual size_t ShowPromptWipeDataMenu(const std::vector<std::string>& backup_headers,
+                                        const std::vector<std::string>& backup_items,
+                                        const std::function<int(int, bool)>& key_handler) = 0;
+
   // Resets the key interrupt status.
   void ResetKeyInterruptStatus() {
     key_interrupted_ = false;
diff --git a/wear_ui.cpp b/wear_ui.cpp
index 8f3bc7bbe..0611f94c9 100644
--- a/wear_ui.cpp
+++ b/wear_ui.cpp
@@ -95,13 +95,14 @@ void WearRecoveryUI::update_progress_locked() {
 
 void WearRecoveryUI::SetStage(int /* current */, int /* max */) {}
 
-void WearRecoveryUI::StartMenu(const std::vector<std::string>& headers,
-                               const std::vector<std::string>& items, size_t initial_selection) {
-  std::lock_guard<std::mutex> lg(updateMutex);
+std::unique_ptr<Menu> WearRecoveryUI::CreateMenu(const std::vector<std::string>& text_headers,
+                                                 const std::vector<std::string>& text_items,
+                                                 size_t initial_selection) const {
   if (text_rows_ > 0 && text_cols_ > 0) {
-    menu_ = std::make_unique<TextMenu>(scrollable_menu_, text_rows_ - menu_unusable_rows_ - 1,
-                                       text_cols_ - 1, headers, items, initial_selection,
-                                       char_height_, *this);
-    update_screen_locked();
+    return std::make_unique<TextMenu>(scrollable_menu_, text_rows_ - menu_unusable_rows_ - 1,
+                                      text_cols_ - 1, text_headers, text_items, initial_selection,
+                                      char_height_, *this);
   }
+
+  return nullptr;
 }
diff --git a/wear_ui.h b/wear_ui.h
index b80cfd758..429af69d2 100644
--- a/wear_ui.h
+++ b/wear_ui.h
@@ -36,8 +36,9 @@ class WearRecoveryUI : public ScreenRecoveryUI {
   // Recovery, build id and etc) and the bottom lines that may otherwise go out of the screen.
   const int menu_unusable_rows_;
 
-  void StartMenu(const std::vector<std::string>& headers, const std::vector<std::string>& items,
-                 size_t initial_selection) override;
+  std::unique_ptr<Menu> CreateMenu(const std::vector<std::string>& text_headers,
+                                   const std::vector<std::string>& text_items,
+                                   size_t initial_selection) const override;
 
   int GetProgressBaseline() const override;