Enable setting primary display for 3+ displays

ash/display/window_tree_host_manager.cc

 // Copyright (c) 2012 The Chromium Authors. All rights reserved.
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file.
 
 #include "ash/display/window_tree_host_manager.h"
 
 #include <algorithm>
 #include <cmath>
 #include <map>
 #include <memory>
@@ skipping 145 matching lines @@
   ui::CursorController::GetInstance()->ClearCursorConfigForWindow(
       host->GetAcceleratedWidget());
 #endif
 }
 
 aura::Window* GetWindow(AshWindowTreeHost* ash_host) {
   CHECK(ash_host->AsWindowTreeHost());
   return ash_host->AsWindowTreeHost()->window();
 }
 
+void SwapRecursive(
+    const std::map<int64_t, display::DisplayPlacement*>& id_to_placement,
+    int64_t current_primary_id,
+    int64_t display_id) {
+  if (display_id == current_primary_id)
+    return;
+
+  DCHECK(id_to_placement.count(display_id));
+  display::DisplayPlacement* placement = id_to_placement.at(display_id);
+  DCHECK(placement);
+  SwapRecursive(id_to_placement, current_primary_id,
+                placement->parent_display_id);
+  placement->Swap();
+}
+
 }  // namespace
 
 // A utility class to store/restore focused/active window
 // when the display configuration has changed.
 class FocusActivationStore {
  public:
   FocusActivationStore()
       : activation_client_(nullptr),
         capture_client_(nullptr),
         focus_client_(nullptr),
@@ skipping 217 matching lines @@
   return controllers;
 }
 
 void WindowTreeHostManager::SetPrimaryDisplayId(int64_t id) {
   // TODO(oshima): Move primary display management to DisplayManager.
   DCHECK_NE(display::Display::kInvalidDisplayID, id);
   if (id == display::Display::kInvalidDisplayID || primary_display_id == id ||
       window_tree_hosts_.size() < 2) {
     return;
   }
-  // TODO(oshima): Implement swapping primary for 2> displays.
-  if (GetDisplayManager()->GetNumDisplays() > 2)
-    return;
 
   const display::Display& new_primary_display =
       GetDisplayManager()->GetDisplayForId(id);
   if (!new_primary_display.is_valid()) {
     LOG(ERROR) << "Invalid or non-existent display is requested:"
                << new_primary_display.ToString();
     return;
   }
 
   display::DisplayManager* display_manager = GetDisplayManager();
@@ skipping 25 matching lines @@
   GetRootWindowSettings(GetWindow(non_primary_host))->display_id =
       old_primary_display.id();
 
   const display::DisplayLayout& layout =
       GetDisplayManager()->GetCurrentDisplayLayout();
   // The requested primary id can be same as one in the stored layout
   // when the primary id is set after new displays are connected.
   // Only update the layout if it is requested to swap primary display.
   if (layout.primary_id != new_primary_display.id()) {
     std::unique_ptr<display::DisplayLayout> swapped_layout(layout.Copy());
-    swapped_layout->placement_list[0].Swap();
+
+    std::map<int64_t, display::DisplayPlacement*> id_to_placement;
+    for (auto& placement : swapped_layout->placement_list)
+      id_to_placement[placement.display_id] = &placement;
+    SwapRecursive(id_to_placement, primary_display_id,
+                  new_primary_display.id());
+
+    std::sort(swapped_layout->placement_list.begin(),
+              swapped_layout->placement_list.end(),
+              [](const display::DisplayPlacement& d1,
+                 const display::DisplayPlacement& d2) {
+                return d1.display_id < d2.display_id;
+              });
+
     swapped_layout->primary_id = new_primary_display.id();
     display::DisplayIdList list = display_manager->GetCurrentDisplayIdList();
     GetDisplayManager()->layout_store()->RegisterLayoutForDisplayIdList(
         list, std::move(swapped_layout));
   }
 
   primary_display_id = new_primary_display.id();
 
   UpdateWorkAreaOfDisplayNearestWindow(GetWindow(primary_host),
                                        old_primary_display.GetWorkAreaInsets());
@@ skipping 408 matching lines @@
   SetDisplayPropertiesOnHost(ash_host, display);
 
 #if defined(OS_CHROMEOS)
   if (switches::ConstrainPointerToRoot())
     ash_host->ConfineCursorToRootWindow();
 #endif
   return ash_host;
 }
 
 }  // namespace ash