Make OnWmDisplayRemoved() take a WindowTreeHost

ui/aura/mus/window_tree_client.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 "ui/aura/mus/window_tree_client.h"
 
 #include <stddef.h>
 
 #include <string>
 #include <utility>
@@ skipping 1283 matching lines @@
 }
 
 void WindowTreeClient::WmNewDisplayAdded(const display::Display& display,
                                          ui::mojom::WindowDataPtr root_data,
                                          bool parent_drawn) {
   WmNewDisplayAddedImpl(display, std::move(root_data), parent_drawn);
 }
 
 void WindowTreeClient::WmDisplayRemoved(int64_t display_id) {
   DCHECK(window_manager_delegate_);
-  // TODO: route to WindowTreeHost.
-  /*
-  for (Window* root : roots_) {
-    if (root->display_id() == display_id) {
-      window_manager_delegate_->OnWmDisplayRemoved(root);
+  for (WindowMus* root : roots_) {
+    DCHECK(root->GetWindow()->GetHost());
+    WindowTreeHostMus* window_tree_host =
+        static_cast<WindowTreeHostMus*>(root->GetWindow()->GetHost());
+    if (window_tree_host->display_id() == display_id) {
+      window_manager_delegate_->OnWmDisplayRemoved(window_tree_host);
       return;
     }
   }
-  */
 }
 
 void WindowTreeClient::WmDisplayModified(const display::Display& display) {
   DCHECK(window_manager_delegate_);
   // TODO(sky): this should likely route to WindowTreeHost.
   window_manager_delegate_->OnWmDisplayModified(display);
 }
 
 // TODO(riajiang): Convert between pixel and DIP for window bounds properly.
 // (http://crbug.com/646942)
@@ skipping 260 matching lines @@
   return ScheduleInFlightChange(base::MakeUnique<InFlightCaptureChange>(
       this, capture_synchronizer_.get(), window));
 }
 
 uint32_t WindowTreeClient::CreateChangeIdForFocus(WindowMus* window) {
   return ScheduleInFlightChange(base::MakeUnique<InFlightFocusChange>(
       this, focus_synchronizer_.get(), window));
 }
 
 }  // namespace aura