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 1b57f14f5b129b8f4b0614063a38f6bc6f5b712f..3bce454fd2bbf51f1c96afc10b5f9caa4777021e 100644
--- a/ui/aura/mus/window_tree_client.cc
+++ b/ui/aura/mus/window_tree_client.cc
@@ -160,12 +160,24 @@ void ConvertEventLocationToDip(int64_t display_id, ui::LocatedEvent* event) {
   event->set_root_location(root_location);
 }
 
-// Set the |target| to be the target window of this |event| and send it to
-// the EventProcessor.
-void DispatchEventToTarget(ui::Event* event, WindowMus* target) {
-  ui::Event::DispatcherApi dispatch_helper(event);
-  dispatch_helper.set_target(target->GetWindow());
-  GetWindowTreeHostMus(target)->SendEventToProcessor(event);
+// Set the |target| to be the target window of this |event| if it is not null,
+// otherwise update event location and leave the event target empty, and send
+// the |event| to the EventProcessor.
+void DispatchEventToTarget(ui::Event* event,
+                           WindowTreeHostMus* host,
+                           WindowMus* target) {
+  if (target) {
+    ui::Event::DispatcherApi dispatch_helper(event);
+    dispatch_helper.set_target(target->GetWindow());
+  } else if (event->IsLocatedEvent()) {
+    // The target window specified by the window server no longer exists. So it
+    // is necessary to find the target using the local event targeters. For
+    // that to work correctly, the event-location needs to be in the root
+    // window's coordinate space.
+    event->AsLocatedEvent()->set_location(
+        event->AsLocatedEvent()->root_location());
+  }
+  host->SendEventToProcessor(event);
 }
 
 }  // namespace
@@ -1156,19 +1168,9 @@ void WindowTreeClient::OnWindowInputEvent(uint32_t event_id,
                                           bool matches_pointer_watcher) {
   DCHECK(event);
 
-  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;
-  }
+  // |window| may be null if it has already been deleted as OnWindowInputEvent
+  // is being called asynchronously.
+  WindowMus* window = GetWindowByServerId(window_id);
 
   if (matches_pointer_watcher && has_pointer_watcher_) {
     DCHECK(event->IsPointerEvent());
@@ -1178,28 +1180,41 @@ void WindowTreeClient::OnWindowInputEvent(uint32_t event_id,
                                       window ? window->GetWindow() : nullptr);
   }
 
-  // TODO: use |display_id| to find host and send there.
-  if (!window || !window->GetWindow()->GetHost()) {
+  // Since this |window| might have already been deleted by the time
+  // OnWindowInputEvent is called, we use the |display_id| to find the host
+  // for event dispatching.
+  WindowTreeHostMus* host = window
+                                ? GetWindowTreeHostMus(window)
+                                : GetWindowTreeHostMusWithDisplayId(display_id);
+  if (!host) {
     tree_->OnWindowInputEventAck(event_id, ui::mojom::EventResult::UNHANDLED);
     return;
   }
 
+  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;
+  }
+
   EventAckHandler ack_handler(CreateEventResultCallback(event_id));
   // TODO(moshayedi): crbug.com/617222. No need to convert to ui::MouseEvent or
   // ui::TouchEvent once we have proper support for pointer events.
   if (event->IsMousePointerEvent()) {
     if (event->type() == ui::ET_POINTER_WHEEL_CHANGED) {
       ui::MouseWheelEvent mapped_event(*event->AsPointerEvent());
-      DispatchEventToTarget(&mapped_event, window);
+      DispatchEventToTarget(&mapped_event, host, window);
     } else {
       ui::MouseEvent mapped_event(*event->AsPointerEvent());
-      DispatchEventToTarget(&mapped_event, window);
+      DispatchEventToTarget(&mapped_event, host, window);
     }
   } else if (event->IsTouchPointerEvent()) {
     ui::TouchEvent mapped_event(*event->AsPointerEvent());
-    DispatchEventToTarget(&mapped_event, window);
+    DispatchEventToTarget(&mapped_event, host, window);
   } else {
-    DispatchEventToTarget(event.get(), window);
+    DispatchEventToTarget(event.get(), host, window);
   }
   ack_handler.set_handled(event->handled());
 }
@@ -1356,15 +1371,9 @@ 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);
 }
 
 void WindowTreeClient::WmDisplayModified(const display::Display& display) {
@@ -1855,4 +1864,18 @@ uint32_t WindowTreeClient::CreateChangeIdForFocus(WindowMus* window) {
       this, focus_synchronizer_.get(), window));
 }
 
+WindowTreeHostMus* WindowTreeClient::GetWindowTreeHostMusWithDisplayId(
+    int64_t display_id) {
+  // TODO(riajiang): Figure out how to get the correct WindowTreeHost we need
+  // in external window mode.
+  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