bounds_tracker: WindowBoundsTracker maintains the Observing window list

- Let WindowBoundsTracker track the window activation changes, and
  observe the window that would have its own WindowState.
- Track the window's root window changes inside WindowBoundsTracker,
  remap or restore the window on this change.
  - Only applying the remap/restore on root window changes because of
    moving the active window through the shortcut `alt+search+m` or
    display removal.
- Remove the callsites of RemapOrRestore because of window's root
  window changes. As now it is handled by WindowBoundsTracker itself
  instead.

Bug: b/313937160, b/314006763
Change-Id: I2a6f1b3eac7ef33c1797468f0ee2ef5d5cb92240
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/5073114
Reviewed-by: Ahmed Fakhry <afakhry@chromium.org>
Commit-Queue: Min Chen <minch@chromium.org>
Cr-Commit-Position: refs/heads/main@{#1233045}
diff --git a/ash/display/display_move_window_util.cc b/ash/display/display_move_window_util.cc
index c14bc6a..e88085a 100644
--- a/ash/display/display_move_window_util.cc
+++ b/ash/display/display_move_window_util.cc
@@ -10,6 +10,7 @@
 
 #include "ash/accessibility/accessibility_controller_impl.h"
 #include "ash/shell.h"
+#include "ash/wm/bounds_tracker/window_bounds_tracker.h"
 #include "ash/wm/mru_window_tracker.h"
 #include "ash/wm/window_util.h"
 #include "base/containers/contains.h"
@@ -76,6 +77,10 @@
                               origin_display_id, display::CompareDisplayIds);
   int64_t target_display_id =
       itr == display_id_list.end() ? display_id_list[0] : *itr;
+
+  if (auto* window_bounds_tracker = Shell::Get()->window_bounds_tracker()) {
+    window_bounds_tracker->set_moving_window_between_displays(window);
+  }
   window_util::MoveWindowToDisplay(window, target_display_id);
   Shell::Get()->accessibility_controller()->TriggerAccessibilityAlert(
       AccessibilityAlert::WINDOW_MOVED_TO_ANOTHER_DISPLAY);
diff --git a/ash/display/window_tree_host_manager.cc b/ash/display/window_tree_host_manager.cc
index a5f3419..b9c83e0 100644
--- a/ash/display/window_tree_host_manager.cc
+++ b/ash/display/window_tree_host_manager.cc
@@ -674,10 +674,9 @@
   Shell::Get()->OnRootWindowWillShutdown(root_being_deleted);
   aura::Window* primary_root_after_host_deletion =
       GetRootWindowForDisplayId(GetPrimaryDisplayId());
-  controller->MoveWindowsTo(primary_root_after_host_deletion);
   // Delete most of root window related objects, but don't delete
   // root window itself yet because the stack may be using it.
-  controller->Shutdown();
+  controller->Shutdown(primary_root_after_host_deletion);
   if (primary_tree_host_for_replace_ == host_to_delete)
     primary_tree_host_for_replace_ = nullptr;
   DCHECK_EQ(primary_root_after_host_deletion, Shell::GetPrimaryRootWindow());
diff --git a/ash/root_window_controller.cc b/ash/root_window_controller.cc
index 1294ad0..20a8d18d 100644
--- a/ash/root_window_controller.cc
+++ b/ash/root_window_controller.cc
@@ -202,24 +202,20 @@
   gfx::Rect restore_bounds;
   const bool has_restore_bounds = state && state->HasRestoreBounds();
 
+  auto* window_bounds_tracker = Shell::Get()->window_bounds_tracker();
+  // `WindowBoundsTracker` will handle the window's bounds on root window
+  // changes if the feature `kWindowBoundsTracker` is enabled.
   const bool update_bounds =
-      state && (state->IsNormalStateType() || state->IsMinimized());
+      state && (state->IsNormalStateType() || state->IsMinimized()) &&
+      !window_bounds_tracker;
   gfx::Rect work_area_in_new_parent =
       screen_util::GetDisplayWorkAreaBoundsInParent(new_parent);
 
   gfx::Rect local_bounds;
-  auto* window_bounds_tracker = Shell::Get()->window_bounds_tracker();
   if (update_bounds) {
-    if (window_bounds_tracker) {
-      local_bounds = window_bounds_tracker->RemapOrRestore(
-          window, display::Screen::GetScreen()
-                      ->GetDisplayNearestWindow(new_parent)
-                      .id());
-    } else {
-      local_bounds = state->window()->bounds();
-      MoveOriginRelativeToSize(src_size, dst_size, &local_bounds);
-      local_bounds.AdjustToFit(work_area_in_new_parent);
-    }
+    local_bounds = state->window()->bounds();
+    MoveOriginRelativeToSize(src_size, dst_size, &local_bounds);
+    local_bounds.AdjustToFit(work_area_in_new_parent);
   }
 
   if (has_restore_bounds) {
@@ -232,10 +228,6 @@
     window_bounds_tracker->AddWindowDisplayIdOnDisplayRemoval(window);
   }
 
-  // TODO: Let `WindowBoundsTracker` maintain an observing windows list, then
-  // the window's bounds can be updated correctly inside
-  // `OnWindowRemovingRootWindow` and `OnWindowAddedToRootWindow` because of
-  // reparenting.
   new_parent->AddChild(window);
 
   // Snapped windows have bounds handled by the layout manager in AddChild().
@@ -522,7 +514,7 @@
     RootWindowController::root_window_controllers_ = nullptr;
 
 RootWindowController::~RootWindowController() {
-  Shutdown();
+  Shutdown(/*destination_root=*/nullptr);
   DCHECK(!wallpaper_widget_controller_.get());
   work_area_insets_.reset();
   ash_host_.reset();
@@ -699,9 +691,13 @@
   return screen_rotation_animator_.get();
 }
 
-void RootWindowController::Shutdown() {
+void RootWindowController::Shutdown(aura::Window* destination_root) {
   is_shutting_down_ = true;
 
+  if (destination_root) {
+    MoveWindowsTo(destination_root);
+  }
+
   // Destroy the `screen_rotation_animator_` now to avoid any potential crashes
   // if there's any ongoing animation. See http://b/293667233.
   screen_rotation_animator_.reset();
@@ -815,23 +811,6 @@
   ::wm::SetTooltipClient(GetRootWindow(), nullptr);
 }
 
-void RootWindowController::MoveWindowsTo(aura::Window* dst) {
-  // Suspend unnecessary updates of the shelf visibility indefinitely since it
-  // is going away.
-  if (GetShelfLayoutManager()) {
-    GetShelfLayoutManager()->SuspendVisibilityUpdateForShutdown();
-  }
-
-  // Clear the workspace controller to avoid a lot of unnecessary operations
-  // when window are removed.
-  // TODO(afakhry): Should we also clear the WorkspaceLayoutManagers of the pip,
-  // always-on-top, and other containers?
-  aura::Window* root = GetRootWindow();
-  ClearWorkspaceControllers(root);
-
-  ReparentAllWindows(root, dst);
-}
-
 void RootWindowController::InitTouchHuds() {
   // Enable touch debugging features when each display is initialized.
   base::CommandLine* command_line = base::CommandLine::ForCurrentProcess();
@@ -1067,6 +1046,23 @@
   capture_client_ = std::make_unique<::wm::ScopedCaptureClient>(root_window);
 }
 
+void RootWindowController::MoveWindowsTo(aura::Window* dst) {
+  // Suspend unnecessary updates of the shelf visibility indefinitely since it
+  // is going away.
+  if (GetShelfLayoutManager()) {
+    GetShelfLayoutManager()->SuspendVisibilityUpdateForShutdown();
+  }
+
+  // Clear the workspace controller to avoid a lot of unnecessary operations
+  // when window are removed.
+  // TODO(afakhry): Should we also clear the WorkspaceLayoutManagers of the pip,
+  // always-on-top, and other containers?
+  aura::Window* root = GetRootWindow();
+  ClearWorkspaceControllers(root);
+
+  ReparentAllWindows(root, dst);
+}
+
 void RootWindowController::Init(RootWindowType root_window_type) {
   aura::Window* root_window = GetRootWindow();
   // Create |split_view_controller_| for every display.
diff --git a/ash/root_window_controller.h b/ash/root_window_controller.h
index 6349f4c..71135488 100644
--- a/ash/root_window_controller.h
+++ b/ash/root_window_controller.h
@@ -141,6 +141,8 @@
     return root_window_layout_manager_;
   }
 
+  bool is_shutting_down() const { return is_shutting_down_; }
+
   // Returns parameters of the work area associated with this root window.
   WorkAreaInsets* work_area_insets() { return work_area_insets_.get(); }
 
@@ -206,19 +208,14 @@
   // Deletes associated objects and clears the state, but doesn't delete
   // the root window yet. This is used to delete a secondary displays'
   // root window safely when the display disconnect signal is received,
-  // which may come while we're in the nested run loop.
-  void Shutdown();
+  // which may come while we're in the nested run loop. Child windows of the
+  // root window of this controller will be moved to `destination_root` if
+  // provided.
+  void Shutdown(aura::Window* destination_root);
 
   // Deletes all child windows and performs necessary cleanup.
   void CloseChildWindows();
 
-  // Moves child windows to |dest|.
-  // TODO(afakhry): Consider renaming this function to avoid misuse. It is only
-  // called by WindowTreeHostManager::DeleteHost(), and has destructive side
-  // effects like deleting the workspace controllers, so it shouldn't be called
-  // for something else.
-  void MoveWindowsTo(aura::Window* dest);
-
   // Initialize touch HUDs if necessary.
   void InitTouchHuds();
 
@@ -285,6 +282,9 @@
   // Takes ownership of |ash_host|.
   explicit RootWindowController(AshWindowTreeHost* ash_host);
 
+  // Moves child windows to `dest`.
+  void MoveWindowsTo(aura::Window* dest);
+
   // Initializes the RootWindowController based on |root_window_type|.
   void Init(RootWindowType root_window_type);
 
diff --git a/ash/shell.cc b/ash/shell.cc
index 9f470434..b6b39abf 100644
--- a/ash/shell.cc
+++ b/ash/shell.cc
@@ -1287,10 +1287,6 @@
 
   InitializeDisplayManager();
 
-  if (features::IsWindowBoundsTrackerEnabled()) {
-    window_bounds_tracker_ = std::make_unique<WindowBoundsTracker>();
-  }
-
   // RefreshFontParams depends on display prefs.
   display_manager_->RefreshFontParams();
 
@@ -1373,6 +1369,12 @@
   focus_controller_ = std::make_unique<::wm::FocusController>(focus_rules_);
   focus_controller_->AddObserver(this);
 
+  // `WindowBoundsTracker` depends on `FocusController`, as it needs to track
+  // the window's activation changes.
+  if (features::IsWindowBoundsTrackerEnabled()) {
+    window_bounds_tracker_ = std::make_unique<WindowBoundsTracker>();
+  }
+
   overview_controller_ = std::make_unique<OverviewController>();
 
   // `GameDashboardController` has dependencies on `OverviewController` and
diff --git a/ash/wm/bounds_tracker/window_bounds_tracker.cc b/ash/wm/bounds_tracker/window_bounds_tracker.cc
index a64f454..be1a72a 100644
--- a/ash/wm/bounds_tracker/window_bounds_tracker.cc
+++ b/ash/wm/bounds_tracker/window_bounds_tracker.cc
@@ -4,6 +4,7 @@
 
 #include "ash/wm/bounds_tracker/window_bounds_tracker.h"
 
+#include "ash/root_window_controller.h"
 #include "ash/shell.h"
 #include "ash/wm/window_state.h"
 #include "ash/wm/window_util.h"
@@ -11,6 +12,7 @@
 #include "ui/display/manager/display_manager.h"
 #include "ui/display/screen.h"
 #include "ui/gfx/geometry/vector2d_conversions.h"
+#include "ui/wm/public/activation_client.h"
 
 namespace ash {
 
@@ -140,64 +142,78 @@
 
 }  // namespace
 
-WindowBoundsTracker::WindowBoundsTracker() = default;
+WindowBoundsTracker::WindowBoundsTracker() {
+  Shell::Get()->activation_client()->AddObserver(this);
+}
 
 WindowBoundsTracker::~WindowBoundsTracker() {
-  ResetBoundsDatabase();
+  bounds_database_.clear();
+  window_observations_.RemoveAllObservations();
+  Shell::Get()->activation_client()->RemoveObserver(this);
 }
 
 void WindowBoundsTracker::OnWindowDestroying(aura::Window* window) {
   RemoveWindowFromBoundsDatabase(window);
 }
 
-gfx::Rect WindowBoundsTracker::RemapOrRestore(aura::Window* window,
-                                              int64_t target_display_id) {
-  WindowState* window_state = WindowState::Get(window);
-  const gfx::Rect bounds_in_parent = window->bounds();
-  // TODO: Taking care of the windows in other window states.
-  if (!window_state->IsNormalStateType()) {
-    return bounds_in_parent;
+void WindowBoundsTracker::OnWindowAddedToRootWindow(aura::Window* window) {
+  // Set `window` to the remapping bounds calculated and stored to
+  // `bounds_database_` inside `OnWindowRemovingFromRootWindow`. If there is no
+  // remapping or restoring bounds can be found for `window`, which means it has
+  // never been moved to another display without user-assigned bounds.
+  const auto iter = bounds_database_.find(window);
+  if (iter == bounds_database_.end()) {
+    return;
   }
-
-  auto* screen = display::Screen::GetScreen();
-  const display::Display source_display =
-      screen->GetDisplayNearestWindow(window);
-  const int64_t source_display_id = source_display.id();
-  gfx::Rect source_work_area = source_display.GetLocalWorkArea();
-
-  const auto& window_bounds_map = UpdateBoundsDatabaseOfWindow(
-      window, source_display_id, source_display.rotation(), source_work_area);
+  const auto& window_bounds_map = iter->second;
   CHECK(!window_bounds_map.empty());
-
-  display::Display target_display;
-  screen->GetDisplayWithDisplayId(target_display_id, &target_display);
-  CHECK(target_display.is_valid());
-  const gfx::Rect target_work_area = target_display.GetLocalWorkArea();
+  display::Display target_display =
+      display::Screen::GetScreen()->GetDisplayNearestWindow(window);
   const WindowDisplayInfo target_window_display_info(
-      target_display_id, target_display.rotation(), target_work_area);
-  const auto iter = window_bounds_map.find(target_window_display_info);
+      target_display.id(), target_display.rotation(),
+      target_display.GetLocalWorkArea());
+  const auto bounds_iter = window_bounds_map.find(target_window_display_info);
+  CHECK(bounds_iter != window_bounds_map.end());
 
-  if (iter != window_bounds_map.end()) {
-    // Restores window to its bounds stored with `target_window_display_info`.
-    return iter->second;
+  window->SetBounds(bounds_iter->second);
+}
+
+void WindowBoundsTracker::OnWindowRemovingFromRootWindow(
+    aura::Window* window,
+    aura::Window* new_root) {
+  // Check whether we should remap or restore `window` on its root window
+  // changes. Only needed if 1) the window was moved between displays through
+  // the shortcut `kMoveActiveWindowBetweenDisplays` 2) removing the window's
+  // host display and the window will be moved to the current primary display.
+  // As in these two scenarios, the window is moving to another display without
+  // user assigned bounds.
+  const bool is_moving_window_between_displays =
+      window == moving_window_between_displays_;
+  const bool should_remap_or_restore =
+      is_moving_window_between_displays ||
+      RootWindowController::ForWindow(window->GetRootWindow())
+          ->is_shutting_down();
+  if (!should_remap_or_restore) {
+    return;
   }
 
-  // Otherwise, calculating the remapping bounds.
+  RemapOrRestore(
+      window,
+      display::Screen::GetScreen()->GetDisplayNearestWindow(new_root).id());
+  // Reset `moving_window_between_displays_` after finishing the remap or
+  // restore on it.
+  if (is_moving_window_between_displays) {
+    moving_window_between_displays_ = nullptr;
+  }
+}
 
-  // Step 1: Anchor point redesign, aka, keep the window's physical position on
-  // different screen orientations.
-  gfx::Rect remapped_bounds = AdjustBoundsForRotation(
-      bounds_in_parent, source_display, target_display, source_work_area);
-
-  // Step 2: Adjust on work area size changes. The relative position from the
-  // center of the window to the center of the work area should be the same.
-  AdjustBoundsForWorkArea(source_work_area, target_work_area, remapped_bounds);
-
-  // Step 3: Offscreen protection. The window should be fully visible inside the
-  // target display configuration.
-  remapped_bounds.AdjustToFit(target_work_area);
-
-  return remapped_bounds;
+void WindowBoundsTracker::OnWindowActivated(ActivationReason reason,
+                                            aura::Window* gained_active,
+                                            aura::Window* lost_active) {
+  if (WindowState::Get(gained_active) &&
+      !window_observations_.IsObservingSource(gained_active)) {
+    window_observations_.AddObservation(gained_active);
+  }
 }
 
 void WindowBoundsTracker::AddWindowDisplayIdOnDisplayRemoval(
@@ -212,9 +228,20 @@
   auto* display_manager = Shell::Get()->display_manager();
   auto iter = window_to_display_map_.begin();
   while (iter != window_to_display_map_.end()) {
-    const auto display_id = iter->second;
-    if (display_manager->IsDisplayIdValid(display_id)) {
-      window_util::MoveWindowToDisplay(iter->first, display_id);
+    const auto candidate_old_display_id = iter->second;
+    if (display_manager->IsDisplayIdValid(candidate_old_display_id)) {
+      auto* window = iter->first;
+      // TODO(b/314160218): Do not store the bounds if it is not user-assigned.
+      // Store the window's bounds in the source display before moving it to the
+      // target display.
+      const display::Display source_display =
+          display::Screen::GetScreen()->GetDisplayNearestWindow(window);
+      UpdateBoundsDatabaseOfWindow(
+          window,
+          WindowDisplayInfo(source_display.id(), source_display.rotation(),
+                            source_display.GetLocalWorkArea()),
+          window->bounds());
+      window_util::MoveWindowToDisplay(window, candidate_old_display_id);
       iter = window_to_display_map_.erase(iter);
     } else {
       ++iter;
@@ -242,32 +269,76 @@
 // -----------------------------------------------------------------------------
 // WindowBoundsTracker:
 
-void WindowBoundsTracker::ResetBoundsDatabase() {
-  for (const auto& [window, _] : bounds_database_) {
-    window->RemoveObserver(this);
+void WindowBoundsTracker::RemapOrRestore(aura::Window* window,
+                                         int64_t target_display_id) {
+  WindowState* window_state = WindowState::Get(window);
+  const gfx::Rect bounds_in_parent = window->bounds();
+  // TODO: Taking care of the windows in other window states.
+  if (!window_state->IsNormalStateType()) {
+    return;
   }
-  bounds_database_.clear();
+
+  auto* screen = display::Screen::GetScreen();
+  const display::Display source_display =
+      screen->GetDisplayNearestWindow(window);
+  const int64_t source_display_id = source_display.id();
+  gfx::Rect source_work_area = source_display.GetLocalWorkArea();
+
+  const auto& window_bounds_map = UpdateBoundsDatabaseOfWindow(
+      window,
+      WindowDisplayInfo(source_display_id, source_display.rotation(),
+                        source_work_area),
+      window->bounds());
+  CHECK(!window_bounds_map.empty());
+
+  display::Display target_display;
+  screen->GetDisplayWithDisplayId(target_display_id, &target_display);
+  CHECK(target_display.is_valid());
+  const gfx::Rect target_work_area = target_display.GetLocalWorkArea();
+  const WindowDisplayInfo target_window_display_info(
+      target_display_id, target_display.rotation(), target_work_area);
+  const auto iter = window_bounds_map.find(target_window_display_info);
+
+  if (iter != window_bounds_map.end()) {
+    return;
+  }
+
+  // Otherwise, calculating the remapping bounds.
+
+  // Step 1: Anchor point redesign, aka, keep the window's physical position on
+  // different screen orientations.
+  gfx::Rect remapped_bounds = AdjustBoundsForRotation(
+      bounds_in_parent, source_display, target_display, source_work_area);
+
+  // Step 2: Adjust on work area size changes. The relative position from the
+  // center of the window to the center of the work area should be the same.
+  AdjustBoundsForWorkArea(source_work_area, target_work_area, remapped_bounds);
+
+  // Step 3: Offscreen protection. The window should be fully visible inside the
+  // target display configuration.
+  remapped_bounds.AdjustToFit(target_work_area);
+
+  UpdateBoundsDatabaseOfWindow(window, target_window_display_info,
+                               remapped_bounds);
+  return;
 }
 
 void WindowBoundsTracker::RemoveWindowFromBoundsDatabase(aura::Window* window) {
   const auto count = bounds_database_.erase(window);
   CHECK(count);
-  window->RemoveObserver(this);
+  if (window_observations_.IsObservingSource(window)) {
+    window_observations_.RemoveObservation(window);
+  }
 }
 
 base::flat_map<WindowBoundsTracker::WindowDisplayInfo, gfx::Rect>&
 WindowBoundsTracker::UpdateBoundsDatabaseOfWindow(
     aura::Window* window,
-    int64_t display_id,
-    display::Display::Rotation rotation,
-    const gfx::Rect& work_area) {
+    const WindowDisplayInfo& window_display_info,
+    const gfx::Rect& bounds) {
   auto& window_bounds_map = bounds_database_[window];
-  if (window_bounds_map.empty()) {
-    window->AddObserver(this);
-  }
-  window_bounds_map.insert_or_assign(
-      WindowBoundsTracker::WindowDisplayInfo(display_id, rotation, work_area),
-      window->bounds());
+  CHECK(window_observations_.IsObservingSource(window));
+  window_bounds_map.insert_or_assign(window_display_info, bounds);
   return window_bounds_map;
 }
 
diff --git a/ash/wm/bounds_tracker/window_bounds_tracker.h b/ash/wm/bounds_tracker/window_bounds_tracker.h
index a14e03c..eebf2a4 100644
--- a/ash/wm/bounds_tracker/window_bounds_tracker.h
+++ b/ash/wm/bounds_tracker/window_bounds_tracker.h
@@ -8,8 +8,10 @@
 #include <unordered_map>
 
 #include "base/containers/flat_map.h"
+#include "base/scoped_multi_source_observation.h"
 #include "chromeos/ui/base/display_util.h"
 #include "ui/aura/window_observer.h"
+#include "ui/wm/public/activation_change_observer.h"
 
 namespace aura {
 class Window;
@@ -24,25 +26,28 @@
 // configuration. E.g., remapping the window if its host display being removed
 // and restoring it if reconnecting the display.
 // Note: `PersistentWindowController` will be disabled with this one enabled.
-class WindowBoundsTracker : public aura::WindowObserver {
+class WindowBoundsTracker : public aura::WindowObserver,
+                            public wm::ActivationChangeObserver {
  public:
   WindowBoundsTracker();
   WindowBoundsTracker(const WindowBoundsTracker&) = delete;
   WindowBoundsTracker& operator=(const WindowBoundsTracker&) = delete;
   ~WindowBoundsTracker() override;
 
+  void set_moving_window_between_displays(aura::Window* window) {
+    moving_window_between_displays_ = window;
+  }
+
   // aura::WindowObserver:
   void OnWindowDestroying(aura::Window* window) override;
+  void OnWindowAddedToRootWindow(aura::Window* window) override;
+  void OnWindowRemovingFromRootWindow(aura::Window* window,
+                                      aura::Window* new_root) override;
 
-  // Gets the window's restoring or remapping bounds in parent coordinates.
-  // Restoring bounds is from `bounds_database_` with the key
-  // `DisplayWindowInfo(target_display_id, target_rotation,
-  // target_work_area)` if it exists. Otherwise, calculating the window's
-  // remapping bounds in this target display configuration. There are three
-  // mechanisms of the calculation: 1) keep the window's physical position on
-  // screen rotation 2) keep the same relative position to the center point of
-  // the work area 3) offscreen protection.
-  gfx::Rect RemapOrRestore(aura::Window* window, int64_t target_display_id);
+  // wm::ActivationChangeObserver:
+  void OnWindowActivated(ActivationReason reason,
+                         aura::Window* gained_active,
+                         aura::Window* lost_active) override;
 
   // Adds `window` and its host display id to `window_to_display_map_` before
   // removing its host display.
@@ -84,32 +89,52 @@
 
   using WindowBoundsMap = base::flat_map<WindowDisplayInfo, gfx::Rect>;
 
-  // Resets `bounds_database_`.
-  void ResetBoundsDatabase();
+  // Stores the window's bounds in its current display for restoring the window
+  // back to this display later. Calculates and stores the window's remapping
+  // bounds inside the target display configuration. There are three mechanisms
+  // of calculating the remapping bounds 1) keep the window's physical position
+  // on screen rotation 2) keep the same relative position to the center point
+  // of the work area 3) offscreen protection.
+  //
+  // Remapping will be applied to a window if it is moved to a display that it
+  // has never been there before, and no user-assigned bounds is assigned. And
+  // restoring will be applied if the window has been moved back to a display
+  // configuration that it has been there before.
+  //
+  // Note: This function should be called before `window` being moved to the
+  // target display.
+  void RemapOrRestore(aura::Window* window, int64_t target_display_id);
 
   // Stops observing `window` and removes it from the `bounds_database_`.
   void RemoveWindowFromBoundsDatabase(aura::Window* window);
 
-  // Updates the window's bounds stored in `bounds_database_` to the key
-  // `DisplayWindowInfo(display_id, rotation, work_area)`. Returns the bounds
-  // database of `window` stored in `bounds_database_`.
+  // Updates the window's bounds stored in `bounds_database_` on the key
+  // `window_display_info` to the given `bounds`. Returns the bounds database of
+  // `window` stored in `bounds_database_`.
   WindowBoundsMap& UpdateBoundsDatabaseOfWindow(
       aura::Window* window,
-      int64_t display_id,
-      display::Display::Rotation rotation,
-      const gfx::Rect& work_area);
+      const WindowDisplayInfo& window_display_info,
+      const gfx::Rect& bounds);
 
   // Stores the window's host display id when removing its host display, which
   // will be used to restore the window when its host display being reconnected
   // later.
   base::flat_map<aura::Window*, int64_t> window_to_display_map_;
 
+  // The window that is being moved between displays through the shortcut
+  // `kMoveActiveWindowBetweenDisplays`.
+  raw_ptr<aura::Window, ExperimentalAsh> moving_window_between_displays_ =
+      nullptr;
+
   // TODO: Figure out how we can redesign this data structure, then extra data
   // structures like `window_to_display_map_` above can be removed.
   // The database that stores the window's bounds in each display configuration.
   // `WindowDisplayInfo` defines the display configuration changes that we are
   // tracking. Note: stored window bounds are in parent coordinates.
   std::unordered_map<aura::Window*, WindowBoundsMap> bounds_database_;
+
+  base::ScopedMultiSourceObservation<aura::Window, aura::WindowObserver>
+      window_observations_{this};
 };
 
 }  // namespace ash
diff --git a/ash/wm/bounds_tracker/window_bounds_tracker_unittests.cc b/ash/wm/bounds_tracker/window_bounds_tracker_unittests.cc
index 534fef0..d16b44f 100644
--- a/ash/wm/bounds_tracker/window_bounds_tracker_unittests.cc
+++ b/ash/wm/bounds_tracker/window_bounds_tracker_unittests.cc
@@ -144,24 +144,26 @@
   EXPECT_TRUE(second_display_work_area.Contains(window->GetBoundsInScreen()));
   EXPECT_EQ(second_center_point, window->GetBoundsInScreen().CenterPoint());
 
-  // Moves the window to the top left center of the primary display.
+  // Creates another window at the top left center of the primary display.
   const gfx::Point top_left_center(first_center_point.x() / 2,
                                    first_center_point.y() / 2);
   const gfx::Point origin(
       top_left_center -
       gfx::Vector2d(window_size.width() / 2, window_size.height() / 2));
-  window->SetBoundsInScreen(gfx::Rect(origin, window_size), first_display);
+  aura::Window* window2 =
+      CreateTestWindowInShellWithBounds(gfx::Rect(origin, window_size));
+  wm::ActivateWindow(window2);
 
-  // Using the shortcut to move the window to the secondary display, it should
+  // Using the shortcut to move `window2` to the secondary display, it should
   // stay at the top left center of the secondary display.
   PressAndReleaseKey(ui::VKEY_M, ui::EF_COMMAND_DOWN | ui::EF_ALT_DOWN);
-  EXPECT_TRUE(second_display_work_area.Contains(window->GetBoundsInScreen()));
+  EXPECT_TRUE(second_display_work_area.Contains(window2->GetBoundsInScreen()));
   const gfx::Point second_local_work_area_center =
       secondary_display.GetLocalWorkArea().CenterPoint();
   const gfx::Point second_top_left_center(
       second_local_work_area_center.x() / 2,
       second_local_work_area_center.y() / 2);
-  EXPECT_EQ(second_top_left_center, window->bounds().CenterPoint());
+  EXPECT_EQ(second_top_left_center, window2->bounds().CenterPoint());
 }
 
 // Tests that window's bounds stored in the same display configuration can be
@@ -267,6 +269,7 @@
   const gfx::Size window_size(200, 100);
   const gfx::Rect initial_bounds(gfx::Point(900, 0), window_size);
   aura::Window* window = CreateTestWindowInShellWithBounds(initial_bounds);
+  wm::ActivateWindow(window);
 
   const int64_t primary_id = GetPrimaryDisplay().id();
   const int64_t secondary_id = GetSecondaryDisplay().id();
@@ -315,4 +318,44 @@
   EXPECT_EQ(updated_bounds_in_1st, window->GetBoundsInScreen());
 }
 
+TEST_F(WindowBoundsTrackerTest, RootWindowChanges) {
+  UpdateDisplay("400x300,600x500");
+
+  display::Display first_display = GetPrimaryDisplay();
+  display::Display secondary_display = GetSecondaryDisplay();
+  const gfx::Rect first_display_work_area = first_display.work_area();
+  const gfx::Rect second_display_work_area = secondary_display.work_area();
+
+  // Initially, the window is half-offscreen inside the 2nd display.
+  const gfx::Size window_size(200, 100);
+  const gfx::Rect initial_bounds(gfx::Point(900, 0), window_size);
+  aura::Window* window = CreateTestWindowInShellWithBounds(initial_bounds);
+  wm::ActivateWindow(window);
+  EXPECT_EQ(window->GetRootWindow(), Shell::GetAllRootWindows()[1]);
+
+  // Drag it to the center of the 1st display.
+  const gfx::Point first_center_point = first_display_work_area.CenterPoint();
+  const gfx::Rect first_center_bounds(
+      gfx::Point(first_center_point.x() - window_size.width() / 2,
+                 first_center_point.y() - window_size.height() / 2),
+      window_size);
+  window->SetBoundsInScreen(first_center_bounds, first_display);
+  EXPECT_EQ(window->GetRootWindow(), Shell::GetAllRootWindows()[0]);
+
+  // Using the shortcut to move the window back to the 2nd display. It should be
+  // remapped to the center of the display instead of initial half-offscreen
+  // bounds. As it was dragged from the 2nd to 1st display, its previous
+  // half-offscreen should not be stored in the bounds database.
+  PressAndReleaseKey(ui::VKEY_M, ui::EF_COMMAND_DOWN | ui::EF_ALT_DOWN);
+  EXPECT_NE(initial_bounds, window->GetBoundsInScreen());
+  EXPECT_EQ(second_display_work_area.CenterPoint(),
+            window->GetBoundsInScreen().CenterPoint());
+
+  // Using the shortcut to move the window to the 1st display. It should be
+  // restored to the center of the 1st display. As it was moved from the 1st to
+  // the 2nd through the shortcut before.
+  PressAndReleaseKey(ui::VKEY_M, ui::EF_COMMAND_DOWN | ui::EF_ALT_DOWN);
+  EXPECT_EQ(first_center_bounds, window->GetBoundsInScreen());
+}
+
 }  // namespace ash
diff --git a/ash/wm/window_util.cc b/ash/wm/window_util.cc
index d2a0c00..ca93256 100644
--- a/ash/wm/window_util.cc
+++ b/ash/wm/window_util.cc
@@ -21,7 +21,6 @@
 #include "ash/shelf/shelf.h"
 #include "ash/shell.h"
 #include "ash/shell_delegate.h"
-#include "ash/wm/bounds_tracker/window_bounds_tracker.h"
 #include "ash/wm/float/float_controller.h"
 #include "ash/wm/mru_window_tracker.h"
 #include "ash/wm/overview/overview_controller.h"
@@ -273,16 +272,7 @@
     window_state->SetRestoreBoundsInScreen(restore_bounds);
   }
 
-  auto* window_bounds_tracker = Shell::Get()->window_bounds_tracker();
-  gfx::Rect remapped_bounds;
-  if (window_bounds_tracker) {
-    remapped_bounds = window_bounds_tracker->RemapOrRestore(window, display_id);
-  }
-
   container->AddChild(window);
-  if (window_bounds_tracker) {
-    window->SetBounds(remapped_bounds);
-  }
   return true;
 }