Merge "Fix stack reference is not updated if the task is organized" into rvc-dev

services/core/java/com/android/server/wm/Task.java

@@ -1431,15 +1431,15 @@ class Task extends WindowContainer<WindowContainer> {
 
     @Override
     void removeChild(WindowContainer child) {
-        // A rootable child task that is now being removed from an organized task. Making sure
-        // the stack references is keep updated.
-        if (mTaskOrganizer != null && mCreatedByOrganizer && child.asTask() != null) {
-            getDisplayArea().removeStackReferenceIfNeeded((ActivityStack) child);
-        }
         removeChild(child, "removeChild");
     }
 
     void removeChild(WindowContainer r, String reason) {
+        // A rootable child task that is now being removed from an organized task. Making sure
+        // the stack references is keep updated.
+        if (mTaskOrganizer != null && mCreatedByOrganizer && r.asTask() != null) {
+            getDisplayArea().removeStackReferenceIfNeeded((ActivityStack) r);
+        }
         if (!mChildren.contains(r)) {
             Slog.e(TAG, "removeChild: r=" + r + " not found in t=" + this);
             return;

services/tests/wmtests/src/com/android/server/wm/ActivityStackTests.java

@@ -233,6 +233,25 @@ public class ActivityStackTests extends ActivityTestsBase {
         assertEquals(homeRoot.getParent(), secondaryTask.getParent());
     }
 
+    @Test
+    public void testRemoveOrganizedTask_UpdateStackReference() {
+        final ActivityStack rootHomeTask = mDefaultTaskDisplayArea.getRootHomeTask();
+        final ActivityRecord homeActivity = new ActivityBuilder(mService)
+                .setStack(rootHomeTask)
+                .setCreateTask(true)
+                .build();
+        final ActivityStack secondaryStack = (ActivityStack) WindowContainer.fromBinder(
+                mService.mTaskOrganizerController.createRootTask(rootHomeTask.getDisplayId(),
+                        WINDOWING_MODE_SPLIT_SCREEN_SECONDARY).token.asBinder());
+
+        rootHomeTask.reparent(secondaryStack, POSITION_TOP);
+        assertEquals(secondaryStack, rootHomeTask.getParent());
+
+        // This should call to {@link TaskDisplayArea#removeStackReferenceIfNeeded}.
+        homeActivity.removeImmediately();
+        assertNull(mDefaultTaskDisplayArea.getRootHomeTask());
+    }
+
     @Test
     public void testStackInheritsDisplayWindowingMode() {
         final ActivityStack primarySplitScreen = mDefaultTaskDisplayArea.createStack(