Ensure default dialog button focus remains after a dialog update.

ui/views/window/dialog_client_view.cc

 // Copyright (c) 2012 The Chromium Authors. All rights reserved.
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file.
 
 #include "ui/views/window/dialog_client_view.h"
 
 #include <algorithm>
 
 #include "build/build_config.h"
 #include "ui/base/material_design/material_design_controller.h"
 #include "ui/events/keycodes/keyboard_codes.h"
 #include "ui/views/background.h"
 #include "ui/views/border.h"
 #include "ui/views/controls/button/blue_button.h"
 #include "ui/views/controls/button/checkbox.h"
 #include "ui/views/controls/button/custom_button.h"
 #include "ui/views/controls/button/label_button.h"
 #include "ui/views/controls/button/md_text_button.h"
 #include "ui/views/layout/grid_layout.h"
 #include "ui/views/layout/layout_provider.h"
 #include "ui/views/style/platform_style.h"
+#include "ui/views/view_observer.h"
 #include "ui/views/widget/widget.h"
 #include "ui/views/window/dialog_delegate.h"
 
 namespace views {
 
 namespace {
 
 // The group used by the buttons.  This name is chosen voluntarily big not to
 // conflict with other groups that could be in the dialog content.
 const int kButtonGroup = 6666;
@@ skipping 11 matching lines @@
 }
 
 // Returns the bounding box required to contain |size1| and |size2|, placed one
 // atop the other.
 gfx::Size GetBoundingSizeForVerticalStack(const gfx::Size& size1,
                                           const gfx::Size& size2) {
   return gfx::Size(std::max(size1.width(), size2.width()),
                    size1.height() + size2.height());
 }
 
+// ViewDeletionObserver implements an observer to track the deletion of the
+// view in focus.
+class ViewDeletionObserver : public ViewObserver {
+ public:
+  explicit ViewDeletionObserver(View* observed_view)
+      : observed_view_(observed_view) {
+    if (observed_view_)
+      observed_view_->AddObserver(this);
+  }
+
+  ~ViewDeletionObserver() override {
+    if (observed_view_)
+      observed_view_->RemoveObserver(this);
+  }
+
+  // ViewObserver:
+  void OnViewIsDeleting(View* observed_view) override {
+    DCHECK_EQ(observed_view, observed_view_);
+    observed_view_ = nullptr;
+    observed_view_->RemoveObserver(this);
+  }
+
+  View* observed_view() { return observed_view_; }
+
+ private:
+  View* observed_view_ = nullptr;
+
+  DISALLOW_COPY_AND_ASSIGN(ViewDeletionObserver);
+};
+
 }  // namespace
 
 // Simple container to bubble child view changes up the view hierarchy.
 class DialogClientView::ButtonRowContainer : public View {
  public:
   explicit ButtonRowContainer(DialogClientView* owner) : owner_(owner) {}
 
   // View:
   void ChildPreferredSizeChanged(View* child) override {
     owner_->ChildPreferredSizeChanged(child);
@@ skipping 256 matching lines @@
   View* third = ok_button_;
   if (kIsOkButtonOnLeftSide)
     std::swap(second, third);
   return {{first, second, third}};
 }
 
 void DialogClientView::SetupLayout() {
   base::AutoReset<bool> auto_reset(&adding_or_removing_views_, true);
   GridLayout* layout = new GridLayout(button_row_container_);
   layout->set_minimum_size(minimum_size_);
+  FocusManager* focus_manager = GetFocusManager();
+  ViewDeletionObserver deletion_observer(focus_manager->GetFocusedView());
 
   // Clobber any existing LayoutManager since it has weak references to child
   // Views which may be removed by SetupViews().
   button_row_container_->SetLayoutManager(layout);
   SetupViews();
   const std::array<View*, kNumButtons> views = GetButtonRowViews();
 
   // Visibility changes on |extra_view_| must be observed to re-Layout. However,
   // when hidden it's not included in the button row (it can't influence layout)
   // and it can't be added to |button_row_container_| (GridLayout complains).
@@ skipping 67 matching lines @@
   // Checkbox (which extends LabelButton). Otherwise, link everything.
   bool skip_first_link =
       views[0] && (!CustomButton::AsCustomButton(views[0]) ||
                    views[0]->GetClassName() == Checkbox::kViewClassName);
   if (skip_first_link)
     column_set->LinkColumnSizes(link[1], link[2], -1);
   else
     column_set->LinkColumnSizes(link[0], link[1], link[2], -1);
 
   layout->AddPaddingRow(kFixed, insets.bottom());
+
+  // The default focus is lost when child views are added back into the dialog.
+  // This restores focus if the button is still available.
+  View* previously_focused_view = deletion_observer.observed_view();
+  if (previously_focused_view && !focus_manager->GetFocusedView() &&
+      Contains(previously_focused_view)) {
+    previously_focused_view->RequestFocus();
+  }
 }
 
 void DialogClientView::SetupViews() {
   button_row_container_->RemoveAllChildViews(false /* delete children */);
   // If SetupLayout() "stored" a hidden |extra_view_| in |this|, ensure it can
   // be re-added to the layout when becoming visible.
   if (extra_view_)
     RemoveChildView(extra_view_);
 
   UpdateDialogButton(&ok_button_, ui::DIALOG_BUTTON_OK);
   UpdateDialogButton(&cancel_button_, ui::DIALOG_BUTTON_CANCEL);
 
   if (extra_view_)
     return;
 
   extra_view_ = GetDialogDelegate()->CreateExtraView();
   if (extra_view_)
     extra_view_->SetGroup(kButtonGroup);
 }
 
 }  // namespace views