Change OnWindowInputEvent to use display_id to find the host and update event root_location in WS.

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 3f69a8f49f89f93adec2bc890b2c226047c729c9..305903463045f488e3698ec0016082af67217036 100644
--- a/ui/aura/mus/window_tree_client.cc
+++ b/ui/aura/mus/window_tree_client.cc
@@ -1136,18 +1136,6 @@ void WindowTreeClient::OnWindowInputEvent(uint32_t event_id,
 
   WindowMus* window = GetWindowByServerId(window_id);  // May be null.
 
-  if (event->IsKeyEvent()) {
-    DCHECK(!matches_pointer_watcher);  // PointerWatcher isn't for key events.
-    if (!window || !window->GetWindow()->GetHost()) {
-      tree_->OnWindowInputEventAck(event_id, ui::mojom::EventResult::UNHANDLED);
-      return;
-    }
-    InputMethodMus* input_method = GetWindowTreeHostMus(window)->input_method();
-    input_method->DispatchKeyEvent(event->AsKeyEvent(),
-                                   CreateEventResultCallback(event_id));
-    return;
-  }
-
   if (matches_pointer_watcher && has_pointer_watcher_) {
     DCHECK(event->IsPointerEvent());
     std::unique_ptr<ui::Event> event_in_dip(ui::Event::Clone(*event));
@@ -1156,15 +1144,30 @@ void WindowTreeClient::OnWindowInputEvent(uint32_t event_id,
                                       window ? window->GetWindow() : nullptr);
   }
 
-  // TODO: use |display_id| to find host and send there.
+  WindowTreeHostMus* host = nullptr;
   if (!window || !window->GetWindow()->GetHost()) {
-    tree_->OnWindowInputEventAck(event_id, ui::mojom::EventResult::UNHANDLED);
-    return;
+    // If |window| has already been deleted by the time OnWindowInputEvent is
+    // called, we fall back to use the |display_id| to find the window_tree_host
+    // for event dispatching.
+    host = GetWindowTreeHostMusWithDisplayId(display_id);
+    if (!host) {
+      tree_->OnWindowInputEventAck(event_id, ui::mojom::EventResult::UNHANDLED);
+      return;
+    }
+    window = WindowMus::Get(host->window());
+  } else {
+    host = GetWindowTreeHostMus(window);

[sky, 2017/02/14 23:31:14]

I don't think you're guaranteed the window has a host. For example, if between
the dispatch and this function being called the window was removed then the host
will be null. Can we always lookup the host?

[riajiang, 2017/02/15 19:35:05]

Oh makes sense. I changed it to always lookup the host with display_id and use
the root window if window was removed before this?

   }
-
-  WindowTreeHostMus* host = GetWindowTreeHostMus(window);
   DCHECK(host);
 
+  if (event->IsKeyEvent()) {
+    DCHECK(!matches_pointer_watcher);  // PointerWatcher isn't for key events.
+    InputMethodMus* input_method = host->input_method();
+    input_method->DispatchKeyEvent(event->AsKeyEvent(),
+                                   CreateEventResultCallback(event_id));
+    return;
+  }
+
   // The location of the event is relative to |window|. As the event is handed
   // to WindowTreeHost we need it to be relative to WindowTreeHost.
   if (event->IsLocatedEvent()) {
@@ -1347,14 +1350,10 @@ void WindowTreeClient::WmNewDisplayAdded(const display::Display& display,
 
 void WindowTreeClient::WmDisplayRemoved(int64_t display_id) {
   DCHECK(window_manager_delegate_);
-  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;
-    }
+  WindowTreeHostMus* host = GetWindowTreeHostMusWithDisplayId(display_id);
+  if (host) {
+    window_manager_delegate_->OnWmDisplayRemoved(host);
+    return;
   }
 }
 
@@ -1851,4 +1850,16 @@ uint32_t WindowTreeClient::CreateChangeIdForFocus(WindowMus* window) {
       this, focus_synchronizer_.get(), window));
 }
 
+WindowTreeHostMus* WindowTreeClient::GetWindowTreeHostMusWithDisplayId(
+    int64_t display_id) {
+  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)
+      return window_tree_host;
+  }
+  return nullptr;
+}
+
 }  // namespace aura