Fix stack reference is not updated if the task is organized

There is an additional stack level for split-screen mode, the child
of the organizer may be the instance referenced by task display area.
So if the child is removed, the stack reference should also be updated.

The corresponding operations of removeChild should be put in the
redirected implementation method, otherwise it will be missed when
propagating the removal to parent, that may cause reusing a detached
task and then crash because it does not have parent and display.

Bug: 154679904
Test: ActivityStackTests#testRemoveOrganizedTask_UpdateStackReference
Change-Id: I9c4032409345c772baa055e32da9eaba9a274b53
diff --git a/services/tests/wmtests/src/com/android/server/wm/ActivityStackTests.java b/services/tests/wmtests/src/com/android/server/wm/ActivityStackTests.java
index 6c209e4..5227f3c 100644
--- a/services/tests/wmtests/src/com/android/server/wm/ActivityStackTests.java
+++ b/services/tests/wmtests/src/com/android/server/wm/ActivityStackTests.java
@@ -234,6 +234,25 @@
     }
 
     @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(
                 WINDOWING_MODE_UNDEFINED, ACTIVITY_TYPE_STANDARD, true /* onTop */);