DesktopAura: Track windows "owned" via the DesktopWindowTreeHost

ui/views/widget/native_widget_aura.cc

Index: ui/views/widget/native_widget_aura.cc
diff --git a/ui/views/widget/native_widget_aura.cc b/ui/views/widget/native_widget_aura.cc
index e493ea209646face0149d26149764d78cddc4bbf..bad1b4288c8f3b4703f9e228d7466c2479b23332 100644
--- a/ui/views/widget/native_widget_aura.cc
+++ b/ui/views/widget/native_widget_aura.cc
@@ -788,6 +788,11 @@ std::string NativeWidgetAura::GetName() const {
   return window_ ? window_->GetName() : std::string();
 }
 
+Widget::Widgets NativeWidgetAura::GetAllOwnedTopLevelWidgets() const {
+  // Only the root window is top level.
+  return Widget::Widgets();
+}
+
 ////////////////////////////////////////////////////////////////////////////////
 // NativeWidgetAura, aura::WindowDelegate implementation:
 
@@ -1141,8 +1146,8 @@ void NativeWidgetPrivate::GetAllChildWidgets(gfx::NativeView native_view,
                                              Widget::Widgets* children) {
   {
     // Code expects widget for |native_view| to be added to |children|.
-    NativeWidgetPrivate* native_widget = static_cast<NativeWidgetPrivate*>(
-        GetNativeWidgetForNativeView(native_view));
+    NativeWidgetPrivate* native_widget =
+        GetNativeWidgetForNativeView(native_view);
     if (native_widget && native_widget->GetWidget())
       children->insert(native_widget->GetWidget());
   }
@@ -1156,19 +1161,34 @@ void NativeWidgetPrivate::GetAllChildWidgets(gfx::NativeView native_view,
 
 // static
 void NativeWidgetPrivate::GetAllOwnedWidgets(gfx::NativeView native_view,
-                                             Widget::Widgets* owned) {
+                                             Widget::Widgets* owned,
+                                             bool include_toplevel) {
   // Add all owned widgets.
   for (aura::Window* transient_child : wm::GetTransientChildren(native_view)) {
-    NativeWidgetPrivate* native_widget = static_cast<NativeWidgetPrivate*>(
-        GetNativeWidgetForNativeView(transient_child));
+    NativeWidgetPrivate* native_widget =
+        GetNativeWidgetForNativeView(transient_child);
     if (native_widget && native_widget->GetWidget())
       owned->insert(native_widget->GetWidget());
-    GetAllOwnedWidgets(transient_child, owned);
+    GetAllOwnedWidgets(transient_child, owned, include_toplevel);
   }
 
   // Add all child windows.
   for (aura::Window* child : native_view->children())
     GetAllChildWidgets(child, owned);
+
+  if (!include_toplevel)
+    return;
+
+  // Discover potential top-level children owned via the WindowTreeHost.
+  NativeWidgetPrivate* native_widget =
+      GetNativeWidgetForNativeView(native_view);
+  if (!native_widget)
+    return;
+
+  for (Widget* child : native_widget->GetAllOwnedTopLevelWidgets()) {
+    owned->insert(child);
+    GetAllOwnedWidgets(child->GetNativeView(), owned, true);
+  }
 }
 
 // static