Fix bug where removed (but not deleted) windows are not reflected in the tree properly

ash/common/devtools/ash_devtools_dom_agent.cc

Index: ash/common/devtools/ash_devtools_dom_agent.cc
diff --git a/ash/common/devtools/ash_devtools_dom_agent.cc b/ash/common/devtools/ash_devtools_dom_agent.cc
index fca24efa4799afe22978610c1f5cd75c60405b37..f3f3c85322409cdc3427728a1e89feb8c0b004bd 100644
--- a/ash/common/devtools/ash_devtools_dom_agent.cc
+++ b/ash/common/devtools/ash_devtools_dom_agent.cc
@@ -101,28 +101,34 @@ ui::devtools::protocol::Response AshDevToolsDOMAgent::getDocument(
   return ui::devtools::protocol::Response::OK();
 }
 
-// Need to remove node in OnWindowDestroying because the window parent reference
-// is gone in OnWindowDestroyed
-void AshDevToolsDOMAgent::OnWindowDestroying(WmWindow* window) {
-  RemoveWindowNode(window, window->GetParent());
-}
-
+// Handles removing windows.
+void AshDevToolsDOMAgent::OnWindowTreeChanging(WmWindow* window,
+                                               const TreeChangeParams& params) {
+  // Only trigger this when window == params.old_parent.
+  // Only removals are handled here. Removing a node can occur as a result of
+  // reorganizing a window or just destroying it. OnWindowTreeChanged
+  // is only called if there is a new_parent. The only case this method isn't
+  // called is when adding a node because old_parent is then null.
+  // Finally, We only trigger this  0 or 1 times as an old_parent will
+  // either exist and only call this callback once, or not at all.
+  if (window == params.old_parent)
+    RemoveWindowNode(params.target);
+}
+
+// Handles adding windows.
 void AshDevToolsDOMAgent::OnWindowTreeChanged(WmWindow* window,
                                               const TreeChangeParams& params) {
-  // Only trigger this when window == root window.
-  // Removals are handled on OnWindowDestroying.
-  if (window != window->GetRootWindow() || !params.new_parent)
-    return;
-  // If there is an old_parent + new_parent, then this window is being moved
-  // which requires a remove followed by an add. If only new_parent
-  // exists, then a new window is being created so only add is called.
-  if (params.old_parent)
-    RemoveWindowNode(params.target, params.old_parent);
-  AddWindowNode(params.target);
+  // Only trigger this when window == params.new_parent.
+  // If there is an old_parent + new_parent, then this window's node was
+  // removed in OnWindowTreeChanging and will now be added to the new_parent.
+  // If there is only a new_parent, OnWindowTreeChanging is never called and
+  // the window is only added here.
+  if (window == params.new_parent)
+    AddWindowNode(params.target);
 }
 
 void AshDevToolsDOMAgent::OnWindowStackingChanged(WmWindow* window) {
-  RemoveWindowNode(window, window->GetParent());
+  RemoveWindowNode(window);
   AddWindowNode(window);
 }
 
@@ -154,6 +160,7 @@ AshDevToolsDOMAgent::BuildInitialTree() {
 }
 
 void AshDevToolsDOMAgent::AddWindowNode(WmWindow* window) {
+  DCHECK(window_to_node_id_map_.count(window->GetParent()));
   WmWindow* prev_sibling = FindPreviousSibling(window);
   frontend()->childNodeInserted(
       window_to_node_id_map_[window->GetParent()],
@@ -161,15 +168,17 @@ void AshDevToolsDOMAgent::AddWindowNode(WmWindow* window) {
       BuildTreeForWindow(window));
 }
 
-void AshDevToolsDOMAgent::RemoveWindowNode(WmWindow* window,
-                                           WmWindow* old_parent) {
-  window->RemoveObserver(this);
+void AshDevToolsDOMAgent::RemoveWindowNode(WmWindow* window) {
+  WmWindow* parent = window->GetParent();
+  DCHECK(parent);
+  DCHECK(window_to_node_id_map_.count(parent));
   WindowToNodeIdMap::iterator it = window_to_node_id_map_.find(window);
   DCHECK(it != window_to_node_id_map_.end());
 
   int node_id = it->second;
-  int parent_id = old_parent ? window_to_node_id_map_[old_parent] : 0;
+  int parent_id = window_to_node_id_map_[parent];
 
+  window->RemoveObserver(this);
   window_to_node_id_map_.erase(it);
   frontend()->childNodeRemoved(parent_id, node_id);
 }