Enable setting primary display for 3+ displays

ash/display/window_tree_host_manager.cc

Index: ash/display/window_tree_host_manager.cc
diff --git a/ash/display/window_tree_host_manager.cc b/ash/display/window_tree_host_manager.cc
index d23d5c80493cd26610c8f3db8baa312d5ce38ee1..61d9fed858016dfb25c97ed48b89dc0900656991 100644
--- a/ash/display/window_tree_host_manager.cc
+++ b/ash/display/window_tree_host_manager.cc
@@ -163,6 +163,21 @@ aura::Window* GetWindow(AshWindowTreeHost* ash_host) {
   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
@@ -400,9 +415,6 @@ void WindowTreeHostManager::SetPrimaryDisplayId(int64_t 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);
@@ -448,7 +460,20 @@ void WindowTreeHostManager::SetPrimaryDisplayId(int64_t id) {
   // 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(