Makes WindowTree::SetFocus(null) work

components/mus/ws/window_server.cc

 // Copyright 2014 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 "components/mus/ws/window_server.h"
 
 #include "base/logging.h"
 #include "base/stl_util.h"
 #include "components/mus/ws/display.h"
 #include "components/mus/ws/display_binding.h"
@@ skipping 222 matching lines @@
 
   // We've been supplied a WindowManagerFactory and no displays have been
   // created yet. Treat this as a signal to create a Display.
   // TODO(sky): we need a better way to determine this, most likely a switch.
   delegate_->CreateDefaultDisplays();
 }
 
 void WindowServer::SetFocusedWindow(ServerWindow* window) {
   // TODO(sky): this should fail if there is modal dialog active and |window|
   // is outside that.
-  Display* focused_display = nullptr;
-  for (Display* display : display_manager_->displays()) {
-    if (display->GetFocusedWindow()) {
-      focused_display = display;
-      break;
-    }
+  ServerWindow* currently_focused = GetFocusedWindow();
+  Display* focused_display =
+      currently_focused
+          ? display_manager_->GetDisplayContaining(currently_focused)
+          : nullptr;
+  if (!window) {
+    if (focused_display)
+      focused_display->SetFocusedWindow(nullptr);
+    return;
   }
   Display* display = display_manager_->GetDisplayContaining(window);
   DCHECK(display);  // It's assumed callers do validation before calling this.
   display->SetFocusedWindow(window);
   // If the focus actually changed, and focus was in another display, then we
   // need to notify the previously focused display so that it cleans up state
   // and notifies appropriately.
   if (window && display->GetFocusedWindow() && display != focused_display &&
       focused_display) {
     focused_display->SetFocusedWindow(nullptr);
   }
 }
 
+ServerWindow* WindowServer::GetFocusedWindow() {
+  for (Display* display : display_manager_->displays()) {
+    ServerWindow* focused_window = display->GetFocusedWindow();
+    if (focused_window)
+      return focused_window;
+  }
+  return nullptr;
+}
+
 uint32_t WindowServer::GenerateWindowManagerChangeId(
     WindowTree* source,
     uint32_t client_change_id) {
   const uint32_t wm_change_id = next_wm_change_id_++;
   in_flight_wm_change_map_[wm_change_id] = {source->id(), client_change_id};
   return wm_change_id;
 }
 
 void WindowServer::WindowManagerChangeCompleted(
     uint32_t window_manager_change_id,
@@ skipping 351 matching lines @@
 void WindowServer::OnFirstDisplayReady() {
   delegate_->OnFirstDisplayReady();
 }
 
 void WindowServer::OnNoMoreDisplays() {
   delegate_->OnNoMoreDisplays();
 }
 
 }  // namespace ws
 }  // namespace mus