Make OnWmDisplayRemoved() take a WindowTreeHost

ui/aura/mus/window_tree_client.cc

Index: ui/aura/mus/window_tree_client.cc
diff --git a/ui/aura/mus/window_tree_client.cc b/ui/aura/mus/window_tree_client.cc
index eff5c9d1abd81ee49dc4eb1ad21a38655a68cea8..8b0e9f8fc8d549960ef0a49655b2abce31eac246 100644
--- a/ui/aura/mus/window_tree_client.cc
+++ b/ui/aura/mus/window_tree_client.cc
@@ -1301,15 +1301,15 @@ void WindowTreeClient::WmNewDisplayAdded(const display::Display& display,
 
 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) {